mod mcp;

use mcp::{read_message, send_response, send_error, JsonRpcResponse};
use serde_json::{json, Value};
use tokio::net::UnixStream;
use tokio::io::{AsyncReadExt, AsyncWriteExt};

#[tokio::main]
async fn main() {
    loop {
        let msg = match read_message().await {
            Some(m) => m,
            None => break,
        };

        tokio::spawn(async move {
            let id = msg.id.clone().unwrap_or(json!(null));

            match msg.method.as_str() {
                "initialize" => {
                    send_response(JsonRpcResponse {
                        jsonrpc: "2.0".to_string(),
                        id,
                        result: Some(json!({
                            "capabilities": {},
                            "serverInfo": {
                                "name": "mcp-memory-linux-nvim",
                                "version": "0.1.0"
                            }
                        })),
                        error: None,
                    }).await;
                }
                "tools/list" => {
                    send_response(JsonRpcResponse {
                        jsonrpc: "2.0".to_string(),
                        id,
                        result: Some(json!({
                            "tools": [
                                {
                                    "name": "nvim_goto_line",
                                    "description": "Open a file and jump to a specific line",
                                    "inputSchema": {
                                        "type": "object",
                                        "properties": {
                                            "file": { "type": "string" },
                                            "line": { "type": "integer" }
                                        },
                                        "required": ["file", "line"]
                                    }
                                },
                                {
                                    "name": "nvim_get_active_buffer",
                                    "description": "Get the contents of the currently active Neovim buffer",
                                    "inputSchema": {
                                        "type": "object",
                                        "properties": {}
                                    }
                                },
                                {
                                    "name": "nvim_get_cursor",
                                    "description": "Get the current cursor position (line and column) in the active Neovim buffer",
                                    "inputSchema": {
                                        "type": "object",
                                        "properties": {}
                                    }
                                },
                                {
                                    "name": "nvim_get_visual_selection",
                                    "description": "Get the text that is currently highlighted or was last highlighted in Visual mode",
                                    "inputSchema": {
                                        "type": "object",
                                        "properties": {}
                                    }
                                },
                                {
                                    "name": "nvim_set_diagnostics",
                                    "description": "Push a diagnostic message (like an LSP warning) to a specific line in the active buffer",
                                    "inputSchema": {
                                        "type": "object",
                                        "properties": {
                                            "line": { "type": "integer" },
                                            "message": { "type": "string" }
                                        },
                                        "required": ["line", "message"]
                                    }
                                }
                            ]
                        })),
                        error: None,
                    }).await;
                }
                "tools/call" => {
                    let params = msg.params.clone().unwrap_or(json!({}));
                    let name = params.get("name").and_then(|v| v.as_str()).unwrap_or("");
                    let args = params.get("arguments").cloned().unwrap_or(json!({}));

                    match name {
                        "nvim_goto_line" => {
                            let file = args.get("file").and_then(|v| v.as_str()).unwrap_or("");
                            let line = args.get("line").and_then(|v| v.as_i64()).unwrap_or(1);

                            let cmd = format!("edit {} | {} | normal! zz", file, line);
                            match send_nvim_command(&cmd).await {
                                Ok(_) => {
                                    send_response(JsonRpcResponse {
                                        jsonrpc: "2.0".to_string(),
                                        id,
                                        result: Some(json!({
                                            "content": [
                                                { "type": "text", "text": format!("Successfully jumped to {}:{}", file, line) }
                                            ]
                                        })),
                                        error: None,
                                    }).await;
                                }
                                Err(e) => {
                                    send_error(id, -32603, &format!("Failed to execute command: {}", e)).await;
                                }
                            }
                        }
                        "nvim_get_active_buffer" => {
                            match get_nvim_active_buffer().await {
                                Ok(content) => {
                                    send_response(JsonRpcResponse {
                                        jsonrpc: "2.0".to_string(),
                                        id,
                                        result: Some(json!({
                                            "content": [
                                                { "type": "text", "text": content }
                                            ]
                                        })),
                                        error: None,
                                    }).await;
                                }
                                Err(e) => {
                                    send_error(id, -32603, &format!("Failed to get active buffer: {}", e)).await;
                                }
                            }
                        }
                        "nvim_get_cursor" => {
                            match get_nvim_cursor().await {
                                Ok(content) => {
                                    send_response(JsonRpcResponse {
                                        jsonrpc: "2.0".to_string(),
                                        id,
                                        result: Some(json!({
                                            "content": [
                                                { "type": "text", "text": content }
                                            ]
                                        })),
                                        error: None,
                                    }).await;
                                }
                                Err(e) => {
                                    send_error(id, -32603, &format!("Failed to get cursor: {}", e)).await;
                                }
                            }
                        }
                        "nvim_get_visual_selection" => {
                            match get_nvim_visual_selection().await {
                                Ok(content) => {
                                    send_response(JsonRpcResponse {
                                        jsonrpc: "2.0".to_string(),
                                        id,
                                        result: Some(json!({
                                            "content": [
                                                { "type": "text", "text": content }
                                            ]
                                        })),
                                        error: None,
                                    }).await;
                                }
                                Err(e) => {
                                    send_error(id, -32603, &format!("Failed to get visual selection: {}", e)).await;
                                }
                            }
                        }
                        "nvim_set_diagnostics" => {
                            let line = args.get("line").and_then(|v| v.as_i64()).unwrap_or(1);
                            let message = args.get("message").and_then(|v| v.as_str()).unwrap_or("");

                            match set_nvim_diagnostics(line, message).await {
                                Ok(_) => {
                                    send_response(JsonRpcResponse {
                                        jsonrpc: "2.0".to_string(),
                                        id,
                                        result: Some(json!({
                                            "content": [
                                                { "type": "text", "text": format!("Successfully set diagnostic on line {}", line) }
                                            ]
                                        })),
                                        error: None,
                                    }).await;
                                }
                                Err(e) => {
                                    send_error(id, -32603, &format!("Failed to set diagnostic: {}", e)).await;
                                }
                            }
                        }
                        _ => {
                            send_error(id, -32601, "Tool not found").await;
                        }
                    }
                }
                _ => {
                    // Ignore other methods
                }
            }
        });
    }
}

