//! A `subagent` turn on the background slot. Do not edit.

mod support;

use bxcore::{LogRecord, SessionId};
use loopd::baseline::Baseline;
use loopd::channel::Context;
use loopd::llama::Client;
use loopd::memory::NoMemory;
use loopd::session::Session;
use loopd::tools::Registry;
use loopd::turn::{Runtime, run_turn};
use std::sync::Arc;
use support::{FakeServer, Home, Reply, ScriptedPort};

const CHAT: &str = "/v1/chat/completions";

fn sse_chunks(chunks: &[serde_json::Value]) -> Reply {
    let mut body = String::new();
    for chunk in chunks {
        body.push_str("data: ");
        body.push_str(&chunk.to_string());
        body.push_str("\n\n");
    }
    body.push_str("data: [DONE]\n\n");
    Reply::raw(format!(
        "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
        body.len()
    ))
}

fn timings() -> serde_json::Value {
    serde_json::json!({"cache_n": 0, "prompt_n": 10, "predicted_n": 1})
}

fn sse_text(text: &str) -> Reply {
    sse_chunks(&[
        serde_json::json!({
            "choices": [{"finish_reason": null, "index": 0, "delta": {"content": text}}],
            "id": "t",
            "timings": timings()
        }),
        serde_json::json!({
            "choices": [{"finish_reason": "stop", "index": 0, "delta": {}}],
            "id": "t",
            "timings": timings()
        }),
    ])
}

fn sse_call_tool(tool: &str, arguments: serde_json::Value) -> Reply {
    let args = serde_json::json!({"name": tool, "arguments": arguments}).to_string();
    sse_named("call_tool", &args)
}

fn sse_named(name: &str, arguments: &str) -> Reply {
    sse_chunks(&[
        serde_json::json!({
            "choices": [{
                "finish_reason": null,
                "index": 0,
                "delta": {
                    "tool_calls": [{
                        "index": 0,
                        "id": "c1",
                        "type": "function",
                        "function": {"name": name, "arguments": arguments}
                    }]
                }
            }],
            "id": "t",
            "timings": timings()
        }),
        serde_json::json!({
            "choices": [{"finish_reason": "tool_calls", "index": 0, "delta": {}}],
            "id": "t",
            "timings": timings()
        }),
    ])
}

fn world(replies: Vec<Reply>) -> (Home, FakeServer, Arc<Context>) {
    let home = Home::new();
    home.write("compact.md", "Summarise.\n");
    let server = FakeServer::start();
    let cfg = home.config(&server.socket);
    server.route(CHAT, replies);
    let ctx = Arc::new(Context::new(
        cfg.clone(),
        Client::new(cfg),
        Box::new(ScriptedPort::new(vec![])),
        Registry::m5c(),
        Box::new(NoMemory),
    ));
    (home, server, ctx)
}

fn runtime(ctx: &Context) -> Runtime<'_> {
    Runtime {
        cfg: &ctx.cfg,
        client: &ctx.client,
        port: &*ctx.port,
        registry: &ctx.registry,
        ctx: Some(ctx),
    }
}

fn parent(ctx: &Context, id: &str) -> Session {
    let baseline = Baseline::assemble(&ctx.cfg, &ctx.registry).unwrap();
    Session::create(
        &ctx.cfg.paths.home,
        SessionId::new(id).unwrap(),
        baseline,
        ctx.cfg.slots.main,
    )
    .unwrap()
}

fn last_tool_text(session: &Session) -> String {
    for record in session.records().iter().rev() {
        if let LogRecord::ToolResult { content, .. } = record {
            return content.clone();
        }
    }
    panic!("no tool result");
}

#[test]
fn a_child_runs_on_the_background_slot_and_the_parent_prefix_still_extends() {
    let (home, server, ctx) = world(vec![
        sse_call_tool("subagent", serde_json::json!({"task": "look around"})),
        sse_text("the drawer is empty"),
        sse_text("ok"),
    ]);
    let mut session = parent(&ctx, "mm-thread");
    let rt = runtime(&ctx);
    let mut events = Vec::new();
    run_turn(&mut session, &rt, "search", None, &mut |e| {
        events.push(e.clone())
    })
    .unwrap();

    let sent = server.requests_to(CHAT);
    assert_eq!(sent.len(), 3);
    assert_eq!(sent[0].json()["id_slot"], ctx.cfg.slots.main);
    assert_eq!(sent[1].json()["id_slot"], ctx.cfg.slots.background);
    assert_eq!(sent[2].json()["id_slot"], ctx.cfg.slots.main);

    let first = sent[0].json()["messages"].as_array().unwrap().clone();
    let third = sent[2].json()["messages"].as_array().unwrap().clone();
    assert!(third.len() > first.len());
    assert_eq!(&third[..first.len()], &first[..]);
    let tool = third[first.len() + 1].clone();
    assert_eq!(tool["role"], "tool");
    assert!(
        tool["content"]
            .as_str()
            .unwrap()
            .starts_with("sub-mm-thread-1 completed\n"),
        "{tool}"
    );
    assert!(
        tool["content"]
            .as_str()
            .unwrap()
            .contains("the drawer is empty"),
        "{tool}"
    );

    let child_body = sent[1].json();
    let child_msgs = child_body["messages"].as_array().unwrap();
    let last = child_msgs.last().unwrap();
    assert_eq!(last["role"], "user");
    let content = last["content"].as_str().unwrap();
    assert!(content.starts_with("It is "), "{content}");
    assert!(content.contains("look around"), "{content}");

    let child = Session::open(&home.dir, SessionId::new("sub-mm-thread-1").unwrap()).unwrap();
    assert_eq!(child.slot(), ctx.cfg.slots.background);
}

