//! Job log shipping: merge the job's output, mask secrets, batch it, and POST it with byte offsets.

use std::collections::VecDeque;
use std::time::Duration;

use forge_ci::mask::MaskStream;
use tokio::io::{AsyncBufReadExt, AsyncRead, BufReader};
use tokio::sync::mpsc;
use tokio::time::{Instant, MissedTickBehavior};

use crate::client::{ClientError, ForgeClient, LogAck};

/// A batch never grows past this many bytes. The forge accepts chunks up to 1 MiB.
pub const MAX_BATCH: usize = 64 * 1024;
/// A line longer than this is sent in pieces, so one line cannot grow without bound.
const MAX_LINE: usize = 16 * 1024;
/// How many resyncs one batch may take before the pump gives up.
const MAX_RESYNCS: u32 = 3;

/// Errors from shipping logs.
#[derive(Debug, thiserror::Error)]
pub enum LogError {
    #[error(transparent)]
    Client(#[from] ClientError),
    #[error("forge keeps disagreeing about the log size")]
    Resync,
}

/// Timing and retry settings for the pump.
#[derive(Clone, Copy, Debug)]
pub struct PumpConfig {
    pub flush_every: Duration,
    pub retries: u32,
    pub retry_delay: Duration,
}

impl Default for PumpConfig {
    fn default() -> Self {
        Self {
            flush_every: Duration::from_secs(1),
            retries: 5,
            retry_delay: Duration::from_secs(2),
        }
    }
}

/// Bytes kept for resends. Acknowledged bytes are trimmed past this, oldest first. Unacknowledged
/// bytes are never trimmed (they are at most one batch).
pub const RETAIN_CAP: usize = 8 * 1024 * 1024;
/// Written in place of bytes the forge lost and the runner no longer holds.
pub const GAP_MARKER: &[u8] = b"[log gap]\n";

/// The job's log as the runner holds it, in forge coordinates (byte positions in the log the
/// forge stores).
///
/// - `offset` is the forge's confirmed size. Bytes below it were acknowledged with a 200.
/// - `start..end` is the retained window. It holds every byte from `start` up to `end`: the
///   acknowledged tail kept for resends, and the unacknowledged batch after `offset`.
/// - A 409 that reports size `S` inside the window resends from `S`. A size outside it leaves a
///   gap, which is marked with `GAP_MARKER`.
///
/// Pure state, so it can be tested without a network.
#[derive(Debug)]
pub struct LogBuffer {
    start: u64,
    offset: u64,
    data: VecDeque<u8>,
    cap: usize,
}

impl Default for LogBuffer {
    fn default() -> Self {
        Self::new()
    }
}

impl LogBuffer {
    pub fn new() -> Self {
        Self::with_cap(RETAIN_CAP)
    }

    /// A buffer that keeps at most `cap` bytes of acknowledged output. `cap` is at least one batch.
    pub fn with_cap(cap: usize) -> Self {
        Self {
            start: 0,
            offset: 0,
            data: VecDeque::new(),
            cap: cap.max(MAX_BATCH),
        }
    }

    /// The forge's confirmed size.
    pub fn offset(&self) -> u64 {
        self.offset
    }

    /// The forge position of the oldest retained byte.
    pub fn retained_start(&self) -> u64 {
        self.start
    }

    /// The position just past the last byte the runner holds.
    pub fn end(&self) -> u64 {
        self.start + self.data.len() as u64
    }

    /// Bytes the forge has not acknowledged.
    pub fn pending(&self) -> usize {
        (self.end() - self.offset) as usize
    }

    /// The next batch to send: up to `MAX_BATCH` unacknowledged bytes from `offset`.
    pub fn batch(&self) -> Vec<u8> {
        let from = (self.offset - self.start) as usize;
        let n = self.pending().min(MAX_BATCH);
        self.data.range(from..from + n).copied().collect()
    }

    /// Room left before the pending bytes reach one batch.
    pub fn room(&self) -> usize {
        MAX_BATCH.saturating_sub(self.pending())
    }

    /// Appends as much of `bytes` as fits, and returns how many bytes were taken.
    pub fn take(&mut self, bytes: &[u8]) -> usize {
        let n = self.room().min(bytes.len());
        self.data.extend(&bytes[..n]);
        n
    }

    /// The batch returned by `batch` was stored. Advances `offset` and trims the acknowledged
    /// tail to the cap.
    pub fn acked(&mut self) {
        self.offset += self.pending().min(MAX_BATCH) as u64;
        self.trim();
    }

    fn trim(&mut self) {
        let excess = self.data.len().saturating_sub(self.cap);
        let acked = (self.offset - self.start) as usize;
        let drop = excess.min(acked);
        if drop > 0 {
            self.data.drain(..drop);
            self.start += drop as u64;
            tracing::debug!(
                dropped = drop,
                "trimmed acknowledged log bytes past the retain cap"
            );
        }
    }

