//! The runner against a mock forge (hyper on a local TCP port) with `HostExecutor`.
//!
//! Enabled for `cargo test` by the self dev-dependency in Cargo.toml, which turns on the
//! `test-host-executor` feature. There is no cfg gate here, so a build without the feature fails
//! to compile instead of silently skipping these tests.

use std::collections::{BTreeMap, HashMap};
use std::convert::Infallible;
use std::path::{Path, PathBuf};
use std::process::Command as StdCommand;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use bytes::Bytes;
use forge_ci::proto::{FinishRequest, JobAssignment, Outcome, RegisterResponse};
use forge_runner::config::{Limits, RunnerConfig};
use forge_runner::executor::HostExecutor;
use forge_runner::logs::PumpConfig;
use forge_runner::runner::{RunOptions, Runner};
use http_body_util::{BodyExt, Full};
use hyper::body::Incoming;
use hyper::header::{AUTHORIZATION, HeaderMap};
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper::{Method, Request, Response, StatusCode};
use hyper_util::rt::TokioIo;
use tokio::net::TcpListener;
use tokio_util::sync::CancellationToken;

const RUNNER_SECRET: &str = "runner-secret-for-tests";
const SECRET_VALUE: &str = "s3cr3t-value-echoed-by-job";
const LOG_SIZE: &str = "X-Forge-Log-Size";
const LOG_OFFSET: &str = "X-Forge-Log-Offset";

/// What the mock forge saw and what it should do.
#[derive(Default)]
struct MockState {
    /// Jobs handed out by claim, in order.
    pending: Vec<JobAssignment>,
    /// Every job by id, so job calls can check the job token.
    by_id: HashMap<String, JobAssignment>,
    /// The stored log, exactly as the forge would keep it.
    log: Vec<u8>,
    log_posts: usize,
    conflicts: usize,
    /// On log POST number `.0`, the forge first stores the first `.1` bytes of the body and then
    /// answers 409, as if an earlier attempt had been applied but its response was lost.
    ghost: Option<(usize, usize)>,
    /// On log POST number `n`, the forge answers 200 but stores nothing, as if the batch was
    /// dropped after it was acknowledged. A later POST then gets 409 with the smaller size.
    drop_post: Option<usize>,
    heartbeats: usize,
    /// Heartbeats from number `n` on answer `cancel: true`.
    cancel_from_heartbeat: Option<usize>,
    artifacts: BTreeMap<String, Vec<u8>>,
    finished: HashMap<String, FinishRequest>,
}

type Shared = Arc<Mutex<MockState>>;

struct Reply {
    status: StatusCode,
    log_size: Option<u64>,
    body: Vec<u8>,
}

impl Reply {
    fn empty(status: StatusCode) -> Self {
        Self {
            status,
            log_size: None,
            body: Vec::new(),
        }
    }

    fn json<T: serde::Serialize>(value: &T) -> Self {
        Self {
            status: StatusCode::OK,
            log_size: None,
            body: serde_json::to_vec(value).unwrap_or_default(),
        }
    }
}

fn bearer(headers: &HeaderMap) -> String {
    headers
        .get(AUTHORIZATION)
        .and_then(|v| v.to_str().ok())
        .and_then(|v| v.strip_prefix("Bearer "))
        .unwrap_or_default()
        .to_owned()
}

fn route(
    st: &mut MockState,
    method: &Method,
    path: &str,
    auth: &str,
    offset: Option<u64>,
    body: &[u8],
) -> Reply {
    let parts: Vec<&str> = path.trim_matches('/').split('/').collect();
    match (method.as_str(), parts.as_slice()) {
        ("POST", ["api", "v1", "ci", "runners", "register"]) => Reply::json(&RegisterResponse {
            runner_id: "runner-1".to_owned(),
            runner_secret: RUNNER_SECRET.to_owned(),
        }),
        ("POST", ["api", "v1", "ci", "runners", "claim"]) => {
            if auth != RUNNER_SECRET {
                return Reply::empty(StatusCode::UNAUTHORIZED);
            }
            if st.pending.is_empty() {
                return Reply::empty(StatusCode::NO_CONTENT);
            }
            let job = st.pending.remove(0);
            Reply::json(&job)
        }
        ("POST" | "PUT" | "GET", ["api", "v1", "ci", "jobs", job_id, rest @ ..]) => {
            let Some(job) = st.by_id.get(*job_id) else {
                return Reply::empty(StatusCode::NOT_FOUND);
            };
            if auth != job.job_token {
                return Reply::empty(StatusCode::UNAUTHORIZED);
            }
            match (method.as_str(), rest) {
                ("POST", ["logs"]) => log_post(st, offset, body),
                ("POST", ["heartbeat"]) => {
                    st.heartbeats += 1;
                    let cancel = st.cancel_from_heartbeat.is_some_and(|n| st.heartbeats >= n);
                    Reply::json(&forge_ci::proto::HeartbeatResponse { cancel })
                }
                ("PUT", ["artifacts", name]) => {
                    let valid = !name.is_empty()
                        && name.len() <= 128
                        && !name.starts_with('.')
                        && name
                            .bytes()
                            .all(|b| b.is_ascii_alphanumeric() || b"._+-".contains(&b));
                    if !valid {
                        return Reply::empty(StatusCode::BAD_REQUEST);
                    }
                    st.artifacts.insert((*name).to_owned(), body.to_vec());
                    Reply::empty(StatusCode::OK)
                }
                ("POST", ["finish"]) => match serde_json::from_slice::<FinishRequest>(body) {
                    Ok(req) => {
                        st.finished.insert((*job_id).to_owned(), req);
                        Reply::empty(StatusCode::OK)
                    }
                    Err(_) => Reply::empty(StatusCode::BAD_REQUEST),
                },
                _ => Reply::empty(StatusCode::NOT_FOUND),
            }
        }
        _ => Reply::empty(StatusCode::NOT_FOUND),
    }
}

