//! End-to-end test: drive the real binary's `exec` mode against a mock
//! OpenAI-compatible SSE server. Verifies the full loop: model call →
//! tool_call → shell execution → tool result → final answer.

use std::io::{Read, Write};
use std::net::TcpListener;
use std::process::Command;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

fn sse_chunk(json: &str) -> String {
    format!("data: {json}\n\n")
}

/// Minimal HTTP/1.1 server speaking just enough chat-completions SSE for the
/// agent. First request returns a run_shell_command tool call that writes a
/// marker file; second request returns a plain completion.
fn spawn_mock_server(marker_path: String) -> (std::net::SocketAddr, Arc<AtomicUsize>) {
    let listener = TcpListener::bind("127.0.0.1:0").expect("bind");
    let addr = listener.local_addr().unwrap();
    let hits = Arc::new(AtomicUsize::new(0));
    let hits2 = hits.clone();

    std::thread::spawn(move || {
        for stream in listener.incoming() {
            let Ok(mut stream) = stream else { break };
            let n = hits2.fetch_add(1, Ordering::SeqCst);
            let marker = marker_path.clone();
            std::thread::spawn(move || {
                // Read request until end of headers, then body by content-length.
                let mut buf = Vec::new();
                let mut tmp = [0u8; 4096];
                let header_end;
                loop {
                    let r = stream.read(&mut tmp).unwrap_or(0);
                    if r == 0 {
                        return;
                    }
                    buf.extend_from_slice(&tmp[..r]);
                    if let Some(pos) = find_subsequence(&buf, b"\r\n\r\n") {
                        header_end = pos + 4;
                        break;
                    }
                }
                let headers = String::from_utf8_lossy(&buf[..header_end]).to_string();
                let content_length: usize = headers
                    .lines()
                    .find_map(|l| {
                        let (k, v) = l.split_once(':')?;
                        if k.eq_ignore_ascii_case("content-length") {
                            v.trim().parse().ok()
                        } else {
                            None
                        }
                    })
                    .unwrap_or(0);
                while buf.len() < header_end + content_length {
                    let r = stream.read(&mut tmp).unwrap_or(0);
                    if r == 0 {
                        break;
                    }
                    buf.extend_from_slice(&tmp[..r]);
                }

                let body = if n == 0 {
                    // Tool call: write marker file via the shell tool.
                    let args = serde_json_escape(&format!(
                        "{{\"command\": \"echo e2e-ok > {marker}\", \"description\": \"write marker\"}}"
                    ));
                    let mut s = String::new();
                    s.push_str(&sse_chunk(&format!(
                        "{{\"choices\":[{{\"delta\":{{\"role\":\"assistant\",\"content\":\"Writing marker.\"}},\"index\":0}}]}}"
                    )));
                    s.push_str(&sse_chunk(&format!(
                        "{{\"choices\":[{{\"delta\":{{\"tool_calls\":[{{\"index\":0,\"id\":\"call_1\",\"type\":\"function\",\"function\":{{\"name\":\"run_shell_command\",\"arguments\":{args}}}}}]}},\"index\":0}}]}}"
                    )));
                    s.push_str(&sse_chunk(
                        "{\"choices\":[{\"delta\":{},\"index\":0,\"finish_reason\":\"tool_calls\"}],\"usage\":{\"prompt_tokens\":100,\"completion_tokens\":20}}",
                    ));
                    s.push_str("data: [DONE]\n\n");
                    s
                } else {
                    let mut s = String::new();
                    s.push_str(&sse_chunk(
                        "{\"choices\":[{\"delta\":{\"content\":\"Marker file written and verified. Task complete.\"},\"index\":0}]}",
                    ));
                    s.push_str(&sse_chunk(
                        "{\"choices\":[{\"delta\":{},\"index\":0,\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":150,\"completion_tokens\":12}}",
                    ));
                    s.push_str("data: [DONE]\n\n");
                    s
                };

                let resp = format!(
                    "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
                    body.len(),
                    body
                );
                let _ = stream.write_all(resp.as_bytes());
                let _ = stream.flush();
            });
        }
    });

    (addr, hits)
}

fn find_subsequence(haystack: &[u8], needle: &[u8]) -> Option<usize> {
    haystack.windows(needle.len()).position(|w| w == needle)
}

/// Wrap a raw string as a JSON string literal.
fn serde_json_escape(s: &str) -> String {
    let mut out = String::from("\"");
    for c in s.chars() {
        match c {
            '"' => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            c => out.push(c),
        }
    }
    out.push('"');
    out
}

