use serde_json::{json, Value};
use std::collections::HashMap;
use std::env;
use std::io::{self, BufRead, Read, Write};
use std::time::{SystemTime, UNIX_EPOCH};

fn now_ms() -> i64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_millis() as i64)
        .unwrap_or(0)
}

fn sql_escape(value: &str) -> String {
    value.replace('\\', "\\\\").replace('\'', "\\'")
}

fn to_string_vec(value: &Value) -> Vec<String> {
    match value {
        Value::Array(items) => items
            .iter()
            .filter_map(|item| item.as_str())
            .map(|s| s.trim().to_string())
            .filter(|s| !s.is_empty())
            .collect(),
        Value::String(raw) => raw
            .split(',')
            .map(|part| part.trim().to_string())
            .filter(|s| !s.is_empty())
            .collect(),
        _ => vec![],
    }
}

struct ManticoreClient {
    base_url: String,
    table: String,
    timeout_secs: u64,
}

impl ManticoreClient {
    fn from_env() -> Self {
        let base_url = env::var("MANTICORE_HTTP_URL")
            .unwrap_or_else(|_| "http://127.0.0.1:9308".to_string())
            .trim_end_matches('/')
            .to_string();
        let table = env::var("MANTICORE_MEMORY_TABLE")
            .unwrap_or_else(|_| "gia_memory_items".to_string())
            .trim()
            .to_string();
        let timeout_secs = env::var("MANTICORE_HTTP_TIMEOUT")
            .ok()
            .and_then(|v| v.parse::<u64>().ok())
            .unwrap_or(5);
        Self {
            base_url,
            table,
            timeout_secs,
        }
    }

    fn sql(&self, query: &str) -> Result<Value, String> {
        let endpoint = format!("{}/sql", self.base_url);
        let response = ureq::post(&endpoint)
            .timeout(std::time::Duration::from_secs(self.timeout_secs))
            .set("Content-Type", "application/x-www-form-urlencoded")
            .send_string(&format!(
                "mode=raw&query={}",
                urlencoding::encode(query).into_owned()
            ))
            .map_err(|err| err.to_string())?;
        let body = response
            .into_string()
            .map_err(|err| format!("manticore response read failed: {err}"))?;
        serde_json::from_str::<Value>(&body)
            .map_err(|err| format!("manticore response parse failed: {err}"))
    }

    fn ensure_table(&self) -> Result<(), String> {
        let query = format!(
            "CREATE TABLE IF NOT EXISTS {} (id BIGINT,memory_uuid STRING,user_id BIGINT,conversation_id STRING,memory_kind STRING,status STRING,updated_ts BIGINT,summary TEXT,body TEXT)",
            self.table
        );
        self.sql(&query).map(|_| ())
    }
}

fn tool_specs() -> Value {
    json!([
        {
            "name": "manticore.status",
            "description": "Report Manticore connectivity and table status.",
            "inputSchema": {
                "type": "object",
                "additionalProperties": false,
                "properties": {}
            }
        },
        {
            "name": "manticore.query",
            "description": "Run fast full-text retrieval against the Manticore memory table.",
            "inputSchema": {
                "type": "object",
                "additionalProperties": false,
                "required": ["query"],
                "properties": {
                    "query": {"type": "string"},
                    "user_id": {"type": "integer"},
                    "conversation_id": {"type": "string"},
                    "statuses": {
                        "oneOf": [
                            {"type": "string"},
                            {"type": "array", "items": {"type": "string"}}
                        ]
                    },
                    "limit": {"type": "integer", "minimum": 1, "maximum": 100}
                }
            }
        },
        {
            "name": "manticore.reindex",
            "description": "Run table maintenance operations for fast reads (flush + optimize).",
            "inputSchema": {
                "type": "object",
                "additionalProperties": false,
                "properties": {
                    "flush_ramchunk": {"type": "boolean"},
                    "optimize": {"type": "boolean"}
                }
            }
        }
    ])
}

