//! Very small proof-of-concept request router for the MCP prototype server.

use mcp_types::CallToolRequestParams;
use mcp_types::CallToolResult;
use mcp_types::CallToolResultContent;
use mcp_types::ClientRequest;
use mcp_types::JSONRPCBatchRequest;
use mcp_types::JSONRPCBatchResponse;
use mcp_types::JSONRPCError;
use mcp_types::JSONRPCErrorError;
use mcp_types::JSONRPCMessage;
use mcp_types::JSONRPCNotification;
use mcp_types::JSONRPCRequest;
use mcp_types::JSONRPCResponse;
use mcp_types::ListToolsResult;
use mcp_types::ModelContextProtocolRequest;
use mcp_types::RequestId;
use mcp_types::ServerCapabilitiesTools;
use mcp_types::ServerNotification;
use mcp_types::TextContent;
use mcp_types::Tool;
use mcp_types::ToolInputSchema;
use mcp_types::JSONRPC_VERSION;
use serde_json::json;
use schemars::schema_for;
use tokio::task;

// Import types from codex-core.
use codex_core::codex_wrapper::init_codex;
use codex_core::config::Config as CodexConfig;
use codex_core::protocol::{Event, EventMsg};

// Helper to convert a Codex Event into an MCP JSON-RPC notification.
fn codex_event_to_notification(event: &Event) -> JSONRPCMessage {
    JSONRPCMessage::Notification(JSONRPCNotification {
        jsonrpc: JSONRPC_VERSION.into(),
        method: "codex/event".into(),
        params: Some(serde_json::to_value(event).expect("Event must serialize")),
    })
}
use tokio::sync::mpsc;

pub(crate) struct MessageProcessor {
    outgoing: mpsc::Sender<JSONRPCMessage>,
    initialized: bool,
}

impl MessageProcessor {
    /// Create a new `MessageProcessor`, retaining a handle to the outgoing
    /// `Sender` so handlers can enqueue messages to be written to stdout.
    pub(crate) fn new(outgoing: mpsc::Sender<JSONRPCMessage>) -> Self {
        Self {
            outgoing,
            initialized: false,
        }
    }

    pub(crate) fn process_request(&mut self, request: JSONRPCRequest) {
        // Hold on to the ID so we can respond.
        let request_id = request.id.clone();

        let client_request = match ClientRequest::try_from(request) {
            Ok(client_request) => client_request,
            Err(e) => {
                tracing::warn!("Failed to convert request: {e}");
                return;
            }
        };

        // Dispatch to a dedicated handler for each request type.
        match client_request {
            ClientRequest::InitializeRequest(params) => {
                self.handle_initialize(request_id, params);
            }
            ClientRequest::PingRequest(params) => {
                self.handle_ping(request_id, params);
            }
            ClientRequest::ListResourcesRequest(params) => {
                self.handle_list_resources(params);
            }
            ClientRequest::ListResourceTemplatesRequest(params) => {
                self.handle_list_resource_templates(params);
            }
            ClientRequest::ReadResourceRequest(params) => {
                self.handle_read_resource(params);
            }
            ClientRequest::SubscribeRequest(params) => {
                self.handle_subscribe(params);
            }
            ClientRequest::UnsubscribeRequest(params) => {
                self.handle_unsubscribe(params);
            }
            ClientRequest::ListPromptsRequest(params) => {
                self.handle_list_prompts(params);
            }
            ClientRequest::GetPromptRequest(params) => {
                self.handle_get_prompt(params);
            }
            ClientRequest::ListToolsRequest(params) => {
                self.handle_list_tools(request_id, params);
            }
            ClientRequest::CallToolRequest(params) => {
                self.handle_call_tool(request_id, params);
            }
            ClientRequest::SetLevelRequest(params) => {
                self.handle_set_level(params);
            }
            ClientRequest::CompleteRequest(params) => {
                self.handle_complete(params);
            }
        }
    }

    /// Handle a standalone JSON-RPC response originating from the peer.
    pub(crate) fn process_response(&mut self, response: JSONRPCResponse) {
        tracing::info!("<- response: {:?}", response);
    }

