//! oak-engine — the protocol-transport push engine for the oakmark throughput

//! benchmark.

//!

//! One process is spawned per synthetic agent by the Go harness and driven

//! serially over line-delimited JSON on stdin/stdout. Each command produces

//! exactly one JSON response line (flushed immediately). All diagnostics go to

//! stderr so stdout stays a clean JSON stream.

//!

//! It reuses `oak_core` to build blobs / trees / commits so every content

//! address is correct by construction (the server re-verifies them), and

//! hand-writes the HTTP push/pull against `oak_core::protocol` DTOs. It never

//! shells out to the `oak` CLI.

use std::io::{BufRead, Write};

use base64::engine::general_purpose::STANDARD as B64;

use base64::Engine as _;

use oak_core::protocol::{

    self, BlobCheckRequest, BlobCheckResponse, BlobData, BranchPushData, CommitData, FileChangeData,

    PullResponse, PushRequest, PushResponse,

};

use oak_core::{

    build_tree, Blob, ChangeType, Commit, FileChange, FileMode, Hash, ManifestEntry,

};

use serde_json::{json, Value};

/// Per-agent engine state, filled in by `init` and mutated as pushes land.

struct Engine {

    http: reqwest::blocking::Client,

    remote: String, // no trailing slash

    owner: String,

    name: String,

    token: String, // empty => anonymous (no auth header)

    author: String,

    branch: String,

    /// The branch's last-known remote head (None until the branch has a

    /// commit on the server). Used for the CAS `expected_*` fields.

    branch_head: Option<String>,

    /// Parent commit hash for the next commit we build. Seeded from the branch

    /// head if the branch exists, else from `main`'s head (the fork point), and

    /// advanced to each accepted commit so the branch is a proper chain.

    parent_next: Option<String>,

    /// Full path -> blob-hash snapshot of the current branch head, so we can

    /// emit accurate Added/Modified/Deleted file changes on the next push.

    prev_files: std::collections::HashMap<String, String>,

}

impl Engine {

    fn auth<'a>(&self, rb: reqwest::blocking::RequestBuilder) -> reqwest::blocking::RequestBuilder {

        if self.token.is_empty() {

            rb

        } else {

            rb.header("authorization", format!("Bearer {}", self.token))

        }

    }

    /// `GET /pull?branch_name=<b>&force=true` and return the parsed response.

    fn pull(&self, branch_name: &str) -> Result<PullResponse, String> {

        let url = format!(

            "{}/api/{}/{}/pull?branch_name={}&force=true",

            self.remote, self.owner, self.name, branch_name

        );

        let resp = self

            .auth(self.http.get(&url))

            .send()

            .map_err(|e| format!("pull transport error: {e}"))?;

        let status = resp.status();

        if !status.is_success() {

            return Err(format!("pull {branch_name} returned HTTP {}", status.as_u16()));

        }

        resp.json::<PullResponse>()

            .map_err(|e| format!("pull decode error: {e}"))

    }

}

/// One decoded file from a push command.

struct InFile {

    path: String,

    content: Vec<u8>,

}

fn main() {

    let stdin = std::io::stdin();

    let mut out = std::io::stdout();

    let mut engine: Option<Engine> = None;

    for line in stdin.lock().lines() {

        let line = match line {

            Ok(l) => l,

            Err(e) => {

                eprintln!("oak-engine: stdin read error: {e}");

                break;

            }

        };

        let trimmed = line.trim();

        if trimmed.is_empty() {

            continue;

        }

        let cmd: Value = match serde_json::from_str(trimmed) {

            Ok(v) => v,

            Err(e) => {

                write_line(&mut out, &json!({"ok": false, "err": format!("bad json: {e}")}));

                continue;

            }

        };

        let op = cmd.get("op").and_then(Value::as_str).unwrap_or("");

        match op {

            "init" => {

                let resp = do_init(&cmd);

                match resp {

                    Ok(e) => {

                        engine = Some(e);

                        write_line(&mut out, &json!({"ok": true}));

                    }

                    Err(err) => write_line(&mut out, &json!({"ok": false, "err": err})),

                }

            }

            "push" => {

                let resp = match engine.as_mut() {

                    Some(e) => do_push(e, &cmd),

                    None => json!({"ok": false, "cas": false, "new_head": Value::Null, "err": "push before init"}),

                };

                write_line(&mut out, &resp);

            }

            "clone" => {

                let resp = match engine.as_ref() {

                    Some(e) => do_clone(e),

                    None => json!({"ok": false, "clone_ms": 0.0, "err": "clone before init"}),

                };

                write_line(&mut out, &resp);

            }

            "quit" => {

                std::process::exit(0);

            }

            other => {

                write_line(

                    &mut out,

                    &json!({"ok": false, "err": format!("unknown op: {other}")}),

                );

            }

        }

    }

}

