//! SHCortex MCP Server — Model Context Protocol over stdio (JSON-RPC 2.0).
//! Exposes: compress, retrieve, prompt, mem_status.

use std::io::{BufRead, BufReader, Write};
use std::sync::Mutex;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use shcortex::Shcortex;

#[derive(Debug, Deserialize)]
struct RpcRequest {
    #[serde(default)]
    id: Option<Value>,
    method: String,
    #[serde(default)]
    params: Option<Value>,
}

#[derive(Debug, Serialize)]
struct RpcResponse {
    jsonrpc: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    id: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    result: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    error: Option<RpcError>,
}

#[derive(Debug, Serialize)]
struct RpcError { code: i32, message: String }

fn tool_definitions() -> Vec<Value> {
    vec![
        json!({"name":"compress","description":"Store text into SHCortex memory.","inputSchema":{"type":"object","properties":{"text":{"type":"string","description":"Text to compress"}},"required":["text"]}}),
        json!({"name":"retrieve","description":"Query persistent memory via FTS5+LIKE.","inputSchema":{"type":"object","properties":{"question":{"type":"string","description":"Search query"}},"required":["question"]}}),
        json!({"name":"prompt","description":"Build structured system+user prompt.","inputSchema":{"type":"object","properties":{"question":{"type":"string","description":"Question to build prompt for"}},"required":["question"]}}),
        json!({"name":"mem_status","description":"RAM and SQLite memory stats.","inputSchema":{"type":"object","properties":{}}}),
    ]
}

fn handle(state: &Mutex<Shcortex>, req: RpcRequest) -> RpcResponse {
    match req.method.as_str() {
        "initialize" => {
            RpcResponse {
                jsonrpc: "2.0".into(), id: req.id,
                result: Some(json!({"protocolVersion":"2024-11-05","serverInfo":{"name":"shcortex-mcp","version":"0.2.0"},"capabilities":{"tools":{}}})),
                error: None,
            }
        },
        "notifications/initialized" => RpcResponse { jsonrpc: "2.0".into(), id: None, result: Some(Value::Null), error: None },
        "tools/list" => {
            RpcResponse {
                jsonrpc: "2.0".into(), id: req.id,
                result: Some(json!({"tools": tool_definitions()})),
                error: None,
            }
        },
        "tools/call" => {
            let params = req.params.unwrap_or(Value::Null);
            let name = params.get("name").and_then(|n| n.as_str()).unwrap_or("");
            let args = params.get("arguments").cloned().unwrap_or(Value::Null);
            let mut s = state.lock().unwrap();
            let result = match name {
                "compress" => {
                    let text = args.get("text").and_then(|t| t.as_str()).unwrap_or("");
                    s.compress(text);
                    json!({"memcons": s.memcons().len(), "atinks": s.atinks().len()})
                }
                "retrieve" => {
                    let q = args.get("question").and_then(|t| t.as_str()).unwrap_or("");
                    json!({"context": s.retrieve(q)})
                }
                "prompt" => {
                    let q = args.get("question").and_then(|t| t.as_str()).unwrap_or("");
                    let (system, user) = s.build_prompt(q);
                    json!({"system": system, "user": user})
                }
                "mem_status" => {
                    let st = s.status();
                    json!({"ram_memcons":st.ram_memcons,"ram_atinks":st.ram_atinks,"db_memcons":st.db_memcons,"db_atinks":st.db_atinks,"db_hectores":st.db_hectores,"db_size_kb":st.db_size_kb})
                }
                _ => json!({"error": format!("Unknown tool: {name}")}),
            };
            RpcResponse {
                jsonrpc: "2.0".into(), id: req.id,
                result: Some(json!({"content":[{"type":"text","text":result.to_string()}]})),
                error: None,
            }
        },
        _ => RpcResponse {
            jsonrpc: "2.0".into(), id: req.id,
            result: None,
            error: Some(RpcError { code: -32601, message: format!("Method not found: {}", req.method) }),
        },
    }
}

fn main() {
    let db = std::env::var("SHCORTEX_DB").unwrap_or_else(|_| {
        let home = std::env::var("HOME").or_else(|_| std::env::var("USERPROFILE")).unwrap_or_default();
        format!("{home}/.shcortex/memory.db")
    });
    let state = Mutex::new(Shcortex::new(Some(db)));
    let stdin = std::io::stdin();
    let reader = BufReader::new(stdin.lock());
    let stdout = std::io::stdout();
    for line in reader.lines() {
        let line = match line { Ok(l) => l, Err(_) => break };
        if line.trim().is_empty() { continue; }
        let req: RpcRequest = match serde_json::from_str(&line) {
            Ok(r) => r,
            Err(e) => {
                let err = RpcResponse { jsonrpc: "2.0".into(), id: None, result: None, error: Some(RpcError { code: -32700, message: format!("Parse error: {e}") }) };
                let _ = writeln!(stdout.lock(), "{}", serde_json::to_string(&err).unwrap());
                let _ = stdout.lock().flush();
                continue;
            }
        };
        let is_notif = req.id.is_none() && req.method.starts_with("notifications/");
        let resp = handle(&state, req);
        if !is_notif && resp.id.is_some() {
            let _ = writeln!(stdout.lock(), "{}", serde_json::to_string(&resp).unwrap());
            let _ = stdout.lock().flush();
        }
    }
}