/// The forge's log rules: append at size, accept an idempotent retry, otherwise 409 with the size.
fn log_post(st: &mut MockState, offset: Option<u64>, body: &[u8]) -> Reply {
    let Some(offset) = offset else {
        return Reply::empty(StatusCode::BAD_REQUEST);
    };
    st.log_posts += 1;
    if st.drop_post == Some(st.log_posts) {
        return Reply::empty(StatusCode::OK);
    }
    if let Some((nth, prefix)) = st.ghost
        && st.log_posts == nth
    {
        let n = prefix.min(body.len());
        st.log.extend_from_slice(&body[..n]);
    }
    let size = st.log.len() as u64;
    if offset == size {
        st.log.extend_from_slice(body);
        return Reply::empty(StatusCode::OK);
    }
    let end = offset + body.len() as u64;
    if end == size && st.log[offset as usize..] == *body {
        return Reply::empty(StatusCode::OK);
    }
    st.conflicts += 1;
    Reply {
        status: StatusCode::CONFLICT,
        log_size: Some(size),
        body: Vec::new(),
    }
}

async fn handle(
    state: Shared,
    req: Request<Incoming>,
) -> Result<Response<Full<Bytes>>, Infallible> {
    let method = req.method().clone();
    let path = req.uri().path().to_owned();
    let auth = bearer(req.headers());
    let offset = req
        .headers()
        .get(LOG_OFFSET)
        .and_then(|v| v.to_str().ok())
        .and_then(|v| v.parse::<u64>().ok());
    let body = req
        .into_body()
        .collect()
        .await
        .map(|c| c.to_bytes().to_vec())
        .unwrap_or_default();
    // Slow claims down a little so an idle runner does not spin.
    if path.ends_with("/claim") {
        tokio::time::sleep(Duration::from_millis(20)).await;
    }
    let reply = {
        let mut st = match state.lock() {
            Ok(st) => st,
            Err(poisoned) => poisoned.into_inner(),
        };
        route(&mut st, &method, &path, &auth, offset, &body)
    };
    let mut builder = Response::builder().status(reply.status);
    if let Some(size) = reply.log_size {
        builder = builder.header(LOG_SIZE, size.to_string());
    }
    let response = builder
        .header("content-type", "application/json")
        .body(Full::new(Bytes::from(reply.body)))
        .unwrap_or_else(|_| Response::new(Full::new(Bytes::new())));
    Ok(response)
}

/// Starts the mock forge on a random local port and returns its base URL.
async fn start_mock(state: Shared) -> String {
    let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
    let addr = listener.local_addr().expect("addr");
    tokio::spawn(async move {
        loop {
            let Ok((stream, _)) = listener.accept().await else {
                break;
            };
            let st = state.clone();
            tokio::spawn(async move {
                let _ = http1::Builder::new()
                    .serve_connection(
                        TokioIo::new(stream),
                        service_fn(move |req| handle(st.clone(), req)),
                    )
                    .await;
            });
        }
    });
    format!("http://{addr}")
}

/// A git repo with one commit. Returns the commit id.
fn make_repo(dir: &Path) -> String {
    let repo = dir.join("repo");
    std::fs::create_dir_all(&repo).expect("mkdir");
    std::fs::write(repo.join("README.md"), "hello\n").expect("write");
    let git = |args: &[&str]| {
        let status = StdCommand::new("git")
            .arg("-C")
            .arg(&repo)
            .args(args)
            .status()
            .expect("run git");
        assert!(status.success(), "git {args:?} failed");
    };
    git(&["init", "-q"]);
    git(&["add", "-A"]);
    git(&[
        "-c",
        "user.name=test",
        "-c",
        "user.email=test@example.com",
        "commit",
        "-q",
        "-m",
        "init",
    ]);
    // Let a file:// fetch ask for the exact commit by id.
    git(&["config", "uploadpack.allowAnySHA1InWant", "true"]);
    let out = StdCommand::new("git")
        .arg("-C")
        .arg(&repo)
        .args(["rev-parse", "HEAD"])
        .output()
        .expect("rev-parse");
    String::from_utf8_lossy(&out.stdout).trim().to_owned()
}