    /// Handle a fire-and-forget JSON-RPC notification.
    pub(crate) fn process_notification(&mut self, notification: JSONRPCNotification) {
        let server_notification = match ServerNotification::try_from(notification) {
            Ok(n) => n,
            Err(e) => {
                tracing::warn!("Failed to convert notification: {e}");
                return;
            }
        };

        // Similar to requests, route each notification type to its own stub
        // handler so additional logic can be implemented incrementally.
        match server_notification {
            ServerNotification::CancelledNotification(params) => {
                self.handle_cancelled_notification(params);
            }
            ServerNotification::ProgressNotification(params) => {
                self.handle_progress_notification(params);
            }
            ServerNotification::ResourceListChangedNotification(params) => {
                self.handle_resource_list_changed(params);
            }
            ServerNotification::ResourceUpdatedNotification(params) => {
                self.handle_resource_updated(params);
            }
            ServerNotification::PromptListChangedNotification(params) => {
                self.handle_prompt_list_changed(params);
            }
            ServerNotification::ToolListChangedNotification(params) => {
                self.handle_tool_list_changed(params);
            }
            ServerNotification::LoggingMessageNotification(params) => {
                self.handle_logging_message(params);
            }
        }
    }

    /// Handle a batch of requests and/or notifications.
    pub(crate) fn process_batch_request(&mut self, batch: JSONRPCBatchRequest) {
        tracing::info!("<- batch request containing {} item(s)", batch.len());
        for item in batch {
            match item {
                mcp_types::JSONRPCBatchRequestItem::JSONRPCRequest(req) => {
                    self.process_request(req);
                }
                mcp_types::JSONRPCBatchRequestItem::JSONRPCNotification(note) => {
                    self.process_notification(note);
                }
            }
        }
    }

    /// Handle an error object received from the peer.
    pub(crate) fn process_error(&mut self, err: JSONRPCError) {
        tracing::error!("<- error: {:?}", err);
    }

    /// Handle a batch of responses/errors.
    pub(crate) fn process_batch_response(&mut self, batch: JSONRPCBatchResponse) {
        tracing::info!("<- batch response containing {} item(s)", batch.len());
        for item in batch {
            match item {
                mcp_types::JSONRPCBatchResponseItem::JSONRPCResponse(resp) => {
                    self.process_response(resp);
                }
                mcp_types::JSONRPCBatchResponseItem::JSONRPCError(err) => {
                    self.process_error(err);
                }
            }
        }
    }

    fn handle_initialize(
        &mut self,
        id: RequestId,
        params: <mcp_types::InitializeRequest as ModelContextProtocolRequest>::Params,
    ) {
        tracing::info!("initialize -> params: {:?}", params);

        if self.initialized {
            // Already initialised: send JSON-RPC error response.
            let error_msg = JSONRPCMessage::Error(JSONRPCError {
                jsonrpc: JSONRPC_VERSION.into(),
                id,
                error: JSONRPCErrorError {
                    code: -32600, // Invalid Request
                    message: "initialize called more than once".to_string(),
                    data: None,
                },
            });

            if let Err(e) = self.outgoing.try_send(error_msg) {
                tracing::error!("Failed to send initialization error: {e}");
            }
            return;
        }

        self.initialized = true;

        // Build a minimal InitializeResult. Fill with placeholders.
        let result = mcp_types::InitializeResult {
            capabilities: mcp_types::ServerCapabilities {
                completions: None,
                experimental: None,
                logging: None,
                prompts: None,
                resources: None,
                tools: Some(ServerCapabilitiesTools {
                    list_changed: Some(true),
                }),
            },
            instructions: None,
            protocol_version: params.protocol_version.clone(),
            server_info: mcp_types::Implementation {
                name: "codex-mcp-server".to_string(),
                version: mcp_types::MCP_SCHEMA_VERSION.to_string(),
            },
        };

        self.send_response::<mcp_types::InitializeRequest>(id, result);
    }

    fn send_response<T>(&self, id: RequestId, result: T::Result)
    where
        T: ModelContextProtocolRequest,
    {
        let response = JSONRPCMessage::Response(JSONRPCResponse {
            jsonrpc: JSONRPC_VERSION.into(),
            id,
            result: serde_json::to_value(result).unwrap(),
        });

        if let Err(e) = self.outgoing.try_send(response) {
            tracing::error!("Failed to send response: {e}");
        }
    }