fn write_line(out: &mut std::io::Stdout, v: &Value) {

    // One compact JSON object per line, flushed so the harness sees it

    // immediately (it drives us request/response, serially).

    let _ = writeln!(out, "{v}");

    let _ = out.flush();

}

fn do_init(cmd: &Value) -> Result<Engine, String> {

    let get = |k: &str| cmd.get(k).and_then(Value::as_str).unwrap_or("").to_string();

    let remote = get("remote");

    let remote = remote.trim_end_matches('/').to_string();

    let owner = get("owner");

    let name = get("name");

    let token = get("token");

    let author = get("author");

    let branch = get("branch");

    if remote.is_empty() || owner.is_empty() || name.is_empty() || branch.is_empty() {

        return Err("init requires remote, owner, name, branch".to_string());

    }

    let http = reqwest::blocking::Client::builder()

        .build()

        .map_err(|e| format!("http client build: {e}"))?;

    let mut engine = Engine {

        http,

        remote,

        owner,

        name,

        token,

        author,

        branch,

        branch_head: None,

        parent_next: None,

        prev_files: std::collections::HashMap::new(),

    };

    // Does our branch already exist on the server (has a head)? If so we chain

    // onto it; otherwise we fork from main's head.

    let branch_pull = engine.pull(&engine.branch)?;

    if let Some(head) = branch_pull.head.clone() {

        engine.branch_head = Some(head.clone());

        engine.parent_next = Some(head);

        eprintln!("oak-engine[{}]: branch exists, head={:?}", engine.branch, engine.branch_head);

    } else {

        let main_pull = engine.pull("main")?;

        engine.branch_head = None;

        engine.parent_next = main_pull.head.clone();

        eprintln!(

            "oak-engine[{}]: new branch, forking from main head={:?}",

            engine.branch, engine.parent_next

        );

    }

    Ok(engine)

}

fn do_push(engine: &mut Engine, cmd: &Value) -> Value {

    match build_and_push(engine, cmd) {

        Ok(PushOutcome::Ok { new_head }) => {

            json!({"ok": true, "cas": false, "new_head": new_head, "err": Value::Null})

        }

        Ok(PushOutcome::Cas { new_head, message }) => {

            json!({"ok": false, "cas": true, "new_head": new_head, "err": message})

        }

        Err(err) => {

            json!({"ok": false, "cas": false, "new_head": Value::Null, "err": err})

        }

    }

}

enum PushOutcome {

    Ok { new_head: String },

    Cas { new_head: Value, message: String },

}