fn job(job_id: &str, commit: &str, clone_url: &str, run: &str) -> JobAssignment {
    JobAssignment {
        job_id: job_id.to_owned(),
        job_token: format!("token-for-{job_id}"),
        repo: "owner/repo".to_owned(),
        clone_url: clone_url.to_owned(),
        commit: commit.to_owned(),
        git_ref: "refs/heads/main".to_owned(),
        job_name: "build".to_owned(),
        image: "docker.io/library/alpine:3".to_owned(),
        run: run.to_owned(),
        env: BTreeMap::from([("GREETING".to_owned(), "hello-env".to_owned())]),
        secrets: BTreeMap::from([("MY_TOKEN".to_owned(), SECRET_VALUE.to_owned())]),
        artifacts: vec!["dist/**".to_owned()],
        timeout_minutes: 5,
        network: false,
    }
}

fn config(url: &str, work_dir: PathBuf) -> RunnerConfig {
    RunnerConfig {
        url: url.to_owned(),
        runner_id: "runner-1".to_owned(),
        runner_secret: RUNNER_SECRET.to_owned(),
        name: "test-runner".to_owned(),
        labels: vec!["linux".to_owned()],
        work_dir,
        concurrency: 1,
        limits: Limits::default(),
        executor: forge_runner::config::ExecutorKind::default(),
        images: forge_runner::config::ImagesConfig::default(),
        sandbox: forge_runner::config::SandboxConfig::default(),
        registries: Vec::new(),
    }
}

fn fast_options() -> RunOptions {
    RunOptions {
        heartbeat_every: Duration::from_millis(50),
        claim_backoff: Duration::from_millis(50),
        pump: PumpConfig {
            flush_every: Duration::from_millis(50),
            retries: 3,
            retry_delay: Duration::from_millis(10),
        },
        timeout_override: None,
    }
}

/// Polls `check` until it returns true or the deadline passes.
async fn wait_until(what: &str, mut check: impl FnMut() -> bool) {
    let deadline = Instant::now() + Duration::from_secs(60);
    while !check() {
        assert!(Instant::now() < deadline, "timed out waiting for {what}");
        tokio::time::sleep(Duration::from_millis(20)).await;
    }
}

struct Harness {
    state: Shared,
    stop: CancellationToken,
    runner_task: tokio::task::JoinHandle<()>,
    _dir: tempfile::TempDir,
}

/// Starts the mock, a runner with HostExecutor, and queues one job.
async fn start(
    run: &str,
    mut options: RunOptions,
    setup: impl FnOnce(&mut MockState),
    timeout: Option<Duration>,
) -> (Harness, String) {
    let dir = tempfile::tempdir().expect("tempdir");
    let commit = make_repo(dir.path());
    let clone_url = format!("file://{}", dir.path().join("repo").display());
    let job_id = "job-1";
    let state: Shared = Arc::new(Mutex::new(MockState::default()));
    {
        let mut st = state.lock().expect("lock");
        let j = job(job_id, &commit, &clone_url, run);
        st.by_id.insert(j.job_id.clone(), j.clone());
        st.pending.push(j);
        setup(&mut st);
    }
    let url = start_mock(state.clone()).await;
    options.timeout_override = timeout;
    let cfg = config(&url, dir.path().join("work"));
    let runner = Runner::new(cfg, Arc::new(HostExecutor), options).expect("runner");
    let stop = CancellationToken::new();
    let runner_task = {
        let stop = stop.clone();
        tokio::spawn(async move { runner.run(stop).await })
    };
    (
        Harness {
            state,
            stop,
            runner_task,
            _dir: dir,
        },
        job_id.to_owned(),
    )
}

async fn finish_of(h: &Harness, job_id: &str) -> FinishRequest {
    let state = h.state.clone();
    let id = job_id.to_owned();
    wait_until("the job to finish", || {
        state
            .lock()
            .map(|st| st.finished.contains_key(&id))
            .unwrap_or(false)
    })
    .await;
    h.state
        .lock()
        .expect("lock")
        .finished
        .get(job_id)
        .cloned()
        .expect("finish recorded")
}

async fn shut_down(h: Harness) {
    h.stop.cancel();
    let _ = tokio::time::timeout(Duration::from_secs(20), h.runner_task).await;
}

