//! MCP server implementation

use std::io::{self, BufRead, BufReader, Write};

use anyhow::Result;
use serde_json::{json, Value};
use tracing::{debug, error, info, warn};

use crate::{
    auth::{extract_token_from_params, AuthConfig},
    prompts, resources, tools,
    types::{
        JsonRpcRequest, JsonRpcResponse, PromptsCapability, ResourcesCapability,
        ServerCapabilities, ToolsCapability,
    },
};

/// MCP Server with authentication support
pub struct McpServer {
    name: String,
    version: String,
    auth: AuthConfig,
}

impl McpServer {
    /// Create a new MCP server
    pub fn new(name: String, version: String) -> Self {
        let auth = AuthConfig::from_env();
        if auth.is_required() {
            info!("MCP server authentication enabled via KOGRAL_MCP_TOKEN");
        } else {
            info!("MCP server running without authentication");
        }
        Self {
            name,
            version,
            auth,
        }
    }

    /// Run the MCP server with stdio transport
    pub async fn run_stdio(&self) -> Result<()> {
        info!("Starting MCP server: {} v{}", self.name, self.version);
        info!("Transport: stdio");

        let stdin = io::stdin();
        let mut reader = BufReader::new(stdin.lock());
        let mut stdout = io::stdout();

        loop {
            let mut line = String::new();
            match reader.read_line(&mut line) {
                Ok(0) => {
                    // EOF
                    info!("Received EOF, shutting down");
                    break;
                }
                Ok(_) => {
                    let line = line.trim();
                    if line.is_empty() {
                        continue;
                    }

                    debug!("Received request: {}", line);

                    match self.handle_request(line).await {
                        Ok(response) => {
                            let response_json = serde_json::to_string(&response)?;
                            debug!("Sending response: {}", response_json);
                            writeln!(stdout, "{}", response_json)?;
                            stdout.flush()?;
                        }
                        Err(e) => {
                            error!("Error handling request: {}", e);
                            let error_response = JsonRpcResponse::error(
                                None,
                                -32603,
                                format!("Internal error: {}", e),
                            );
                            let error_json = serde_json::to_string(&error_response)?;
                            writeln!(stdout, "{}", error_json)?;
                            stdout.flush()?;
                        }
                    }
                }
                Err(e) => {
                    error!("Error reading from stdin: {}", e);
                    break;
                }
            }
        }

        Ok(())
    }

    /// Handle a JSON-RPC request
    async fn handle_request(&self, request_json: &str) -> Result<JsonRpcResponse> {
        let request: JsonRpcRequest = serde_json::from_str(request_json)?;

        debug!("Handling method: {}", request.method);

        // Verify authentication for protected methods
        // Initialize is allowed without auth for initial setup, ping can be used for
        // health checks
        if !matches!(request.method.as_str(), "initialize" | "ping") {
            if let Err(e) = self.authenticate(&request.params) {
                debug!("Authentication failed: {}", e);
                return Ok(JsonRpcResponse::error(
                    request.id,
                    -32600,
                    format!("Unauthorized: {}", e),
                ));
            }
        }

        let result = match request.method.as_str() {
            "initialize" => self.handle_initialize(request.params).await,
            "tools/list" => self.handle_tools_list().await,
            "tools/call" => self.handle_tools_call(request.params).await,
            "resources/list" => self.handle_resources_list().await,
            "resources/read" => self.handle_resources_read(request.params).await,
            "prompts/list" => self.handle_prompts_list().await,
            "prompts/get" => self.handle_prompts_get(request.params).await,
            "ping" => Ok(json!({"status": "ok"})),
            _ => {
                warn!("Unknown method: {}", request.method);
                return Ok(JsonRpcResponse::error(
                    request.id,
                    -32601,
                    format!("Method not found: {}", request.method),
                ));
            }
        };

        match result {
            Ok(value) => Ok(JsonRpcResponse::success(request.id, value)),
            Err(e) => {
                error!("Error executing method {}: {}", request.method, e);
                Ok(JsonRpcResponse::error(
                    request.id,
                    -32603,
                    format!("Internal error: {}", e),
                ))
            }
        }
    }

    /// Authenticate a request by verifying its token
    fn authenticate(&self, params: &Value) -> Result<()> {
        let token = extract_token_from_params(params);
        self.auth.verify(token)?;
        Ok(())
    }

    /// Handle initialize request
    async fn handle_initialize(&self, _params: Value) -> Result<Value> {
        info!("Initializing MCP server");

        Ok(json!({
            "protocolVersion": "2024-11-05",
            "serverInfo": {
                "name": self.name,
                "version": self.version
            },
            "capabilities": ServerCapabilities {
                tools: Some(ToolsCapability {
                    list_changed: Some(false),
                }),
                resources: Some(ResourcesCapability {
                    subscribe: Some(false),
                    list_changed: Some(false),
                }),
                prompts: Some(PromptsCapability {
                    list_changed: Some(false),
                }),
            }
        }))
    }

    /// Handle tools/list request
    async fn handle_tools_list(&self) -> Result<Value> {
        let tools = tools::list_tools();
        Ok(json!({ "tools": tools }))
    }

    /// Handle tools/call request
    async fn handle_tools_call(&self, params: Value) -> Result<Value> {
        let name = params["name"]
            .as_str()
            .ok_or_else(|| anyhow::anyhow!("Missing tool name"))?;
        let arguments = params["arguments"].clone();

        let result = tools::execute_tool(name, arguments).await?;
        Ok(serde_json::to_value(result)?)
    }

    /// Handle resources/list request
    async fn handle_resources_list(&self) -> Result<Value> {
        let resources = resources::list_resources().await?;
        Ok(json!({ "resources": resources }))
    }

    /// Handle resources/read request
    async fn handle_resources_read(&self, params: Value) -> Result<Value> {
        let uri = params["uri"]
            .as_str()
            .ok_or_else(|| anyhow::anyhow!("Missing resource URI"))?;

        let contents = resources::read_resource(uri).await?;
        Ok(json!({ "contents": [contents] }))
    }

    /// Handle prompts/list request
    async fn handle_prompts_list(&self) -> Result<Value> {
        let prompts = prompts::list_prompts();
        Ok(json!({ "prompts": prompts }))
    }

    /// Handle prompts/get request
    async fn handle_prompts_get(&self, params: Value) -> Result<Value> {
        let name = params["name"]
            .as_str()
            .ok_or_else(|| anyhow::anyhow!("Missing prompt name"))?;
        let arguments = params["arguments"].clone();

        let messages = prompts::execute_prompt(name, arguments).await?;
        Ok(json!({ "messages": messages }))
    }
}
