use std::env;
use std::fs;
use std::io::{BufRead, BufReader, Write};
use std::path::{Path, PathBuf};
use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio};

use serde_json::{json, Value};

const PROTOCOL_VERSION: &str = "2025-11-25";

struct ChioEdge {
    child: Child,
    stdin: ChildStdin,
    stdout: BufReader<ChildStdout>,
}

impl Drop for ChioEdge {
    fn drop(&mut self) {
        let _ = self.child.kill();
        let _ = self.child.wait();
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let project_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    let demo_name = env::args().nth(1).unwrap_or_else(|| "Chio".to_string());
    let receipt_db = project_root.join(".chio").join("receipts.db");
    let target_dir = env::var_os("CARGO_TARGET_DIR")
        .map(PathBuf::from)
        .unwrap_or_else(|| project_root.join(".chio").join("target"));
    fs::create_dir_all(receipt_db.parent().expect("receipt db parent"))?;
    fs::create_dir_all(&target_dir)?;

    let chio_bin = env::var("CHIO_BIN").unwrap_or_else(|_| "chio".to_string());
    let mut edge = spawn_chio_edge(&chio_bin, &project_root, &receipt_db, &target_dir)?;

    send_json(
        &mut edge.stdin,
        &json!({
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {
                "protocolVersion": PROTOCOL_VERSION,
                "capabilities": {},
                "clientInfo": {
                    "name": "{{PACKAGE_NAME}}-demo",
                    "version": "0.1.0"
                }
            }
        }),
    )?;
    let initialize = read_json(&mut edge.stdout)?;
    println!("initialize:");
    println!("{}", serde_json::to_string_pretty(&initialize)?);

    send_json(
        &mut edge.stdin,
        &json!({
            "jsonrpc": "2.0",
            "method": "notifications/initialized"
        }),
    )?;

    send_json(
        &mut edge.stdin,
        &json!({
            "jsonrpc": "2.0",
            "id": 2,
            "method": "tools/list",
            "params": {}
        }),
    )?;
    let tools = read_json(&mut edge.stdout)?;
    println!();
    println!("tools:");
    println!("{}", serde_json::to_string_pretty(&tools)?);

    send_json(
        &mut edge.stdin,
        &json!({
            "jsonrpc": "2.0",
            "id": 3,
            "method": "tools/call",
            "params": {
                "name": "hello_world",
                "arguments": {
                    "name": demo_name
                }
            }
        }),
    )?;
    let tool_call = read_json(&mut edge.stdout)?;
    let greeting = tool_call
        .get("result")
        .and_then(|result| result.get("structuredContent"))
        .and_then(|structured| structured.get("greeting"))
        .and_then(Value::as_str)
        .ok_or("tool result did not include a greeting")?;
    println!();
    println!("tool_call:");
    println!("{}", serde_json::to_string_pretty(&tool_call)?);
    println!();
    println!("governed greeting: {greeting}");

    drop(edge);

    let receipts = Command::new(&chio_bin)
        .arg("--receipt-db")
        .arg(&receipt_db)
        .arg("receipt")
        .arg("list")
        .arg("--limit")
        .arg("1")
        .current_dir(&project_root)
        .output()?;
    if !receipts.status.success() {
        return Err(format!(
            "failed to read receipt store: {}",
            String::from_utf8_lossy(&receipts.stderr)
        )
        .into());
    }

    println!();
    println!("latest receipt:");
    println!("{}", String::from_utf8_lossy(&receipts.stdout).trim());

    Ok(())
}

fn spawn_chio_edge(
    chio_bin: &str,
    project_root: &Path,
    receipt_db: &Path,
    target_dir: &Path,
) -> Result<ChioEdge, Box<dyn std::error::Error>> {
    let mut child = Command::new(chio_bin)
        .arg("--receipt-db")
        .arg(receipt_db)
        .arg("mcp")
        .arg("serve")
        .arg("--policy")
        .arg("policy.yaml")
        .arg("--server-id")
        .arg("hello")
        .arg("cargo")
        .arg("run")
        .arg("--quiet")
        .arg("--bin")
        .arg("hello_server")
        .current_dir(project_root)
        .env("CARGO_TARGET_DIR", target_dir)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::inherit())
        .spawn()?;

    let stdin = child.stdin.take().ok_or("failed to capture Chio stdin")?;
    let stdout = child.stdout.take().ok_or("failed to capture Chio stdout")?;

    Ok(ChioEdge {
        child,
        stdin,
        stdout: BufReader::new(stdout),
    })
}

fn send_json(stdin: &mut ChildStdin, value: &Value) -> Result<(), Box<dyn std::error::Error>> {
    writeln!(stdin, "{}", serde_json::to_string(value)?)?;
    stdin.flush()?;
    Ok(())
}

fn read_json(stdout: &mut BufReader<ChildStdout>) -> Result<Value, Box<dyn std::error::Error>> {
    let mut line = String::new();
    loop {
        line.clear();
        let bytes = stdout.read_line(&mut line)?;
        if bytes == 0 {
            return Err("Chio edge closed stdout before responding".into());
        }
        if line.trim().is_empty() {
            continue;
        }
        return Ok(serde_json::from_str(line.trim())?);
    }
}