fn log_text(h: &Harness) -> String {
    String::from_utf8_lossy(&h.state.lock().expect("lock").log).into_owned()
}

#[tokio::test]
async fn full_flow_masks_secrets_and_uploads_artifacts() {
    let run = "mkdir -p dist && echo built > dist/out.txt && echo \"token is $MY_TOKEN\" && \
               echo \"b64 $(printf %s \"$MY_TOKEN\" | base64)\" && echo \"env is $GREETING\"";
    let (h, job_id) = start(run, fast_options(), |_| {}, None).await;

    let finish = finish_of(&h, &job_id).await;
    assert_eq!(finish.outcome, Outcome::Success, "{finish:?}");
    assert_eq!(finish.exit_code, Some(0));

    let log = log_text(&h);
    assert!(log.starts_with("forge-runner "), "info line first: {log:?}");
    assert!(log.contains("token is ***"), "masked: {log:?}");
    assert!(log.contains("env is hello-env"), "{log:?}");
    assert!(!log.contains(SECRET_VALUE), "secret leaked into the log");
    let b64 = base64_of(SECRET_VALUE);
    assert!(
        !log.contains(&b64),
        "base64 of the secret leaked into the log"
    );

    let artifact = h
        .state
        .lock()
        .expect("lock")
        .artifacts
        .get("dist_out.txt")
        .cloned();
    assert_eq!(artifact.as_deref(), Some(&b"built\n"[..]));

    shut_down(h).await;
}

fn base64_of(text: &str) -> String {
    use base64::Engine;
    base64::engine::general_purpose::STANDARD.encode(text)
}

#[tokio::test]
async fn heartbeat_cancel_stops_the_job() {
    let mut options = fast_options();
    options.heartbeat_every = Duration::from_millis(50);
    let (h, job_id) = start(
        "echo start; sleep 60",
        options,
        |st| st.cancel_from_heartbeat = Some(2),
        None,
    )
    .await;

    let started = Instant::now();
    let finish = finish_of(&h, &job_id).await;
    assert_eq!(finish.outcome, Outcome::Cancelled, "{finish:?}");
    assert!(
        started.elapsed() < Duration::from_secs(30),
        "the job was not stopped"
    );
    shut_down(h).await;
}

#[tokio::test]
async fn timeout_fails_the_job() {
    let (h, job_id) = start(
        "sleep 60",
        fast_options(),
        |_| {},
        Some(Duration::from_millis(400)),
    )
    .await;

    let finish = finish_of(&h, &job_id).await;
    assert_eq!(finish.outcome, Outcome::Failure, "{finish:?}");
    assert!(
        finish
            .message
            .as_deref()
            .unwrap_or_default()
            .starts_with("timed out"),
        "{finish:?}"
    );
    shut_down(h).await;
}

#[tokio::test]
async fn log_offset_mismatch_resyncs_without_loss_or_duplicates() {
    // Log POST #2 is applied by the forge for its first 3 bytes, then answered 409. The runner
    // must drop those 3 bytes and continue, so the stored log has every line exactly once.
    let run = "echo first; sleep 0.3; echo second; sleep 0.3; echo third";
    let (h, job_id) = start(run, fast_options(), |st| st.ghost = Some((2, 3)), None).await;

    let finish = finish_of(&h, &job_id).await;
    assert_eq!(finish.outcome, Outcome::Success, "{finish:?}");

    let log = log_text(&h);
    assert!(log.ends_with("first\nsecond\nthird\n"), "log was {log:?}");
    assert_eq!(
        log.matches("second").count(),
        1,
        "duplicate or lost bytes: {log:?}"
    );
    let conflicts = h.state.lock().expect("lock").conflicts;
    assert!(conflicts >= 1, "the 409 path never ran");
    shut_down(h).await;
}

#[tokio::test]
async fn acknowledged_batch_the_forge_lost_is_resent_in_order() {
    // Log POST #2 is acknowledged but never stored. The next POST gets 409 with the smaller size,
    // and the runner must resend the lost bytes from its retained buffer, with no gap marker.
    let run = "echo first; sleep 0.3; echo second; sleep 0.3; echo third";
    let (h, job_id) = start(run, fast_options(), |st| st.drop_post = Some(2), None).await;

    let finish = finish_of(&h, &job_id).await;
    assert_eq!(finish.outcome, Outcome::Success, "{finish:?}");

    let log = log_text(&h);
    assert!(log.ends_with("first\nsecond\nthird\n"), "log was {log:?}");
    assert_eq!(log.matches("second").count(), 1, "duplicate bytes: {log:?}");
    assert!(
        !log.contains("[log gap]"),
        "resend should not need a gap: {log:?}"
    );
    let conflicts = h.state.lock().expect("lock").conflicts;
    assert!(conflicts >= 1, "the 409 path never ran");
    shut_down(h).await;
}