    fn handle_ping(
        &self,
        id: RequestId,
        params: <mcp_types::PingRequest as mcp_types::ModelContextProtocolRequest>::Params,
    ) {
        tracing::info!("ping -> params: {:?}", params);
        let result = json!({});
        self.send_response::<mcp_types::PingRequest>(id, result);
    }

    fn handle_list_resources(
        &self,
        params: <mcp_types::ListResourcesRequest as mcp_types::ModelContextProtocolRequest>::Params,
    ) {
        tracing::info!("resources/list -> params: {:?}", params);
    }

    fn handle_list_resource_templates(
        &self,
        params:
            <mcp_types::ListResourceTemplatesRequest as mcp_types::ModelContextProtocolRequest>::Params,
    ) {
        tracing::info!("resources/templates/list -> params: {:?}", params);
    }

    fn handle_read_resource(
        &self,
        params: <mcp_types::ReadResourceRequest as mcp_types::ModelContextProtocolRequest>::Params,
    ) {
        tracing::info!("resources/read -> params: {:?}", params);
    }

    fn handle_subscribe(
        &self,
        params: <mcp_types::SubscribeRequest as mcp_types::ModelContextProtocolRequest>::Params,
    ) {
        tracing::info!("resources/subscribe -> params: {:?}", params);
    }

    fn handle_unsubscribe(
        &self,
        params: <mcp_types::UnsubscribeRequest as mcp_types::ModelContextProtocolRequest>::Params,
    ) {
        tracing::info!("resources/unsubscribe -> params: {:?}", params);
    }

    fn handle_list_prompts(
        &self,
        params: <mcp_types::ListPromptsRequest as mcp_types::ModelContextProtocolRequest>::Params,
    ) {
        tracing::info!("prompts/list -> params: {:?}", params);
    }

    fn handle_get_prompt(
        &self,
        params: <mcp_types::GetPromptRequest as mcp_types::ModelContextProtocolRequest>::Params,
    ) {
        tracing::info!("prompts/get -> params: {:?}", params);
    }

    fn handle_list_tools(
        &self,
        id: RequestId,
        params: <mcp_types::ListToolsRequest as mcp_types::ModelContextProtocolRequest>::Params,
    ) {
        tracing::trace!("tools/list -> {params:?}");
        // -----------------------------------------------------------------
        // Build the schema for the Codex tool dynamically using `schemars`.
        // -----------------------------------------------------------------
        let root_schema = schema_for!(CodexConfig);
        let schema_value = serde_json::to_value(&root_schema).expect("schema serializable");

        // Attempt to extract `properties` and `required` from the generated schema.
        let (properties, required) = schema_value
            .get("schema")
            .map(|schema_root| {
                let props = schema_root.get("properties").cloned();
                let req = schema_root
                    .get("required")
                    .and_then(|v| serde_json::from_value::<Vec<String>>(v.clone()).ok());
                (props, req)
            })
            .unwrap_or((None, None));

        let result = ListToolsResult {
            tools: vec![Tool {
                name: "codex".to_string(),
                input_schema: ToolInputSchema {
                    r#type: "object".to_string(),
                    properties,
                    required,
                },
                description: Some(
                    "Run a Codex session. Accepts configuration parameters matching the Codex Config struct.".to_string(),
                ),
                annotations: None,
            }],
            next_cursor: None,
        };

        self.send_response::<mcp_types::ListToolsRequest>(id, result);
    }