#[test]
fn a_second_spawn_in_the_same_parent_uses_the_next_call_id() {
    let (_home, _server, ctx) = world(vec![
        sse_call_tool("subagent", serde_json::json!({"task": "one"})),
        sse_text("first"),
        sse_call_tool("subagent", serde_json::json!({"task": "two"})),
        sse_text("second"),
        sse_text("ok"),
    ]);
    let mut session = parent(&ctx, "mm-thread");
    let rt = runtime(&ctx);
    let mut events = Vec::new();
    run_turn(&mut session, &rt, "both", None, &mut |e| {
        events.push(e.clone())
    })
    .unwrap();
    let texts: Vec<String> = session
        .records()
        .iter()
        .filter_map(|r| match r {
            LogRecord::ToolResult { content, .. } => Some(content.clone()),
            _ => None,
        })
        .collect();
    assert!(
        texts[0].starts_with("sub-mm-thread-1 completed\n"),
        "{texts:?}"
    );
    assert!(
        texts[1].starts_with("sub-mm-thread-2 completed\n"),
        "{texts:?}"
    );
}

#[test]
fn an_overflowing_parent_id_is_invalid_arguments() {
    let parent_id = "a".repeat(59);
    let (_home, server, ctx) = world(vec![
        sse_call_tool("subagent", serde_json::json!({"task": "x"})),
        sse_text("ok"),
    ]);
    let mut session = parent(&ctx, &parent_id);
    let rt = runtime(&ctx);
    let mut events = Vec::new();
    run_turn(&mut session, &rt, "go", None, &mut |e| {
        events.push(e.clone())
    })
    .unwrap();
    let text = last_tool_text(&session);
    assert!(
        text.contains("would not fit") && text.contains("subagent-id"),
        "{text}"
    );
    assert_eq!(server.requests_to(CHAT).len(), 2);
}

#[test]
fn a_nested_child_returns_through_the_middle_tool_result() {
    let (home, server, ctx) = world(vec![
        sse_call_tool("subagent", serde_json::json!({"task": "outer"})),
        sse_call_tool("subagent", serde_json::json!({"task": "inner"})),
        sse_text("leaf"),
        sse_text("middle"),
        sse_text("ok"),
    ]);
    let mut session = parent(&ctx, "mm-thread");
    let rt = runtime(&ctx);
    let mut events = Vec::new();
    run_turn(&mut session, &rt, "nest", None, &mut |e| {
        events.push(e.clone())
    })
    .unwrap();

    let middle = Session::open(&home.dir, SessionId::new("sub-mm-thread-1").unwrap()).unwrap();
    let leaf = Session::open(&home.dir, SessionId::new("sub-sub-mm-thread-1-1").unwrap()).unwrap();
    assert_eq!(leaf.slot(), ctx.cfg.slots.background);
    let middle_text = last_tool_text(&middle);
    assert!(
        middle_text.starts_with("sub-sub-mm-thread-1-1 completed\n"),
        "{middle_text}"
    );
    assert!(middle_text.contains("leaf"), "{middle_text}");
    let parent_text = last_tool_text(&session);
    assert!(
        parent_text.starts_with("sub-mm-thread-1 completed\n"),
        "{parent_text}"
    );
    assert!(parent_text.contains("middle"), "{parent_text}");

    let sent = server.requests_to(CHAT);
    assert_eq!(sent[0].json()["id_slot"], 0);
    assert_eq!(sent[1].json()["id_slot"], 1);
    assert_eq!(sent[2].json()["id_slot"], 1);
    let parent_first = sent[0].json()["messages"].as_array().unwrap().clone();
    let parent_last = sent[4].json()["messages"].as_array().unwrap().clone();
    assert_eq!(&parent_last[..parent_first.len()], &parent_first[..]);
}

#[test]
fn a_child_that_ends_on_a_tool_call_is_failed() {
    let home = Home::new();
    home.write("compact.md", "Summarise.\n");
    let server = FakeServer::start();
    let mut cfg = home.config(&server.socket);
    cfg.r#loop.tool_iterations = 1;
    server.route(
        CHAT,
        vec![
            sse_call_tool("subagent", serde_json::json!({"task": "loop"})),
            sse_named("clock", "{}"),
            sse_named("clock", "{}"),
            sse_text("ok"),
        ],
    );
    let ctx = Arc::new(Context::new(
        cfg.clone(),
        Client::new(cfg),
        Box::new(ScriptedPort::new(vec![])),
        Registry::m5c(),
        Box::new(NoMemory),
    ));
    let mut session = parent(&ctx, "mm-thread");
    let rt = runtime(&ctx);
    let mut events = Vec::new();
    run_turn(&mut session, &rt, "go", None, &mut |e| {
        events.push(e.clone())
    })
    .unwrap();
    let text = last_tool_text(&session);
    assert_eq!(text, "sub-mm-thread-1 failed");
    drop(home);
}