async fn get_socket_path() -> Result<String, String> {
    // 1. Try active_nvim.txt first
    if let Ok(home) = std::env::var("HOME") {
        let path = format!("{}/.gemini/active_nvim.txt", home);
        if let Ok(content) = std::fs::read_to_string(&path) {
            let p = content.trim().to_string();
            if !p.is_empty() {
                return Ok(p);
            }
        }
    }
    
    // 2. Fallback: search /tmp/nvim.*/0
    if let Ok(entries) = std::fs::read_dir("/tmp") {
        for entry in entries.flatten() {
            if let Ok(name) = entry.file_name().into_string() {
                if name.starts_with("nvim.") {
                    let path = entry.path().join("0");
                    if path.exists() {
                        return Ok(path.to_string_lossy().to_string());
                    }
                }
            }
        }
    }
    Err("Could not find Neovim socket".to_string())
}

async fn call_nvim(req: rmpv::Value) -> Result<rmpv::Value, String> {
    let msgid = if let rmpv::Value::Array(ref arr) = req {
        if arr.len() > 1 { arr[1].clone() } else { rmpv::Value::Nil }
    } else { rmpv::Value::Nil };

    let socket_path = get_socket_path().await?;
    let mut stream = UnixStream::connect(socket_path).await.map_err(|e| e.to_string())?;
    
    let mut buf = Vec::new();
    rmpv::encode::write_value(&mut buf, &req).map_err(|e| e.to_string())?;
    stream.write_all(&buf).await.map_err(|e| e.to_string())?;
    
    let mut resp_buf = Vec::new();
    let mut chunk = vec![0u8; 8192];
    let mut offset = 0;

    loop {
        let mut cursor = std::io::Cursor::new(&resp_buf[offset..]);
        match rmpv::decode::read_value(&mut cursor) {
            Ok(val) => {
                offset += cursor.position() as usize;
                
                if let rmpv::Value::Array(ref arr) = val {
                    if arr.len() >= 4 && arr[0] == rmpv::Value::Integer(1.into()) && arr[1] == msgid {
                        return Ok(val);
                    }
                }
                continue;
            },
            Err(_) => {
                let n = stream.read(&mut chunk).await.map_err(|e| e.to_string())?;
                if n == 0 {
                    return Err("Unexpected EOF from Neovim".to_string());
                }
                resp_buf.extend_from_slice(&chunk[..n]);
            }
        }
    }
}

async fn send_nvim_command(cmd: &str) -> Result<(), String> {
    use rmpv::Value as RmpValue;
    let req = RmpValue::Array(vec![
        RmpValue::Integer(0.into()),
        RmpValue::Integer(1.into()), // msgid
        RmpValue::String("nvim_command".into()),
        RmpValue::Array(vec![RmpValue::String(cmd.into())]),
    ]);
    
    let resp = call_nvim(req).await?;
    if let RmpValue::Array(arr) = resp {
        if !arr[2].is_nil() {
            return Err(format!("Neovim error: {:?}", arr[2]));
        }
        return Ok(());
    }
    Err("Invalid response".to_string())
}