    fn handle_call_tool(
        &self,
        id: RequestId,
        params: <mcp_types::CallToolRequest as mcp_types::ModelContextProtocolRequest>::Params,
    ) {
        tracing::info!("tools/call -> params: {:?}", params);
        let CallToolRequestParams { name, arguments } = params;

        // We only support the "codex" tool for now.
        if name != "codex" {
            // Tool not found – return error result so the LLM can react.
            let result = CallToolResult {
                content: vec![CallToolResultContent::TextContent(TextContent {
                    r#type: "text".to_string(),
                    text: format!("Unknown tool '{name}'"),
                    annotations: None,
                })],
                is_error: Some(true),
            };
            self.send_response::<mcp_types::CallToolRequest>(id, result);
            return;
        }

        // Clone outgoing sender to move into async task.
        let outgoing = self.outgoing.clone();

        // Spawn an async task to handle the Codex session so that we do not
        // block the synchronous message-processing loop.
        task::spawn(async move {
            // -----------------------------------------------------------------
            // Step 1: Parse configuration parameters.
            // -----------------------------------------------------------------
            let config: CodexConfig = match arguments {
                Some(json_val) => match serde_json::from_value::<CodexConfig>(json_val) {
                    Ok(cfg) => cfg,
                    Err(e) => {
                        let result = CallToolResult {
                            content: vec![CallToolResultContent::TextContent(TextContent {
                                r#type: "text".to_owned(),
                                text: format!(
                                    "Failed to parse configuration for Codex tool: {e}"
                                ),
                                annotations: None,
                            })],
                            is_error: Some(true),
                        };
                        let _ = outgoing
                            .send(JSONRPCMessage::Response(JSONRPCResponse {
                                jsonrpc: JSONRPC_VERSION.into(),
                                id: id.clone(),
                                result: result.into(),
                            }))
                            .await;
                        return;
                    }
                },
                None => match CodexConfig::load_with_overrides(Default::default()) {
                    Ok(cfg) => cfg,
                    Err(e) => {
                        let result = CallToolResult {
                            content: vec![CallToolResultContent::TextContent(TextContent {
                                r#type: "text".to_string(),
                                text: format!(
                                    "Cannot load default Codex configuration: {e}"
                                ),
                                annotations: None,
                            })],
                            is_error: Some(true),
                        };
                        let _ = outgoing
                            .send(JSONRPCMessage::Response(JSONRPCResponse {
                                jsonrpc: JSONRPC_VERSION.into(),
                                id: id.clone(),
                                result: result.into(),
                            }))
                            .await;
                        return;
                    }
                },
            };

            // -----------------------------------------------------------------
            // Step 2: Start Codex session.
            // -----------------------------------------------------------------
            let (codex, first_event, _ctrl_c) = match init_codex(config).await {
                Ok(res) => res,
                Err(e) => {
                    let result = CallToolResult {
                        content: vec![CallToolResultContent::TextContent(TextContent {
                            r#type: "text".to_string(),
                            text: format!("Failed to start Codex session: {e}"),
                            annotations: None,
                        })],
                        is_error: Some(true),
                    };
                    let _ = outgoing
                        .send(JSONRPCMessage::Response(JSONRPCResponse {
                            jsonrpc: JSONRPC_VERSION.into(),
                            id: id.clone(),
                            result: result.into(),
                        }))
                        .await;
                    return;
                }
            };

            // Send the initial SessionConfigured event as a notification so the
            // client can begin rendering.
            let _ = outgoing.send(codex_event_to_notification(&first_event)).await;

            // We'll track the last AgentMessage so we can fulfil the tool call
            // response when the task completes.
            let mut last_agent_message: Option<String> = None;

            // -----------------------------------------------------------------
            // Step 3: Pump events until we reach a state that requires a tool
            // response.
            // -----------------------------------------------------------------
            loop {
                match codex.next_event().await {
                    Ok(event) => {
                        // Forward all events to the MCP client.
                        let _ = outgoing.send(codex_event_to_notification(&event)).await;

                        match &event.msg {
                            EventMsg::AgentMessage { message } => {
                                last_agent_message = Some(message.clone());
                            }
                            EventMsg::ExecApprovalRequest { .. } => {
                                // Respond to the original call with an exec approval request.
                                let result = CallToolResult {
                                    content: vec![CallToolResultContent::TextContent(TextContent {
                                        r#type: "text".to_string(),
                                        text: "EXEC_APPROVAL_REQUIRED".to_string(),
                                        annotations: None,
                                    })],
                                    is_error: None,
                                };
                                let _ = outgoing
                                    .send(JSONRPCMessage::Response(JSONRPCResponse {
                                        jsonrpc: JSONRPC_VERSION.into(),
                                        id: id.clone(),
                                        result: result.into(),
                                    }))
                                    .await;
                                break;
                            }
                            EventMsg::ApplyPatchApprovalRequest { .. } => {
                                // Respond to the original call with a patch approval request.
                                let result = CallToolResult {
                                    content: vec![CallToolResultContent::TextContent(TextContent {
                                        r#type: "text".to_string(),
                                        text: "PATCH_APPROVAL_REQUIRED".to_string(),
                                        annotations: None,
                                    })],
                                    is_error: None,
                                };
                                let _ = outgoing
                                    .send(JSONRPCMessage::Response(JSONRPCResponse {
                                        jsonrpc: JSONRPC_VERSION.into(),
                                        id: id.clone(),
                                        result: result.into(),
                                    }))
                                    .await;
                                break;
                            }
                            EventMsg::TaskComplete => {
                                // Return the last agent message, if any.
                                let result = if let Some(msg) = last_agent_message {
                                    CallToolResult {
                                        content: vec![CallToolResultContent::TextContent(TextContent {
                                            r#type: "text".to_string(),
                                            text: msg,
                                            annotations: None,
                                        })],
                                        is_error: None,
                                    }
                                } else {
                                    CallToolResult {
                                        content: vec![CallToolResultContent::TextContent(TextContent {
                                            r#type: "text".to_string(),
                                            text: "<no-output>".to_string(),
                                            annotations: None,
                                        })],
                                        is_error: None,
                                    }
                                };

                                let _ = outgoing
                                    .send(JSONRPCMessage::Response(JSONRPCResponse {
                                        jsonrpc: JSONRPC_VERSION.into(),
                                        id: id.clone(),
                                        result: result.into(),
                                    }))
                                    .await;
                                break;
                            }
                            _ => {
                                // Nothing to do; continue pumping.
                            }
                        }
                    }
                    Err(e) => {
                        // Bubble up error to the user via the response.
                        let result = CallToolResult {
                            content: vec![CallToolResultContent::TextContent(TextContent {
                                r#type: "text".to_string(),
                                text: format!("Codex session error: {e}"),
                                annotations: None,
                            })],
                            is_error: Some(true),
                        };
                        let _ = outgoing
                            .send(JSONRPCMessage::Response(JSONRPCResponse {
                                jsonrpc: JSONRPC_VERSION.into(),
                                id: id.clone(),
                                result: result.into(),
                            }))
                            .await;
                        break;
                    }
                }
            }
        });
    }