    /// The forge reports that it holds `held` bytes. Returns true when nothing is pending.
    ///
    /// - If `held` is inside the retained window, the runner resends from `held`.
    /// - Otherwise the bytes between `held` and the window are unavailable. The runner writes
    ///   `GAP_MARKER` at `held`, then continues with what it still holds.
    pub fn resync(&mut self, held: u64) -> bool {
        if held >= self.start && held <= self.end() {
            self.offset = held;
        } else {
            tracing::warn!(
                held,
                retained_start = self.start,
                end = self.end(),
                "forge log size is outside the retained window, marking a log gap"
            );
            let kept: VecDeque<u8> = if held < self.start {
                // Keep the retained bytes after the marker.
                GAP_MARKER
                    .iter()
                    .copied()
                    .chain(self.data.drain(..))
                    .collect()
            } else {
                GAP_MARKER.iter().copied().collect()
            };
            self.data = kept;
            self.start = held;
            self.offset = held;
        }
        self.pending() == 0
    }
}

/// Ships `rx` to the forge: mask each chunk, batch it, and POST it with its offset.
///
/// Returns the number of bytes the forge holds for this job once `rx` is closed and the last
/// batch is stored.
pub async fn pump(
    client: &ForgeClient,
    job_id: &str,
    job_token: &str,
    mut stream: MaskStream,
    mut rx: mpsc::Receiver<Vec<u8>>,
    cfg: PumpConfig,
) -> Result<u64, LogError> {
    let mut buf = LogBuffer::new();
    let mut ticker = tokio::time::interval_at(Instant::now() + cfg.flush_every, cfg.flush_every);
    ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);

    loop {
        tokio::select! {
            msg = rx.recv() => match msg {
                Some(bytes) => {
                    let masked = stream.push(&bytes);
                    append(&mut buf, &masked, client, job_id, job_token, &cfg).await?;
                }
                None => break,
            },
            _ = ticker.tick() => {
                send(&mut buf, client, job_id, job_token, &cfg).await?;
            }
        }
    }

    let tail = stream.finish();
    append(&mut buf, &tail, client, job_id, job_token, &cfg).await?;
    send(&mut buf, client, job_id, job_token, &cfg).await?;
    Ok(buf.offset())
}

async fn append(
    buf: &mut LogBuffer,
    mut bytes: &[u8],
    client: &ForgeClient,
    job_id: &str,
    job_token: &str,
    cfg: &PumpConfig,
) -> Result<(), LogError> {
    while !bytes.is_empty() {
        if buf.room() == 0 {
            send(buf, client, job_id, job_token, cfg).await?;
        }
        let n = buf.take(bytes);
        bytes = &bytes[n..];
    }
    Ok(())
}

/// Sends the pending batch until the forge has it, retrying transient errors.
async fn send(
    buf: &mut LogBuffer,
    client: &ForgeClient,
    job_id: &str,
    job_token: &str,
    cfg: &PumpConfig,
) -> Result<(), LogError> {
    let mut failures = 0u32;
    let mut resyncs = 0u32;
    while buf.pending() > 0 {
        let batch = buf.batch();
        match client
            .post_logs(job_id, job_token, buf.offset(), batch)
            .await
        {
            Ok(LogAck::Accepted) => buf.acked(),
            Ok(LogAck::Resync(held)) => {
                resyncs += 1;
                if resyncs > MAX_RESYNCS {
                    return Err(LogError::Resync);
                }
                tracing::warn!(
                    held,
                    offset = buf.offset(),
                    "log offset mismatch, resyncing"
                );
                buf.resync(held);
            }
            Err(e) => {
                failures += 1;
                if failures > cfg.retries {
                    return Err(e.into());
                }
                tracing::warn!(error = %e, failures, "log upload failed, retrying");
                tokio::time::sleep(cfg.retry_delay * failures).await;
            }
        }
    }
    Ok(())
}