fn call_tool(client: &ManticoreClient, name: &str, arguments: &Value) -> Result<Value, String> {
    match name {
        "manticore.status" => {
            client.ensure_table()?;
            let table_check = client.sql(&format!(
                "SHOW TABLES LIKE '{}'",
                sql_escape(&client.table)
            ))?;
            Ok(json!({
                "backend": "manticore",
                "ok": true,
                "manticore_http_url": client.base_url,
                "manticore_table": client.table,
                "table_check": table_check,
                "ts_ms": now_ms()
            }))
        }
        "manticore.query" => {
            client.ensure_table()?;
            let query = arguments
                .get("query")
                .and_then(Value::as_str)
                .unwrap_or("")
                .trim()
                .to_string();
            if query.is_empty() {
                return Err("query is required".to_string());
            }
            let limit = arguments
                .get("limit")
                .and_then(Value::as_i64)
                .unwrap_or(20)
                .clamp(1, 100);
            let mut where_parts = vec![format!("MATCH('{}')", sql_escape(&query))];
            if let Some(user_id) = arguments.get("user_id").and_then(Value::as_i64) {
                where_parts.push(format!("user_id={user_id}"));
            }
            let conversation_id = arguments
                .get("conversation_id")
                .and_then(Value::as_str)
                .unwrap_or("")
                .trim()
                .to_string();
            if !conversation_id.is_empty() {
                where_parts.push(format!(
                    "conversation_id='{}'",
                    sql_escape(&conversation_id)
                ));
            }
            let statuses = to_string_vec(arguments.get("statuses").unwrap_or(&Value::Null));
            if !statuses.is_empty() {
                let joined = statuses
                    .iter()
                    .map(|s| format!("'{}'", sql_escape(s)))
                    .collect::<Vec<_>>()
                    .join(",");
                where_parts.push(format!("status IN ({joined})"));
            }
            let sql = format!(
                "SELECT memory_uuid,memory_kind,status,conversation_id,updated_ts,summary,WEIGHT() AS score FROM {} WHERE {} ORDER BY score DESC LIMIT {}",
                client.table,
                where_parts.join(" AND "),
                limit
            );
            let payload = client.sql(&sql)?;
            let count = payload
                .get("data")
                .and_then(Value::as_array)
                .map(|rows| rows.len())
                .unwrap_or(0);
            Ok(json!({
                "backend": "manticore",
                "query": query,
                "count": count,
                "hits": payload.get("data").cloned().unwrap_or_else(|| json!([])),
                "raw": payload
            }))
        }
        "manticore.reindex" => {
            client.ensure_table()?;
            let flush = arguments
                .get("flush_ramchunk")
                .and_then(Value::as_bool)
                .unwrap_or(true);
            let optimize = arguments
                .get("optimize")
                .and_then(Value::as_bool)
                .unwrap_or(true);
            let mut actions: Vec<Value> = Vec::new();
            if flush {
                let sql = format!("FLUSH RAMCHUNK {}", client.table);
                let payload = client.sql(&sql)?;
                actions.push(json!({"sql": sql, "result": payload}));
            }
            if optimize {
                let sql = format!("OPTIMIZE TABLE {}", client.table);
                let payload = client.sql(&sql)?;
                actions.push(json!({"sql": sql, "result": payload}));
            }
            Ok(json!({
                "ok": true,
                "actions": actions,
                "ts_ms": now_ms()
            }))
        }
        _ => Err(format!("Unknown tool: {name}")),
    }
}

fn response(id: Value, result: Value) -> Value {
    json!({"jsonrpc":"2.0","id":id,"result":result})
}

fn error(id: Value, code: i32, message: &str) -> Value {
    json!({"jsonrpc":"2.0","id":id,"error":{"code":code,"message":message}})
}

fn write_message(payload: &Value, compat_newline_mode: bool) -> Result<(), String> {
    let raw = serde_json::to_vec(payload).map_err(|e| e.to_string())?;
    let mut stdout = io::stdout();
    if compat_newline_mode {
        stdout
            .write_all(format!("{}\n", String::from_utf8_lossy(&raw)).as_bytes())
            .map_err(|e| e.to_string())?;
    } else {
        stdout
            .write_all(format!("Content-Length: {}\r\n\r\n", raw.len()).as_bytes())
            .map_err(|e| e.to_string())?;
        stdout.write_all(&raw).map_err(|e| e.to_string())?;
    }
    stdout.flush().map_err(|e| e.to_string())
}

