mod mcp;

use mcp::{read_message, send_response, send_error, JsonRpcResponse};
use serde_json::json;
use tokio::net::windows::named_pipe::ClientOptions;
use tokio::io::{AsyncReadExt, AsyncWriteExt};

fn init_logging(app_name: &str) -> Option<tracing_appender::non_blocking::WorkerGuard> {
    let base_dir = std::env::var("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| {
        dirs::home_dir()
            .map(|mut h| {
                h.push(".gemini/mcp_memory");
                h.to_string_lossy().to_string()
            })
            .unwrap_or_else(|| ".gemini/mcp_memory".into())
    });
    let log_dir = std::path::PathBuf::from(base_dir).join("logs");
    std::fs::create_dir_all(&log_dir).unwrap_or_default();
    
    let file_appender = tracing_appender::rolling::daily(log_dir, format!("{}.log", app_name));
    let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
    
    let _ = tracing_subscriber::fmt()
        .with_writer(non_blocking)
        .with_ansi(false)
        .with_max_level(tracing::Level::INFO)
        .try_init();
    
    Some(guard)
}

#[tokio::main]
async fn main() {
    if std::env::args().any(|a| a == "--version" || a == "-V") {
        println!("mcp-memory-win-nvim {}", env!("APP_VERSION"));
        return;
    }
    
    let _guard = init_logging("win-nvim");
    tracing::info!("win-nvim MCP server started");
    let mut stdin = tokio::io::BufReader::new(tokio::io::stdin());
    loop {
        let msg = match read_message(&mut stdin).await {
            Some(m) => {
                tracing::info!("Received message method: {}", m.method);
                m
            },
            None => {
                tracing::info!("Stdin closed, exiting loop");
                break;
            }
        };

        tokio::spawn(async move {
            let id = msg.id.clone().unwrap_or(json!(null));
            let _start_time = std::time::Instant::now();

            match msg.method.as_str() {
                "initialize" => {
                    send_response(JsonRpcResponse {
                        jsonrpc: "2.0".to_string(),
                        id,
                        result: Some(json!({
                            "protocolVersion": "2024-11-05",
                            "capabilities": {
                                "tools": {}
                            },
                            "serverInfo": {
                                "name": "mcp-memory-win-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"]
                                    }
                                },
                                {
                                    "name": "nvim_execute_lua",
                                    "description": "Execute arbitrary Lua code in Neovim and return the result (JSON serialized).",
                                    "inputSchema": {
                                        "type": "object",
                                        "properties": {
                                            "code": { "type": "string" }
                                        },
                                        "required": ["code"]
                                    }
                                },
                                {
                                    "name": "nvim_list_buffers",
                                    "description": "Get a list of all loaded Neovim buffers and their IDs.",
                                    "inputSchema": {
                                        "type": "object",
                                        "properties": {}
                                    }
                                },
                                {
                                    "name": "nvim_get_diagnostics",
                                    "description": "Get all LSP diagnostics (errors, warnings) for the active buffer.",
                                    "inputSchema": {
                                        "type": "object",
                                        "properties": {}
                                    }
                                }
                            ]
                        })),
                        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;
                                }
                            }
                        }
                        
                        "nvim_execute_lua" => {
                            let code = args.get("code").and_then(|v| v.as_str()).unwrap_or("");
                            match execute_nvim_lua(code).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 execute lua: {}", e)).await; }
                            }
                        }
                        "nvim_list_buffers" => {
                            let code = r#"
                                local bufs = vim.api.nvim_list_bufs()
                                local loaded = {}
                                for _, b in ipairs(bufs) do
                                    if vim.api.nvim_buf_is_loaded(b) then
                                        local name = vim.api.nvim_buf_get_name(b)
                                        table.insert(loaded, {id = b, name = name == "" and "[No Name]" or name})
                                    end
                                end
                                return loaded
                            "#;
                            match execute_nvim_lua(code).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 list buffers: {}", e)).await; }
                            }
                        }
                        "nvim_get_diagnostics" => {
                            let code = r#"
                                local diags = vim.diagnostic.get(0)
                                local res = {}
                                for _, d in ipairs(diags) do
                                    table.insert(res, {
                                        line = d.lnum + 1,
                                        col = d.col,
                                        message = d.message,
                                        severity = d.severity
                                    })
                                end
                                return res
                            "#;
                            match execute_nvim_lua(code).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 diagnostics: {}", e)).await; }
                            }
                        }
                        _ => {
                            send_error(id, -32601, "Tool not found").await;
                        }
                    }
                }
                _ => {
                    // Ignore other methods
                }
            }
        });
    }
}

