use serde_json::{Value, json};
use std::io::{BufRead, Write};

use crate::protocol::{
    ERR_INVALID_PARAMS, ERR_METHOD_NOT_FOUND, ERR_PARSE, InitializeResult, MCP_PROTOCOL_VERSION,
    Request, Response, RpcError, ServerCapabilities, ServerInfo, ToolContent, ToolsCallParams,
    ToolsCallResult, ToolsCapability, ToolsListResult,
};
use crate::tools;

const SERVER_NAME: &str = "second-brain";
const SERVER_VERSION: &str = env!("CARGO_PKG_VERSION");

pub struct Server {
    client: reqwest::blocking::Client,
    api_base: String,
    initialized: bool,
}

impl Server {
    pub fn new(client: reqwest::blocking::Client, api_base: String) -> Self {
        Self {
            client,
            api_base,
            initialized: false,
        }
    }

    pub fn handle(&mut self, req: Request) -> Option<Response> {
        let id = match &req.id {
            Some(v) => v.clone(),
            None => {
                if req.method == "notifications/initialized" {
                    self.initialized = true;
                }
                return None;
            }
        };

        if !self.initialized && req.method != "initialize" {
            return Some(Response::err(
                id,
                ERR_INVALID_PARAMS,
                "server not initialized",
            ));
        }

        match req.method.as_str() {
            "initialize" => Some(self.on_initialize(id)),
            "tools/list" => Some(self.on_tools_list(id)),
            "tools/call" => Some(self.on_tools_call(id, req.params)),
            "ping" => Some(Response::ok(id, json!({}))),
            _ => Some(Response::err(
                id,
                ERR_METHOD_NOT_FOUND,
                format!("unknown method: {}", req.method),
            )),
        }
    }

    fn on_initialize(&mut self, id: Value) -> Response {
        self.initialized = true;
        let result = InitializeResult {
            protocol_version: MCP_PROTOCOL_VERSION,
            capabilities: ServerCapabilities {
                tools: ToolsCapability {
                    list_changed: false,
                },
            },
            server_info: ServerInfo {
                name: SERVER_NAME,
                version: SERVER_VERSION,
            },
        };
        Response::ok(id, serde_json::to_value(result).expect("serialize init"))
    }

    fn on_tools_list(&self, id: Value) -> Response {
        let result = ToolsListResult {
            tools: tools::descriptors(),
        };
        Response::ok(id, serde_json::to_value(result).expect("serialize tools"))
    }

    fn on_tools_call(&self, id: Value, params: Value) -> Response {
        let parsed: ToolsCallParams = match serde_json::from_value(params) {
            Ok(p) => p,
            Err(e) => return Response::err(id, ERR_INVALID_PARAMS, e.to_string()),
        };

        let result = tools::dispatch(
            &parsed.name,
            &parsed.arguments,
            &self.client,
            &self.api_base,
        );

        match result {
            Ok(tool_result) => Response::ok(
                id,
                serde_json::to_value(tool_result).expect("serialize result"),
            ),
            Err(e) => {
                let tool_result = ToolsCallResult {
                    content: vec![ToolContent::Text {
                        text: format!("error: {e}"),
                    }],
                    is_error: Some(true),
                };
                Response::ok(
                    id,
                    serde_json::to_value(tool_result).expect("serialize error"),
                )
            }
        }
    }
}

pub fn run_stdio(client: reqwest::blocking::Client, api_base: String) -> std::io::Result<()> {
    let stdin = std::io::stdin();
    let stdout = std::io::stdout();
    let mut out = stdout.lock();
    let mut server = Server::new(client, api_base);

    for line in stdin.lock().lines() {
        let line = line?;
        if line.trim().is_empty() {
            continue;
        }
        match serde_json::from_str::<Request>(&line) {
            Ok(req) => {
                if let Some(resp) = server.handle(req) {
                    let msg = serde_json::to_string(&resp).map_err(|e| {
                        std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string())
                    })?;
                    out.write_all(msg.as_bytes())?;
                    out.write_all(b"\n")?;
                    out.flush()?;
                }
            }
            Err(e) => {
                let resp = Response {
                    jsonrpc: "2.0",
                    id: Value::Null,
                    result: None,
                    error: Some(RpcError {
                        code: ERR_PARSE,
                        message: format!("parse error: {e}"),
                    }),
                };
                let msg = serde_json::to_string(&resp).map_err(|e| {
                    std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string())
                })?;
                out.write_all(msg.as_bytes())?;
                out.write_all(b"\n")?;
                out.flush()?;
            }
        }
    }
    Ok(())
}