async fn get_nvim_active_buffer() -> Result<String, String> {
    use rmpv::Value as RmpValue;
    let req = RmpValue::Array(vec![
        RmpValue::Integer(0.into()),
        RmpValue::Integer(2.into()), // msgid
        RmpValue::String("nvim_buf_get_lines".into()),
        RmpValue::Array(vec![
            RmpValue::Integer(0.into()),
            RmpValue::Integer(0.into()),
            RmpValue::Integer((-1).into()),
            RmpValue::Boolean(true),
        ]),
    ]);
    
    let resp = call_nvim(req).await?;
    if let RmpValue::Array(arr) = resp {
        if !arr[2].is_nil() {
            return Err(format!("Neovim error: {:?}", arr[2]));
        }
        if let RmpValue::Array(lines) = &arr[3] {
            let mut text = String::new();
            for line in lines {
                if let RmpValue::String(s) = line {
                    if let Some(s) = s.as_str() {
                        text.push_str(s);
                        text.push('\n');
                    }
                }
            }
            return Ok(text);
        }
    }
    Err("Invalid response".to_string())
}

async fn get_nvim_cursor() -> Result<String, String> {
    use rmpv::Value as RmpValue;
    let req = RmpValue::Array(vec![
        RmpValue::Integer(0.into()),
        RmpValue::Integer(3.into()), // msgid
        RmpValue::String("nvim_win_get_cursor".into()),
        RmpValue::Array(vec![
            RmpValue::Integer(0.into()),
        ]),
    ]);
    
    let resp = call_nvim(req).await?;
    if let RmpValue::Array(arr) = resp {
        if !arr[2].is_nil() {
            return Err(format!("Neovim error: {:?}", arr[2]));
        }
        if let RmpValue::Array(pos) = &arr[3] {
            if pos.len() == 2 {
                if let (RmpValue::Integer(row), RmpValue::Integer(col)) = (&pos[0], &pos[1]) {
                    return Ok(format!("Line: {}, Column: {}", row, col));
                }
            }
        }
    }
    Err("Invalid response".to_string())
}

async fn get_nvim_visual_selection() -> Result<String, String> {
    let lua_script = r#"
        local _, csrow, cscol, _ = unpack(vim.fn.getpos("'<"))
        local _, cerow, cecol, _ = unpack(vim.fn.getpos("'>"))
        local lines = vim.fn.getline(csrow, cerow)
        if type(lines) == "table" then
            return table.concat(lines, "\n")
        else
            return lines
        end
    "#;
    
    use rmpv::Value as RmpValue;
    let req = RmpValue::Array(vec![
        RmpValue::Integer(0.into()),
        RmpValue::Integer(4.into()), // msgid
        RmpValue::String("nvim_exec_lua".into()),
        RmpValue::Array(vec![
            RmpValue::String(lua_script.into()),
            RmpValue::Array(vec![]),
        ]),
    ]);
    
    let resp = call_nvim(req).await?;
    if let RmpValue::Array(arr) = resp {
        if !arr[2].is_nil() {
            return Err(format!("Neovim error: {:?}", arr[2]));
        }
        if let RmpValue::String(s) = &arr[3] {
            if let Some(text) = s.as_str() {
                return Ok(text.to_string());
            }
        }
    }
    Err("Invalid response".to_string())
}

async fn set_nvim_diagnostics(line: i64, message: &str) -> Result<(), String> {
    let escaped_message = message.replace("\\", "\\\\").replace("\"", "\\\"");
    let lua_script = format!(r#"
        local ns = vim.api.nvim_create_namespace("gemini_diagnostics")
        local diagnostics = {{{{
            lnum = {} - 1,
            col = 0,
            severity = vim.diagnostic.severity.WARN,
            message = "{}",
        }}}}
        vim.diagnostic.set(ns, 0, diagnostics, {{}})
    "#, line, escaped_message);
    
    use rmpv::Value as RmpValue;
    let req = RmpValue::Array(vec![
        RmpValue::Integer(0.into()),
        RmpValue::Integer(5.into()), // msgid
        RmpValue::String("nvim_exec_lua".into()),
        RmpValue::Array(vec![
            RmpValue::String(lua_script.into()),
            RmpValue::Array(vec![]),
        ]),
    ]);
    
    let resp = call_nvim(req).await?;
    if let RmpValue::Array(arr) = resp {
        if !arr[2].is_nil() {
            return Err(format!("Neovim error: {:?}", arr[2]));
        }
        return Ok(());
    }
    Err("Invalid response".to_string())
}