/// Reads `reader` and sends its output to `tx`, one line per message. A line longer than
/// `MAX_LINE` is split. If the receiver goes away, the reader keeps draining so the child does
/// not block on a full pipe.
pub async fn forward_lines<R>(reader: R, tx: mpsc::Sender<Vec<u8>>)
where
    R: AsyncRead + Unpin,
{
    let mut reader = BufReader::new(reader);
    let mut line: Vec<u8> = Vec::new();
    let mut alive = true;
    loop {
        let (take, done) = {
            let available = match reader.fill_buf().await {
                Ok(b) if !b.is_empty() => b,
                _ => break,
            };
            let (take, done) = match available.iter().position(|&b| b == b'\n') {
                Some(i) => (i + 1, true),
                None => (available.len(), false),
            };
            line.extend_from_slice(&available[..take]);
            (take, done)
        };
        reader.consume(take);
        if (done || line.len() >= MAX_LINE) && alive {
            alive = tx.send(std::mem::take(&mut line)).await.is_ok();
        }
    }
    if !line.is_empty() && alive {
        let _ = tx.send(line).await;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn batches_never_exceed_the_cap() {
        let mut buf = LogBuffer::new();
        let data = vec![b'x'; MAX_BATCH + 10];
        let n = buf.take(&data);
        assert_eq!(n, MAX_BATCH);
        assert_eq!(buf.room(), 0);
        assert_eq!(buf.batch().len(), MAX_BATCH);
        buf.acked();
        assert_eq!(buf.offset(), MAX_BATCH as u64);
        assert_eq!(buf.take(&data[n..]), 10);
    }

    #[test]
    fn resync_drops_bytes_the_forge_already_has() {
        // We sent "abcdef" at offset 0 and the response was lost. The forge stored it.
        let mut buf = LogBuffer::new();
        buf.take(b"abcdef");
        assert!(buf.resync(6));
        assert_eq!(buf.offset(), 6);
        assert!(buf.batch().is_empty());
    }

    #[test]
    fn resync_resends_the_uncovered_tail() {
        let mut buf = LogBuffer::new();
        buf.take(b"abcdef");
        assert!(!buf.resync(2));
        assert_eq!(buf.offset(), 2);
        assert_eq!(buf.batch(), b"cdef");
    }

    #[test]
    fn resync_inside_the_window_resends_acknowledged_bytes() {
        // "abc" was acknowledged, "def" is pending. The forge then reports size 1.
        let mut buf = LogBuffer::new();
        buf.take(b"abc");
        buf.acked();
        buf.take(b"def");
        assert!(!buf.resync(1));
        assert_eq!(buf.offset(), 1);
        assert_eq!(buf.batch(), b"bcdef");
        assert_eq!(buf.end(), 6);
    }

    #[test]
    fn resync_ahead_of_everything_we_hold_marks_a_gap() {
        let mut buf = LogBuffer::new();
        buf.take(b"ab");
        assert!(!buf.resync(10));
        assert_eq!(buf.offset(), 10);
        assert_eq!(buf.batch(), GAP_MARKER);
    }

    #[test]
    fn resync_below_the_window_marks_a_gap_and_keeps_the_rest() {
        // Two full batches are acknowledged. A one-batch cap trims the first.
        let mut small = LogBuffer::with_cap(MAX_BATCH);
        small.take(&vec![b'a'; MAX_BATCH]);
        small.acked();
        small.take(&vec![b'b'; MAX_BATCH]);
        small.acked();
        assert_eq!(small.retained_start(), MAX_BATCH as u64);
        // The forge says it holds 1 byte, which is below the window. The bytes 1..start are gone.
        assert!(!small.resync(1));
        assert_eq!(small.offset(), 1);
        let batch = small.batch();
        assert_eq!(&batch[..GAP_MARKER.len()], GAP_MARKER);
        assert!(batch[GAP_MARKER.len()..].iter().all(|&b| b == b'b'));
        assert_eq!(small.end(), 1 + (GAP_MARKER.len() + MAX_BATCH) as u64);
    }

    #[test]
    fn trim_drops_only_acknowledged_bytes() {
        let mut buf = LogBuffer::with_cap(MAX_BATCH);
        buf.take(&vec![b'a'; MAX_BATCH]);
        buf.acked();
        buf.take(b"tail");
        // The cap is one batch, so the acked batch is trimmed once it is acked.
        buf.acked();
        assert_eq!(buf.offset(), (MAX_BATCH + 4) as u64);
        assert!(buf.retained_start() > 0);
        assert!(buf.retained_start() <= buf.offset());
        assert_eq!(buf.end(), buf.offset());
    }

    #[tokio::test]
    async fn forward_lines_splits_on_newlines() {
        let (tx, mut rx) = mpsc::channel(16);
        let input: &[u8] = b"one\ntwo\nthree";
        forward_lines(input, tx).await;
        let mut got = Vec::new();
        while let Some(line) = rx.recv().await {
            got.push(line);
        }
        assert_eq!(
            got,
            vec![b"one\n".to_vec(), b"two\n".to_vec(), b"three".to_vec()]
        );
    }

    #[tokio::test]
    async fn forward_lines_caps_long_lines() {
        let (tx, mut rx) = mpsc::channel(64);
        let input = vec![b'z'; MAX_LINE * 2 + 5];
        forward_lines(&input[..], tx).await;
        let mut total = 0;
        while let Some(piece) = rx.recv().await {
            assert!(piece.len() <= MAX_LINE);
            total += piece.len();
        }
        assert_eq!(total, input.len());
    }
}