    fn handle_set_level(
        &self,
        params: <mcp_types::SetLevelRequest as mcp_types::ModelContextProtocolRequest>::Params,
    ) {
        tracing::info!("logging/setLevel -> params: {:?}", params);
    }

    fn handle_complete(
        &self,
        params: <mcp_types::CompleteRequest as mcp_types::ModelContextProtocolRequest>::Params,
    ) {
        tracing::info!("completion/complete -> params: {:?}", params);
    }

    // ---------------------------------------------------------------------
    // Notification handlers
    // ---------------------------------------------------------------------

    fn handle_cancelled_notification(
        &self,
        params: <mcp_types::CancelledNotification as mcp_types::ModelContextProtocolNotification>::Params,
    ) {
        tracing::info!("notifications/cancelled -> params: {:?}", params);
    }

    fn handle_progress_notification(
        &self,
        params: <mcp_types::ProgressNotification as mcp_types::ModelContextProtocolNotification>::Params,
    ) {
        tracing::info!("notifications/progress -> params: {:?}", params);
    }

    fn handle_resource_list_changed(
        &self,
        params: <mcp_types::ResourceListChangedNotification as mcp_types::ModelContextProtocolNotification>::Params,
    ) {
        tracing::info!(
            "notifications/resources/list_changed -> params: {:?}",
            params
        );
    }

    fn handle_resource_updated(
        &self,
        params: <mcp_types::ResourceUpdatedNotification as mcp_types::ModelContextProtocolNotification>::Params,
    ) {
        tracing::info!("notifications/resources/updated -> params: {:?}", params);
    }

    fn handle_prompt_list_changed(
        &self,
        params: <mcp_types::PromptListChangedNotification as mcp_types::ModelContextProtocolNotification>::Params,
    ) {
        tracing::info!("notifications/prompts/list_changed -> params: {:?}", params);
    }

    fn handle_tool_list_changed(
        &self,
        params: <mcp_types::ToolListChangedNotification as mcp_types::ModelContextProtocolNotification>::Params,
    ) {
        tracing::info!("notifications/tools/list_changed -> params: {:?}", params);
    }

    fn handle_logging_message(
        &self,
        params: <mcp_types::LoggingMessageNotification as mcp_types::ModelContextProtocolNotification>::Params,
    ) {
        tracing::info!("notifications/message -> params: {:?}", params);
    }
}