fn read_message(
    stdin: &mut io::StdinLock<'_>,
    compat_newline_mode: &mut bool,
) -> Result<Option<Value>, String> {
    let mut headers: HashMap<String, String> = HashMap::new();
    let mut pending_body: Vec<u8> = Vec::new();
    loop {
        let mut line: Vec<u8> = Vec::new();
        let bytes = stdin.read_until(b'\n', &mut line).map_err(|e| e.to_string())?;
        if bytes == 0 {
            return Ok(None);
        }
        let trimmed = line
            .iter()
            .copied()
            .skip_while(|b| b.is_ascii_whitespace())
            .collect::<Vec<u8>>();
        if headers.is_empty() && (trimmed.starts_with(b"{") || trimmed.starts_with(b"[")) {
            *compat_newline_mode = true;
            let raw = String::from_utf8_lossy(&line).trim().to_string();
            let parsed = serde_json::from_str::<Value>(&raw).map_err(|e| e.to_string())?;
            return Ok(Some(parsed));
        }

        if line == b"\r\n" || line == b"\n" {
            break;
        }
        let decoded = String::from_utf8_lossy(&line).trim().to_string();
        if let Some((k, v)) = decoded.split_once(':') {
            headers.insert(k.trim().to_lowercase(), v.trim().to_string());
        }
    }

    if let Some(length_raw) = headers.get("content-length") {
        let length = length_raw
            .parse::<usize>()
            .map_err(|_| "invalid content-length".to_string())?;
        if length > 0 {
            pending_body.resize(length, 0);
            stdin.read_exact(&mut pending_body).map_err(|e| e.to_string())?;
            let parsed =
                serde_json::from_slice::<Value>(&pending_body).map_err(|e| e.to_string())?;
            return Ok(Some(parsed));
        }
        return Ok(None);
    }
    Ok(None)
}

fn main() {
    let client = ManticoreClient::from_env();
    let stdin = io::stdin();
    let mut locked = stdin.lock();
    let mut compat_newline_mode = false;

    loop {
        let message = match read_message(&mut locked, &mut compat_newline_mode) {
            Ok(Some(value)) => value,
            Ok(None) => return,
            Err(err) => {
                let _ = write_message(
                    &error(Value::Null, -32000, &format!("read failed: {err}")),
                    compat_newline_mode,
                );
                return;
            }
        };

        let id = message.get("id").cloned().unwrap_or(Value::Null);
        let method = message
            .get("method")
            .and_then(Value::as_str)
            .unwrap_or("")
            .to_string();
        let params = message.get("params").cloned().unwrap_or_else(|| json!({}));

        let response_payload = match method.as_str() {
            "notifications/initialized" => None,
            "initialize" => Some(response(
                id,
                json!({
                    "protocolVersion":"2025-06-18",
                    "serverInfo":{"name":"gia-manticore-mcp-rust","version":"0.1.0"},
                    "capabilities":{"tools":{}}
                }),
            )),
            "ping" => Some(response(id, json!({}))),
            "tools/list" => Some(response(id, json!({"tools": tool_specs()}))),
            "tools/call" => {
                let name = params
                    .get("name")
                    .and_then(Value::as_str)
                    .unwrap_or("")
                    .trim()
                    .to_string();
                let args = params
                    .get("arguments")
                    .cloned()
                    .unwrap_or_else(|| json!({}));
                match call_tool(&client, &name, &args) {
                    Ok(payload) => Some(response(
                        id,
                        json!({"isError":false,"content":[{"type":"text","text":payload.to_string()}]}),
                    )),
                    Err(err) => Some(response(
                        id,
                        json!({"isError":true,"content":[{"type":"text","text":json!({"error":err}).to_string()}]}),
                    )),
                }
            }
            _ => Some(error(id, -32601, &format!("Method not found: {method}"))),
        };

        if let Some(payload) = response_payload {
            if write_message(&payload, compat_newline_mode).is_err() {
                return;
            }
        }
    }
}
