use crate::commands::queries::execute_query_core;
use crate::commands::schema::{get_table_columns_core, list_schemas_core, list_tables_core};
use crate::models::connection::ConnectionConfig;
use crate::state::AppState;
use rmcp::handler::server::tool::ToolRouter;
use rmcp::handler::server::wrapper::Parameters;
use rmcp::model::*;
use rmcp::transport::streamable_http_server::session::local::LocalSessionManager;
use rmcp::transport::streamable_http_server::StreamableHttpServerConfig;
use rmcp::transport::streamable_http_server::StreamableHttpService;
use rmcp::{tool, tool_handler, tool_router, ErrorData as McpError, ServerHandler};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::watch;

// --- Tool parameter types ---

#[derive(Debug, Deserialize, JsonSchema)]
struct ConnectionIdParam {
    /// ID of the connection to use (from list_connections)
    connection_id: String,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct ExecuteQueryParam {
    /// ID of the connection to use
    connection_id: String,
    /// SQL query to execute
    sql: String,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct ListTablesParam {
    /// ID of the connection to use
    connection_id: String,
    /// Schema name (e.g. "public")
    schema: String,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct DescribeTableParam {
    /// ID of the connection to use
    connection_id: String,
    /// Schema name (e.g. "public")
    schema: String,
    /// Table name
    table: String,
}

// --- MCP Server ---

#[derive(Clone)]
pub struct TuskMcpServer {
    state: Arc<AppState>,
    connections_path: PathBuf,
    tool_router: ToolRouter<Self>,
}

#[derive(Serialize)]
struct ConnectionStatus {
    id: String,
    name: String,
    host: String,
    port: u16,
    database: String,
    active: bool,
    read_only: bool,
}

#[tool_router]
impl TuskMcpServer {
    fn new(state: Arc<AppState>, connections_path: PathBuf) -> Self {
        Self {
            state,
            connections_path,
            tool_router: Self::tool_router(),
        }
    }

    #[tool(
        description = "List all configured database connections with their active/read-only status"
    )]
    async fn list_connections(&self) -> Result<CallToolResult, McpError> {
        let configs: Vec<ConnectionConfig> = if self.connections_path.exists() {
            let data = std::fs::read_to_string(&self.connections_path).map_err(|e| {
                McpError::internal_error(format!("Failed to read connections file: {}", e), None)
            })?;
            serde_json::from_str(&data).unwrap_or_default()
        } else {
            vec![]
        };

        let pools = self.state.pools.read().await;
        let read_only_map = self.state.read_only.read().await;

        let statuses: Vec<ConnectionStatus> = configs
            .into_iter()
            .map(|c| {
                let active = pools.contains_key(&c.id);
                let read_only = read_only_map.get(&c.id).copied().unwrap_or(true);
                ConnectionStatus {
                    id: c.id,
                    name: c.name,
                    host: c.host,
                    port: c.port,
                    database: c.database,
                    active,
                    read_only,
                }
            })
            .collect();

        let json = serde_json::to_string_pretty(&statuses)
            .map_err(|e| McpError::internal_error(format!("Serialization error: {}", e), None))?;

        Ok(CallToolResult::success(vec![Content::text(json)]))
    }

    #[tool(description = "Execute a SQL query on a database connection. Respects read-only mode.")]
    async fn execute_query(
        &self,
        Parameters(params): Parameters<ExecuteQueryParam>,
    ) -> Result<CallToolResult, McpError> {
        match execute_query_core(&self.state, &params.connection_id, &params.sql).await {
            Ok(result) => {
                let json = serde_json::to_string_pretty(&result).map_err(|e| {
                    McpError::internal_error(format!("Serialization error: {}", e), None)
                })?;
                Ok(CallToolResult::success(vec![Content::text(json)]))
            }
            Err(e) => Ok(CallToolResult::error(vec![Content::text(e.to_string())])),
        }
    }

    #[tool(description = "List all schemas in the current database (excludes system schemas)")]
    async fn list_schemas(
        &self,
        Parameters(params): Parameters<ConnectionIdParam>,
    ) -> Result<CallToolResult, McpError> {
        match list_schemas_core(&self.state, &params.connection_id).await {
            Ok(schemas) => {
                let json = serde_json::to_string_pretty(&schemas).map_err(|e| {
                    McpError::internal_error(format!("Serialization error: {}", e), None)
                })?;
                Ok(CallToolResult::success(vec![Content::text(json)]))
            }
            Err(e) => Ok(CallToolResult::error(vec![Content::text(e.to_string())])),
        }
    }

    #[tool(description = "List all tables in a schema with row count and size information")]
    async fn list_tables(
        &self,
        Parameters(params): Parameters<ListTablesParam>,
    ) -> Result<CallToolResult, McpError> {
        match list_tables_core(&self.state, &params.connection_id, &params.schema).await {
            Ok(tables) => {
                let json = serde_json::to_string_pretty(&tables).map_err(|e| {
                    McpError::internal_error(format!("Serialization error: {}", e), None)
                })?;
                Ok(CallToolResult::success(vec![Content::text(json)]))
            }
            Err(e) => Ok(CallToolResult::error(vec![Content::text(e.to_string())])),
        }
    }

    #[tool(
        description = "Describe table columns: name, type, nullable, primary key, default value"
    )]
    async fn describe_table(
        &self,
        Parameters(params): Parameters<DescribeTableParam>,
    ) -> Result<CallToolResult, McpError> {
        match get_table_columns_core(
            &self.state,
            &params.connection_id,
            &params.schema,
            &params.table,
        )
        .await
        {
            Ok(columns) => {
                let json = serde_json::to_string_pretty(&columns).map_err(|e| {
                    McpError::internal_error(format!("Serialization error: {}", e), None)
                })?;
                Ok(CallToolResult::success(vec![Content::text(json)]))
            }
            Err(e) => Ok(CallToolResult::error(vec![Content::text(e.to_string())])),
        }
    }
}