#[test]
fn exec_mode_runs_tool_call_end_to_end() {
    let dir = tempfile::tempdir().expect("tempdir");
    let marker = dir.path().join("marker.txt");
    let (addr, hits) = spawn_mock_server(marker.display().to_string());

    let bin = env!("CARGO_BIN_EXE_lucid-train");
    let out = Command::new(bin)
        .args([
            "exec",
            "write the marker file",
            "--base-url",
            &format!("http://{addr}"),
            "--api-key",
            "test-key",
            "--model",
            "mock-model",
            "--cwd",
            &dir.path().display().to_string(),
        ])
        .env("LUCID_TRAIN_HOME", dir.path())
        .output()
        .expect("binary runs");

    let stdout = String::from_utf8_lossy(&out.stdout);
    let stderr = String::from_utf8_lossy(&out.stderr);
    assert!(
        out.status.success(),
        "exit={:?}\nstdout:\n{stdout}\nstderr:\n{stderr}",
        out.status.code()
    );

    // The shell tool actually ran and produced the marker file.
    let marker_content = std::fs::read_to_string(&marker).expect("marker file written by tool");
    assert_eq!(marker_content.trim(), "e2e-ok");

    // Final assistant message made it to stdout.
    assert!(stdout.contains("Task complete"), "stdout:\n{stdout}");
    // Both model calls happened (tool round + final round).
    assert!(hits.load(Ordering::SeqCst) >= 2);
    // Exec cell rendering shows the command run.
    assert!(stdout.contains("echo e2e-ok"), "stdout:\n{stdout}");
}

#[test]
fn exec_mode_blocks_catastrophic_commands() {
    let dir = tempfile::tempdir().expect("tempdir");
    // Server that always asks for `rm -rf /` then finishes.
    let (addr, _hits) = spawn_mock_server_blocked();

    let bin = env!("CARGO_BIN_EXE_lucid-train");
    let out = Command::new(bin)
        .args([
            "exec",
            "destroy everything",
            "--base-url",
            &format!("http://{addr}"),
            "--api-key",
            "test-key",
            "--model",
            "mock-model",
            "--cwd",
            &dir.path().display().to_string(),
        ])
        .env("LUCID_TRAIN_HOME", dir.path())
        .output()
        .expect("binary runs");

    let stdout = String::from_utf8_lossy(&out.stdout);
    // The dangerous command must be blocked, and the run must still complete.
    assert!(stdout.contains("blocked"), "stdout:\n{stdout}");
    assert!(out.status.success());
}

fn spawn_mock_server_blocked() -> (std::net::SocketAddr, Arc<AtomicUsize>) {
    let listener = TcpListener::bind("127.0.0.1:0").expect("bind");
    let addr = listener.local_addr().unwrap();
    let hits = Arc::new(AtomicUsize::new(0));
    let hits2 = hits.clone();

    std::thread::spawn(move || {
        for stream in listener.incoming() {
            let Ok(mut stream) = stream else { break };
            let n = hits2.fetch_add(1, Ordering::SeqCst);
            std::thread::spawn(move || {
                let mut tmp = [0u8; 65536];
                let mut total = Vec::new();
                loop {
                    let r = stream.read(&mut tmp).unwrap_or(0);
                    if r == 0 {
                        break;
                    }
                    total.extend_from_slice(&tmp[..r]);
                    if find_subsequence(&total, b"\r\n\r\n").is_some() {
                        // naive: assume whole body arrives promptly
                        if total.len() > 100 {
                            break;
                        }
                    }
                }
                let body = if n == 0 {
                    let mut s = String::new();
                    s.push_str(&sse_chunk(
                        "{\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"type\":\"function\",\"function\":{\"name\":\"run_shell_command\",\"arguments\":\"{\\\"command\\\": \\\"rm -rf /\\\"}\"}}]},\"index\":0}]}",
                    ));
                    s.push_str(&sse_chunk(
                        "{\"choices\":[{\"delta\":{},\"index\":0,\"finish_reason\":\"tool_calls\"}]}",
                    ));
                    s.push_str("data: [DONE]\n\n");
                    s
                } else {
                    let mut s = String::new();
                    s.push_str(&sse_chunk(
                        "{\"choices\":[{\"delta\":{\"content\":\"Understood, stopping.\"},\"index\":0}]}",
                    ));
                    s.push_str(&sse_chunk(
                        "{\"choices\":[{\"delta\":{},\"index\":0,\"finish_reason\":\"stop\"}]}",
                    ));
                    s.push_str("data: [DONE]\n\n");
                    s
                };
                let resp = format!(
                    "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
                    body.len(),
                    body
                );
                let _ = stream.write_all(resp.as_bytes());
            });
        }
    });

    (addr, hits)
}