async fn get_socket_path() -> Result<String, String> {
    // 1. Primary: Use the active_nvim.txt which is updated by Neovim's BufEnter telemetry
    let profile = std::env::var("USERPROFILE").unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".into());
    let path = format!("{}\\.gemini\\active_nvim.txt", profile);
    
    if let Ok(content) = std::fs::read_to_string(&path) {
        let p = content.trim().to_string();
        if !p.is_empty() {
            // It might be a full pipe path or just the name. If it's just the name, prepend \\.\pipe\
            if p.starts_with(r"\\.\pipe\") {
                return Ok(p);
            } else if p.starts_with("nvim.") {
                return Ok(format!(r"\\.\pipe\{}", p));
            } else {
                // Some other servername format? Try it as is.
                return Ok(p);
            }
        }
    }

    // 2. Fallback to auto-discovery in \\.\pipe\ (only if single instance is running)
    tracing::warn!("active_nvim.txt missing or invalid, falling back to pipe discovery");
    if let Ok(dir) = std::fs::read_dir(r"\\.\pipe\") {
        for entry in dir.flatten() {
            let name = entry.file_name();
            let name_str = name.to_string_lossy();
            if name_str.starts_with("nvim.") {
                return Ok(format!(r"\\.\pipe\{}", name_str));
            }
        }
    }
    
    Err("Could not find active Windows Neovim named pipe".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 };

    tracing::info!("Connecting to neovim pipe");
    let socket_path = get_socket_path().await?;
    let mut client = ClientOptions::new().open(&socket_path).map_err(|e| e.to_string())?;
    
    let mut buf = Vec::new();
    rmpv::encode::write_value(&mut buf, &req).map_err(|e| e.to_string())?;
    tracing::info!("Sending RPC request to neovim (msgid: {})", msgid);
    client.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 {
                        tracing::info!("Received RPC response from neovim (msgid: {})", msgid);
                        return Ok(val);
                    }
                }
                continue;
            },
            Err(_) => {
                let read_future = client.read(&mut chunk);
                match tokio::time::timeout(tokio::time::Duration::from_secs(5), read_future).await {
                    Ok(Ok(n)) => {
                        if n == 0 { return Err("Connection closed".into()); }
                        resp_buf.extend_from_slice(&chunk[..n]);
                    }
                    Ok(Err(e)) => return Err(e.to_string()),
                    Err(_) => {
                        tracing::error!("Timeout waiting for Neovim response (msgid: {})", msgid);
                        return Err("Timeout waiting for Neovim response".into());
                    }
                }
            }
        }
    }
}

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())
}

fn rmpv_to_json(val: &rmpv::Value) -> serde_json::Value {
    match val {
        rmpv::Value::Nil => serde_json::Value::Null,
        rmpv::Value::Boolean(b) => serde_json::json!(b),
        rmpv::Value::Integer(i) => {
            if let Some(n) = i.as_i64() {
                serde_json::json!(n)
            } else if let Some(n) = i.as_u64() {
                serde_json::json!(n)
            } else {
                serde_json::Value::Null
            }
        },
        rmpv::Value::F32(f) => serde_json::json!(f),
        rmpv::Value::F64(f) => serde_json::json!(f),
        rmpv::Value::String(s) => {
            if let Some(str_val) = s.as_str() {
                serde_json::json!(str_val)
            } else {
                serde_json::Value::Null
            }
        },
        rmpv::Value::Array(arr) => {
            let vec: Vec<serde_json::Value> = arr.iter().map(rmpv_to_json).collect();
            serde_json::Value::Array(vec)
        },
        rmpv::Value::Map(map) => {
            let mut obj = serde_json::Map::new();
            for (k, v) in map {
                let key_str = if let rmpv::Value::String(s) = k {
                    s.as_str().unwrap_or("").to_string()
                } else {
                    format!("{:?}", k)
                };
                obj.insert(key_str, rmpv_to_json(v));
            }
            serde_json::Value::Object(obj)
        },
        _ => serde_json::json!(format!("{:?}", val)),
    }
}

async fn execute_nvim_lua(code: &str) -> Result<String, String> {
    use rmpv::Value as RmpValue;
    let req = RmpValue::Array(vec![
        RmpValue::Integer(0.into()),
        RmpValue::Integer(6.into()), // msgid
        RmpValue::String("nvim_exec_lua".into()),
        RmpValue::Array(vec![
            RmpValue::String(code.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 arr.len() > 3 {
            return Ok(serde_json::to_string_pretty(&rmpv_to_json(&arr[3])).unwrap_or_default());
        }
        return Ok("".to_string());
    }
    Err("Invalid response".to_string())
}