use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use rmcp::{
    ErrorData as McpError, ServerHandler, ServiceExt,
    handler::server::{tool::ToolRouter, wrapper::Parameters},
    model::{
        CallToolRequestParams, CallToolResult, Content, Implementation, ListToolsResult,
        PaginatedRequestParams, ProtocolVersion, ServerCapabilities, ServerInfo,
    },
    service::{RequestContext, RoleServer},
    tool, tool_router,
    transport::stdio,
};
use schemars::JsonSchema;
use serde::Deserialize;
use tokio::sync::RwLock;

use crate::executor::{self, ExecConfig};

// =============================================================================
// Public entry point
// =============================================================================

pub async fn run() -> anyhow::Result<()> {
    let server = PythonExecMcp::new();
    let service = server.serve(stdio()).await?;
    service.waiting().await?;
    Ok(())
}

// =============================================================================
// MCP Server
// =============================================================================

#[derive(Clone)]
pub struct PythonExecMcp {
    tool_router: ToolRouter<Self>,
    exec_config: Arc<RwLock<ExecConfig>>,
}

impl PythonExecMcp {
    pub fn new() -> Self {
        Self {
            tool_router: Self::tool_router(),
            exec_config: Arc::new(RwLock::new(ExecConfig::from_env())),
        }
    }
}

// =============================================================================
// ServerHandler
// =============================================================================

impl ServerHandler for PythonExecMcp {
    fn get_info(&self) -> ServerInfo {
        ServerInfo {
            protocol_version: ProtocolVersion::V_2025_03_26,
            capabilities: ServerCapabilities::builder().enable_tools().build(),
            server_info: Implementation::from_build_env(),
            instructions: Some(
                "Python script execution server with safe defaults (alternative to shell tools). \
                 Auto-detects python3 or micropython. \
                 os.system/subprocess/shutil are blocked by default. \
                 Configure via ENV: PYEXEC_ENGINE, PYEXEC_TIMEOUT, PYEXEC_WORKDIR, PYEXEC_DENY. \
                 Use `execute` to run scripts, `session_start` to override settings at runtime."
                    .to_string(),
            ),
        }
    }

    async fn list_tools(
        &self,
        _request: Option<PaginatedRequestParams>,
        _cx: RequestContext<RoleServer>,
    ) -> Result<ListToolsResult, McpError> {
        Ok(ListToolsResult {
            tools: self.tool_router.list_all(),
            next_cursor: None,
            meta: None,
        })
    }

    async fn call_tool(
        &self,
        request: CallToolRequestParams,
        cx: RequestContext<RoleServer>,
    ) -> Result<CallToolResult, McpError> {
        let tool_ctx = rmcp::handler::server::tool::ToolCallContext::new(self, request, cx);
        self.tool_router.call(tool_ctx).await
    }
}

// =============================================================================
// Tool params
// =============================================================================

#[derive(Debug, Deserialize, JsonSchema)]
struct SessionStartParams {
    /// Working directory for script execution.
    #[serde(default)]
    workdir: Option<String>,

    /// Timeout in seconds (default: 30, max: 300).
    #[serde(default)]
    timeout_secs: Option<u64>,

    /// Python binary override (auto-detected: python3 → micropython).
    #[serde(default)]
    engine: Option<String>,

    /// Additional modules to deny (comma-separated). Added on top of defaults.
    #[serde(default)]
    deny: Option<String>,

    /// Restrict execute_file paths to within workdir (default: false).
    #[serde(default)]
    restrict_path: Option<bool>,

    /// Paths to block from open() (colon-separated, prefix match).
    #[serde(default)]
    deny_paths: Option<String>,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct ExecuteParams {
    /// Python script to execute.
    script: String,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct ExecuteFileParams {
    /// Path to a .py file to execute.
    path: String,
}

// =============================================================================
// Tools
// =============================================================================

#[tool_router]
impl PythonExecMcp {
    /// Configure session: engine, timeout, workdir, deny modules.
    #[tool(name = "session_start")]
    async fn session_start(
        &self,
        Parameters(params): Parameters<SessionStartParams>,
    ) -> Result<CallToolResult, McpError> {
        let mut config = self.exec_config.write().await;

        if let Some(ref wd) = params.workdir {
            let path = PathBuf::from(wd);
            if !path.is_dir() {
                return Err(McpError::invalid_params(
                    format!("workdir does not exist: {wd}"),
                    None,
                ));
            }
            config.workdir = Some(path);
        }

        if let Some(secs) = params.timeout_secs {
            config.timeout = Duration::from_secs(secs.min(300));
        }

        if let Some(ref engine) = params.engine {
            config.binary = engine.clone();
        }

        if let Some(ref deny) = params.deny {
            for m in deny.split(',').map(|m| m.trim().to_string()) {
                if !m.is_empty() && !config.deny_modules.contains(&m) {
                    config.deny_modules.push(m);
                }
            }
        }

        if let Some(restrict) = params.restrict_path {
            config.restrict_path = restrict;
        }

        if let Some(ref deny_paths) = params.deny_paths {
            for p in deny_paths.split(':').map(|p| p.trim()) {
                if !p.is_empty() {
                    let path = PathBuf::from(p);
                    if !config.deny_paths.contains(&path) {
                        config.deny_paths.push(path);
                    }
                }
            }
        }

        let summary = format!(
            "Session configured:\n  engine: {}\n  timeout: {}s\n  workdir: {}\n  restrict_path: {}\n  deny (extra): {:?}\n  deny_paths: {:?}",
            config.binary,
            config.timeout.as_secs(),
            config
                .workdir
                .as_deref()
                .map(|p| p.display().to_string())
                .unwrap_or_else(|| "(cwd)".to_string()),
            config.restrict_path,
            config.deny_modules,
            config.deny_paths,
        );

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

    /// Execute a Python script string. Returns stdout, stderr, and exit code.
    #[tool(name = "execute")]
    async fn execute(
        &self,
        Parameters(params): Parameters<ExecuteParams>,
    ) -> Result<CallToolResult, McpError> {
        let config = self.exec_config.read().await.clone();
        Ok(format_result(
            executor::execute(&params.script, &config).await,
        ))
    }

    /// Execute a Python script file. Returns stdout, stderr, and exit code.
    #[tool(name = "execute_file")]
    async fn execute_file(
        &self,
        Parameters(params): Parameters<ExecuteFileParams>,
    ) -> Result<CallToolResult, McpError> {
        let config = self.exec_config.read().await.clone();
        let path = PathBuf::from(&params.path);
        Ok(format_result(executor::execute_file(&path, &config).await))
    }
}

fn format_result(result: Result<executor::ExecResult, executor::ExecError>) -> CallToolResult {
    match result {
        Ok(r) => {
            let mut output = String::new();
            if !r.stdout.is_empty() {
                output.push_str(&r.stdout);
            }
            if !r.stderr.is_empty() {
                if !output.is_empty() {
                    output.push('\n');
                }
                output.push_str("[stderr]\n");
                output.push_str(&r.stderr);
            }
            if output.is_empty() {
                output.push_str("(no output)");
            }
            if r.exit_code != 0 {
                output.push_str(&format!("\n[exit code: {}]", r.exit_code));
            }
            CallToolResult::success(vec![Content::text(output)])
        }
        Err(e) => CallToolResult::error(vec![Content::text(e.to_string())]),
    }
}