#[tool_handler]
impl ServerHandler for TuskMcpServer {
    fn get_info(&self) -> ServerInfo {
        ServerInfo {
            protocol_version: ProtocolVersion::V_2025_03_26,
            capabilities: ServerCapabilities::builder().enable_tools().build(),
            server_info: Implementation {
                name: "tusk-mcp".to_string(),
                version: env!("CARGO_PKG_VERSION").to_string(),
                title: Some("Tusk PostgreSQL MCP Server".to_string()),
                description: None,
                website_url: None,
                icons: None,
            },
            instructions: Some(
                "Tusk MCP server provides access to PostgreSQL databases through the Tusk GUI app. \
                 Use list_connections to see available connections (must be connected in the Tusk GUI first). \
                 Use the connection_id from list_connections in other tools. \
                 Read-only mode is respected — write queries will fail if a connection is in read-only mode."
                    .to_string(),
            ),
        }
    }
}

pub async fn start_mcp_server(
    state: Arc<AppState>,
    connections_path: PathBuf,
    port: u16,
    mut shutdown_rx: watch::Receiver<bool>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    let service = StreamableHttpService::new(
        move || Ok(TuskMcpServer::new(state.clone(), connections_path.clone())),
        Arc::new(LocalSessionManager::default()),
        StreamableHttpServerConfig::default(),
    );

    let router = axum::Router::new().nest_service("/mcp", service);
    let addr = format!("127.0.0.1:{}", port);
    let listener = tokio::net::TcpListener::bind(&addr).await?;

    log::info!("MCP server listening on http://{}/mcp", addr);

    tokio::select! {
        res = axum::serve(listener, router) => {
            res?;
        }
        _ = shutdown_rx.changed() => {
            log::info!("MCP server stopped by shutdown signal");
        }
    }

    Ok(())
}
