use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::sync::Arc;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use cerebro::prelude::*;

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

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

async fn handle_request(engine: &MemoryEngine, req: RpcRequest) -> Option<RpcResponse> {
    let id = req.id.clone()?;

    match req.method.as_str() {
        "initialize" => {
            Some(RpcResponse {
                jsonrpc: "2.0".into(), id,
                result: Some(json!({
                    "protocolVersion": "2024-11-05",
                    "capabilities": { "tools": { "listChanged": false } },
                    "serverInfo": { "name": "cerebro-mcp", "version": "0.1.0" }
                })),
                error: None,
            })
        }
        "tools/list" => {
            Some(RpcResponse {
                jsonrpc: "2.0".into(), id,
                result: Some(json!({
                    "tools": [{
                        "name": "ingest_memory",
                        "description": "Ingests a document into the long term semantic memory.",
                        "inputSchema": {
                            "type": "object",
                            "properties": { "content": { "type": "string" } },
                            "required": ["content"]
                        }
                    }, {
                        "name": "search_memory",
                        "description": "Searches the semantic memory for relevant chunks.",
                        "inputSchema": {
                            "type": "object",
                            "properties": { "query": { "type": "string" }, "top_k": { "type": "number" } },
                            "required": ["query"]
                        }
                    }]
                })),
                error: None,
            })
        }
        "tools/call" => {
            let params = req.params.unwrap_or(json!({}));
            let name = params["name"].as_str().unwrap_or("");
            let args = params["arguments"].clone();

            match name {
                "ingest_memory" => {
                    let text = args["content"].as_str().unwrap_or("");
                    let doc = Document::new(text);
                    if let Err(e) = engine.ingest_document(doc).await {
                        return Some(RpcResponse { jsonrpc: "2.0".into(), id, result: None, error: Some(json!({"code": -32603, "message": e.to_string()})) });
                    }
                    Some(RpcResponse { jsonrpc: "2.0".into(), id, result: Some(json!({ "content": [{ "type": "text", "text": "Success" }] })), error: None })
                }
                "search_memory" => {
                    let query = args["query"].as_str().unwrap_or("");
                    let top_k = args["top_k"].as_u64().unwrap_or(5) as usize;
                    match engine.query(query, top_k).await {
                        Ok(results) => {
                            let mut txt = String::new();
                            for (node, score) in results { txt.push_str(&format!("Score {}: {}\n", score, node.chunk.text)); }
                            Some(RpcResponse { jsonrpc: "2.0".into(), id, result: Some(json!({ "content": [{ "type": "text", "text": txt }] })), error: None })
                        }
                        Err(e) => Some(RpcResponse { jsonrpc: "2.0".into(), id, result: None, error: Some(json!({"code": -32603, "message": e.to_string()})) })
                    }
                }
                _ => Some(RpcResponse { jsonrpc: "2.0".into(), id, result: None, error: Some(json!({"code": -32601, "message": "Method not found"})) })
            }
        }
        "notifications/initialized" => None,
        _ => Some(RpcResponse { jsonrpc: "2.0".into(), id, result: None, error: Some(json!({"code": -32601, "message": "Method not found"})) })
    }
}

#[tokio::main]
async fn main() {
    let chunker = Arc::new(RecursiveCharacterChunker::new(512, 50));
    
    let embedder: Arc<dyn Embedder> = if let Ok(api_key) = std::env::var("OPENAI_API_KEY") {
        Arc::new(OpenAIEmbedder::new(api_key, "text-embedding-3-small"))
    } else {
        Arc::new(MockEmbedder::new(1536))
    };

    let store = Arc::new(MemoryVectorStore::new());
    let engine = MemoryEngine::new(chunker, embedder, store);

    let stdin = tokio::io::stdin();
    let mut stdout = tokio::io::stdout();
    let mut reader = BufReader::new(stdin);
    let mut line = String::new();

    loop {
        line.clear();
        match reader.read_line(&mut line).await {
            Ok(0) => break,
            Ok(_) => {
                if let Ok(req) = serde_json::from_str::<RpcRequest>(&line) {
                    if let Some(resp) = handle_request(&engine, req).await {
                        let mut out = serde_json::to_string(&resp).unwrap();
                        out.push('\n');
                        if stdout.write_all(out.as_bytes()).await.is_err() { break; }
                        let _ = stdout.flush().await;
                    }
                }
            }
            Err(_) => break,
        }
    }
}