fn build_and_push(engine: &mut Engine, cmd: &Value) -> Result<PushOutcome, String> {

    let files = decode_files(cmd)?;

    // 1. Blobs + manifest entries. Blob::new computes the BLAKE3 content hash

    //    the server re-verifies.

    let mut entries: Vec<ManifestEntry> = Vec::with_capacity(files.len());

    // Unique blob hash -> content, for the push payload.

    let mut blob_by_hash: std::collections::HashMap<String, (Vec<u8>, u64)> =

        std::collections::HashMap::new();

    // The new full snapshot: path -> blob hash.

    let mut new_files: std::collections::HashMap<String, String> =

        std::collections::HashMap::new();

    for f in &files {

        let blob = Blob::new(f.content.clone());

        let hh = blob.hash.to_string();

        entries.push(ManifestEntry {

            path: f.path.clone(),

            blob_hash: blob.hash.clone(),

            mode: FileMode::Regular,

        });

        new_files.insert(f.path.clone(), hh.clone());

        blob_by_hash

            .entry(hh)

            .or_insert((blob.content, blob.size));

    }

    // 2. Tree objects (root hash == the commit's manifest_hash).

    let built = build_tree(&entries).map_err(|e| format!("build_tree: {e}"))?;

    let manifest_hash = built.root_hash.clone();

    // 3. File changes vs the previous snapshot (Added/Modified/Deleted). The v1

    //    commit hash excludes this list, but the server stores and validates the

    //    individual rows, so we send accurate ones.

    let file_changes = diff_snapshots(&engine.prev_files, &new_files);

    // 4. Commit (computes its own verified hash from the truncated-micros now()).

    let parent_hash = engine

        .parent_next

        .as_ref()

        .map(|s| Hash::from_hex(s))

        .transpose()

        .map_err(|e| format!("bad parent hash: {e}"))?;

    let commit = Commit::new(

        engine.branch.clone(),

        parent_hash,

        None,

        manifest_hash,

        engine.author.clone(),

        None,

        file_changes,

    )

    .map_err(|e| format!("Commit::new: {e}"))?;

    let commit_data = commit_to_wire(&commit);

    // 5. Blob payload — optionally drop blobs the server already has.

    let mut blobs: Vec<BlobData> = blob_by_hash

        .into_iter()

        .map(|(hash, (content, size))| BlobData {

            hash,

            content,

            size,

            chunks: Vec::new(),

        })

        .collect();

    maybe_drop_present_blobs(engine, &mut blobs);

    // 6. Trees on the wire.

    let trees = built.trees.iter().map(protocol::tree_to_wire).collect();

    // 7. Branch metadata (flat model: parent=main, open).

    let branch_data = BranchPushData {

        name: engine.branch.clone(),

        description: None,

        parent_branch: Some("main".to_string()),

        status: "open".to_string(),

        created_at: chrono::Utc::now().to_rfc3339(),

        close_reason: None,

    };

    // 8. The push request. force=true (each agent owns its branch), but we still

    //    send the tracked heads so the chain is explicit.

    let push_req = PushRequest {

        expected_head: engine.branch_head.clone(),

        expected_branch_head: engine.branch_head.clone(),

        force: true,

        branch: Some(branch_data),

        commits: vec![commit_data],

        blobs,

        trees,

    };

    let url = format!("{}/api/{}/{}/push", engine.remote, engine.owner, engine.name);

    let resp = engine

        .auth(

            engine

                .http

                .post(&url)

                .header("x-oak-user", engine.author.clone())

                .json(&push_req),

        )

        .send()

        .map_err(|e| format!("push transport error: {e}"))?;

    let status = resp.status();

    if !status.is_success() {

        let body = resp.text().unwrap_or_default();

        return Err(format!("push HTTP {}: {}", status.as_u16(), truncate(&body, 300)));

    }

    let pr: PushResponse = resp

        .json()

        .map_err(|e| format!("push decode error: {e}"))?;

    if pr.success {

        let new_head = pr

            .new_head

            .clone()

            .unwrap_or_else(|| commit.hash.to_string());

        // Advance tracked state only on an accepted push.

        engine.branch_head = Some(new_head.clone());

        engine.parent_next = Some(new_head.clone());

        engine.prev_files = new_files;

        Ok(PushOutcome::Ok { new_head })

    } else {

        // HTTP 200 + success:false is a CAS/conflict, not an error. Leave state

        // untouched so the caller can retry against the real head.

        Ok(PushOutcome::Cas {

            new_head: pr

                .new_head

                .map(Value::from)

                .unwrap_or(Value::Null),

            message: format!("success:false: {}", pr.message),

        })

    }

}

/// Best-effort `POST /blobs/check` to drop blobs the server already has. Skipped

/// when the inline payload is below one round trip's worth of bytes (64 KiB).

/// Any failure falls through to "send everything," which is always correct.

fn maybe_drop_present_blobs(engine: &Engine, blobs: &mut Vec<BlobData>) {

    if blobs.is_empty() {

        return;

    }

    const MIN_BYTES: u64 = 64 * 1024;

    let inline: u64 = blobs.iter().map(|b| b.content.len() as u64).sum();

    if inline < MIN_BYTES {

        eprintln!(

            "oak-engine[{}]: skipping blobs/check ({} blob(s), {inline} B < {MIN_BYTES})",

            engine.branch,

            blobs.len()

        );

        return;

    }

    let hashes: Vec<String> = blobs.iter().map(|b| b.hash.clone()).collect();

    let url = format!(

        "{}/api/{}/{}/blobs/check",

        engine.remote, engine.owner, engine.name

    );

    let req = BlobCheckRequest { hashes };

    match engine.auth(engine.http.post(&url).json(&req)).send() {

        Ok(resp) if resp.status().is_success() => match resp.json::<BlobCheckResponse>() {

            Ok(body) => {

                let missing: std::collections::HashSet<String> =

                    body.missing.into_iter().collect();

                let before = blobs.len();

                blobs.retain(|b| missing.contains(&b.hash));

                eprintln!(

                    "oak-engine[{}]: blobs/check kept {}/{}",

                    engine.branch,

                    blobs.len(),

                    before

                );

            }

            Err(e) => eprintln!("oak-engine: blobs/check decode failed ({e}); sending all"),

        },

        Ok(resp) => eprintln!(

            "oak-engine: blobs/check HTTP {}; sending all",

            resp.status().as_u16()

        ),

        Err(e) => eprintln!("oak-engine: blobs/check request failed ({e}); sending all"),

    }

}

