use std::path::PathBuf;
use std::sync::atomic::{AtomicI64, Ordering};

use anyhow::{anyhow, Context, Result};
use rmcp::{
    model::{
        CallToolRequest, CallToolRequestParams, CallToolResult, ClientJsonRpcMessage,
        ClientRequest, JsonRpcRequest, NumberOrString, ServerJsonRpcMessage, ServerResult,
    },
    service::serve_directly,
    service::RoleServer,
    transport::OneshotTransport,
};
use serde_json::{Map, Value};

use crate::tools::LeanCtxServer;

pub struct ContextEngine {
    server: LeanCtxServer,
    next_id: AtomicI64,
}

impl ContextEngine {
    pub fn new() -> Self {
        Self {
            server: LeanCtxServer::new(),
            next_id: AtomicI64::new(1),
        }
    }

    pub fn with_project_root(project_root: impl Into<PathBuf>) -> Self {
        let root = project_root.into().to_string_lossy().to_string();
        Self {
            server: LeanCtxServer::new_with_project_root(Some(&root)),
            next_id: AtomicI64::new(1),
        }
    }

    pub fn from_server(server: LeanCtxServer) -> Self {
        Self {
            server,
            next_id: AtomicI64::new(1),
        }
    }

    pub fn server(&self) -> &LeanCtxServer {
        &self.server
    }

    pub fn manifest(&self) -> Value {
        crate::core::mcp_manifest::manifest_value()
    }

    pub async fn call_tool_value(&self, name: &str, arguments: Option<Value>) -> Result<Value> {
        let result = self.call_tool_result(name, arguments).await?;
        serde_json::to_value(result).map_err(|e| anyhow!("serialize CallToolResult: {e}"))
    }

    pub async fn call_tool_result(
        &self,
        name: &str,
        arguments: Option<Value>,
    ) -> Result<CallToolResult> {
        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
        let req_id = NumberOrString::Number(id);

        let args_obj: Map<String, Value> = match arguments {
            None => Map::new(),
            Some(Value::Object(m)) => m,
            Some(other) => {
                return Err(anyhow!(
                    "tool arguments must be a JSON object (got {other})"
                ))
            }
        };

        let params = CallToolRequestParams::new(name.to_string()).with_arguments(args_obj);
        let call: CallToolRequest = CallToolRequest::new(params);
        let client_req = ClientRequest::CallToolRequest(call);
        let msg = ClientJsonRpcMessage::Request(JsonRpcRequest::new(req_id, client_req));

        let (transport, mut rx) = OneshotTransport::<RoleServer>::new(msg);
        let service = serve_directly(self.server.clone(), transport, None);
        tokio::spawn(async move {
            let _ = service.waiting().await;
        });

        let server_msg =
            match tokio::time::timeout(std::time::Duration::from_mins(2), rx.recv()).await {
                Ok(Some(msg)) => msg,
                Ok(None) => return Err(anyhow!("no response from tool call")),
                Err(_) => return Err(anyhow!("tool call timed out after 120s")),
            };

        match server_msg {
            ServerJsonRpcMessage::Response(r) => match r.result {
                ServerResult::CallToolResult(result) => Ok(result),
                other => Err(anyhow!("unexpected server result: {other:?}")),
            },
            ServerJsonRpcMessage::Error(e) => Err(anyhow!("{e:?}")).context("tool call error"),
            ServerJsonRpcMessage::Notification(_) => Err(anyhow!("unexpected notification")),
            ServerJsonRpcMessage::Request(_) => Err(anyhow!("unexpected request")),
        }
    }

    pub async fn call_tool_text(&self, name: &str, arguments: Option<Value>) -> Result<String> {
        let result = self.call_tool_result(name, arguments).await?;
        let mut out = String::new();
        for c in result.content {
            if let Some(t) = c.as_text() {
                out.push_str(&t.text);
            }
        }
        if out.is_empty() {
            if let Some(v) = result.structured_content {
                out = v.to_string();
            }
        }
        Ok(out)
    }
}

impl Default for ContextEngine {
    fn default() -> Self {
        Self::new()
    }
}
