use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::process::Stdio;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, Command};
use std::sync::Arc;
use tokio::sync::Mutex;

#[derive(Debug, Serialize, Deserialize)]
pub struct JsonRpcRequest {
    pub jsonrpc: String,
    pub method: String,
    pub params: Value,
    pub id: u64,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct JsonRpcResponse {
    pub jsonrpc: String,
    pub result: Option<Value>,
    pub error: Option<Value>,
    pub id: u64,
}

/// A client for the Model Context Protocol (MCP) using stdio transport.
pub struct McpClient {
    pub name: String,
    child: Arc<Mutex<Child>>,
    request_id: Arc<Mutex<u64>>,
}

impl McpClient {
    /// Launches an MCP server and connects via stdio.
    pub async fn spawn(command: &str, args: &[&str]) -> Result<Self, String> {
        let child = Command::new(command)
            .args(args)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::inherit())
            .spawn()
            .map_err(|e| format!("Failed to spawn MCP server: {}", e))?;

        Ok(Self {
            name: command.to_string(),
            child: Arc::new(Mutex::new(child)),
            request_id: Arc::new(Mutex::new(0)),
        })
    }

    /// Launches an MCP server inside a macOS sandbox (sandbox-exec).
    /// Restricts the server to the provided allowed_paths and injects environment variables.
    pub async fn spawn_sandboxed(command: &str, args: &[&str], allowed_paths: Vec<String>, envs: Vec<(String, String)>) -> Result<Self, String> {
        // Construct basic macOS sandbox profile
        let mut profile = String::from("(version 1)\n(deny default)\n(allow process-fork)\n(allow signal)\n(allow sysctl-read)\n");
        
        // Allow reading/writing to the provided paths
        for path in allowed_paths {
            let absolute_path = std::fs::canonicalize(&path).unwrap_or_else(|_| std::path::PathBuf::from(&path));
            let path_str = absolute_path.to_string_lossy();
            profile.push_str(&format!("(allow file-read* file-write* (subpath \"{}\"))\n", path_str));
        }

        // Essential system paths for Node/npx/Standard dynamic linkers
        profile.push_str("(allow file-read* (subpath \"/usr/lib\"))\n");
        profile.push_str("(allow file-read* (subpath \"/usr/share\"))\n");
        profile.push_str("(allow file-read* (subpath \"/private/var/folders\"))\n");
        profile.push_str("(allow file-read* (subpath \"/dev/urandom\"))\n");
        profile.push_str("(allow network-outbound)\n"); // Allow network for npm/API calls
        profile.push_str("(allow process-exec)\n"); // Allow npx to exec node

        let mut sandbox_args = vec!["-p", &profile, command];
        sandbox_args.extend_from_slice(args);

        let mut cmd = Command::new("sandbox-exec");
        cmd.args(&sandbox_args)
            .envs(envs)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::inherit());
            
        let child = cmd.spawn()
            .map_err(|e| format!("Failed to spawn sandboxed MCP server: {}", e))?;

        Ok(Self {
            name: command.to_string(),
            child: Arc::new(Mutex::new(child)),
            request_id: Arc::new(Mutex::new(0)),
        })
    }

    /// Sends an initialization request to the MCP server.
    pub async fn initialize(&self) -> Result<(), String> {
        let params = serde_json::json!({
            "protocolVersion": "2024-11-05",
            "capabilities": {},
            "clientInfo": { "name": "sutra_afm_lib", "version": "0.8.1" }
        });
        self.call("initialize", params).await?;
        self.notification("notifications/initialized", serde_json::json!({})).await?;
        Ok(())
    }

    /// Calls a method on the MCP server and waits for a response.
    pub async fn call(&self, method: &str, params: Value) -> Result<Value, String> {
        let mut id_guard = self.request_id.lock().await;
        *id_guard += 1;
        let id = *id_guard;

        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            method: method.to_string(),
            params,
            id,
        };

        let mut child = self.child.lock().await;
        
        let req_json = serde_json::to_string(&request).map_err(|e| e.to_string())? + "\n";
        
        if let Some(stdin) = child.stdin.as_mut() {
            stdin.write_all(req_json.as_bytes()).await.map_err(|e| e.to_string())?;
            stdin.flush().await.map_err(|e| e.to_string())?;
        } else {
            return Err("Failed to open stdin".to_string());
        }

        let mut line = String::new();
        if let Some(stdout) = child.stdout.as_mut() {
            let mut reader = BufReader::new(stdout);
            reader.read_line(&mut line).await.map_err(|e| e.to_string())?;
        } else {
            return Err("Failed to open stdout".to_string());
        }
        
        if line.is_empty() {
            return Err("MCP server closed connection unexpectedly".to_string());
        }

        let response: JsonRpcResponse = serde_json::from_str(&line).map_err(|e| format!("JSON Parse error: {}. Line: {}", e, line))?;
        if let Some(err) = response.error {
            Err(format!("MCP Error: {}", err))
        } else {
            Ok(response.result.unwrap_or(Value::Null))
        }
    }

    /// Sends a notification (fire-and-forget) to the MCP server.
    pub async fn notification(&self, method: &str, params: Value) -> Result<(), String> {
        let request = serde_json::json!({
            "jsonrpc": "2.0",
            "method": method,
            "params": params
        });

        let mut child = self.child.lock().await;
        let stdin = child.stdin.as_mut().ok_or("Failed to open stdin")?;
        let req_json = serde_json::to_string(&request).map_err(|e| e.to_string())? + "\n";
        stdin.write_all(req_json.as_bytes()).await.map_err(|e| e.to_string())?;
        stdin.flush().await.map_err(|e| e.to_string())?;
        Ok(())
    }

    /// Lists tools available on this MCP server.
    pub async fn list_tools(&self) -> Result<Value, String> {
        self.call("tools/list", serde_json::json!({})).await
    }
}