fn do_clone(engine: &Engine) -> Value {

    let t0 = std::time::Instant::now();

    match engine.pull("main") {

        Ok(pr) => {

            let ms = t0.elapsed().as_secs_f64() * 1000.0;

            eprintln!(

                "oak-engine[{}]: clone (pull main) head={:?} commits={} blobs={} trees={} in {ms:.3}ms",

                engine.branch,

                pr.head,

                pr.commits.len(),

                pr.blobs.len(),

                pr.trees.len()

            );

            json!({"ok": true, "clone_ms": ms, "err": Value::Null})

        }

        Err(e) => {

            let ms = t0.elapsed().as_secs_f64() * 1000.0;

            json!({"ok": false, "clone_ms": ms, "err": e})

        }

    }

}

fn decode_files(cmd: &Value) -> Result<Vec<InFile>, String> {

    let arr = cmd

        .get("files")

        .and_then(Value::as_array)

        .ok_or_else(|| "push requires a files array".to_string())?;

    let mut out = Vec::with_capacity(arr.len());

    for f in arr {

        let path = f

            .get("path")

            .and_then(Value::as_str)

            .ok_or_else(|| "file missing path".to_string())?

            .to_string();

        let b64 = f

            .get("content_b64")

            .and_then(Value::as_str)

            .ok_or_else(|| "file missing content_b64".to_string())?;

        let content = B64

            .decode(b64)

            .map_err(|e| format!("bad base64 for {path}: {e}"))?;

        out.push(InFile { path, content });

    }

    Ok(out)

}

/// Diff two path->blob-hash snapshots into FileChanges. Unchanged files produce

/// no change (matching how `oak commit` records only real changes).

fn diff_snapshots(

    prev: &std::collections::HashMap<String, String>,

    next: &std::collections::HashMap<String, String>,

) -> Vec<FileChange> {

    let mut changes = Vec::new();

    for (path, new_hash) in next {

        match prev.get(path) {

            None => changes.push(FileChange {

                path: path.clone(),

                change_type: ChangeType::Added,

                old_blob_hash: None,

                new_blob_hash: Some(Hash(new_hash.clone())),

                old_path: None,

                old_mode: None,

                new_mode: None,

            }),

            Some(old_hash) if old_hash != new_hash => changes.push(FileChange {

                path: path.clone(),

                change_type: ChangeType::Modified,

                old_blob_hash: Some(Hash(old_hash.clone())),

                new_blob_hash: Some(Hash(new_hash.clone())),

                old_path: None,

                old_mode: None,

                new_mode: None,

            }),

            Some(_) => {} // unchanged

        }

    }

    for (path, old_hash) in prev {

        if !next.contains_key(path) {

            changes.push(FileChange {

                path: path.clone(),

                change_type: ChangeType::Deleted,

                old_blob_hash: Some(Hash(old_hash.clone())),

                new_blob_hash: None,

                old_path: None,

                old_mode: None,

                new_mode: None,

            });

        }

    }

    changes

}

/// Convert a core `Commit` to its wire `CommitData`. Mirrors the CLI's push

/// path exactly so the server reproduces the same commit hash.

fn commit_to_wire(c: &Commit) -> CommitData {

    CommitData {

        hash: c.hash.to_string(),

        branch_name: c.branch_name.clone(),

        parent_hash: c.parent_hash.as_ref().map(|h| h.to_string()),

        merge_parent_hash: c.merge_parent_hash.as_ref().map(|h| h.to_string()),

        manifest_hash: c.manifest_hash.to_string(),

        author: c.author.clone(),

        message: c.message.clone(),

        timestamp: c.timestamp.to_rfc3339(),

        files: c

            .files

            .iter()

            .map(|f| FileChangeData {

                path: f.path.clone(),

                change_type: match f.change_type {

                    ChangeType::Added => "added".to_string(),

                    ChangeType::Modified => "modified".to_string(),

                    ChangeType::Deleted => "deleted".to_string(),

                    ChangeType::Renamed => "renamed".to_string(),

                },

                old_blob_hash: f.old_blob_hash.as_ref().map(|h| h.to_string()),

                new_blob_hash: f.new_blob_hash.as_ref().map(|h| h.to_string()),

                old_path: f.old_path.clone(),

            })

            .collect(),

    }

}

fn truncate(s: &str, n: usize) -> String {

    if s.len() <= n {

        return s.to_string();

    }

    // Back off to the nearest char boundary so we never slice mid-codepoint.

    let mut end = n;

    while end > 0 && !s.is_char_boundary(end) {

        end -= 1;

    }

    format!("{}…", &s[..end])

}