use crate::{Session, ToolCallResponse, ToolDefinition, GenerationOptions, McpClient, ToolCall};
use serde_json::{Value, json};
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use colored::*;

pub type ToolFn = Box<dyn Fn(Value) -> Pin<Box<dyn Future<Output = Result<Value, String>> + Send>> + Send + Sync>;

#[derive(Debug)]
pub enum AgentEvent {
    ToolCallRequested(ToolCall),
    FinalAnswer(String),
}

/// An autonomous agent that can execute tools and maintain a stateful conversation.
/// Now supports MCP (Model Context Protocol) servers and Human-in-the-loop.
pub struct Agent {
    pub session: Session,
    pub tools: Vec<ToolDefinition>,
    tool_impls: HashMap<String, ToolFn>,
    mcp_clients: Vec<McpClient>,
    max_iterations: usize,
    pub hitl_enabled: bool,
}

impl Agent {
    pub fn new(id: &str, instructions: &str) -> Self {
        Self {
            session: Session::new(id, instructions),
            tools: Vec::new(),
            tool_impls: HashMap::new(),
            mcp_clients: Vec::new(),
            max_iterations: 10,
            hitl_enabled: false,
        }
    }

    /// Connects an MCP server to the agent inside a macOS sandbox.
    pub async fn add_mcp_server(&mut self, command: &str, args: &[&str], allowed_paths: Vec<String>, envs: Vec<(String, String)>) -> Result<Vec<String>, String> {
        let client = McpClient::spawn_sandboxed(command, args, allowed_paths, envs).await?;
        client.initialize().await?;
        
        let server_name = client.name.clone();
        let mut discovered_tools = Vec::new();
        let tools_list = client.list_tools().await?;
        if let Some(tools) = tools_list["tools"].as_array() {
            for tool in tools {
                let name = tool["name"].as_str().unwrap_or("").to_string();
                let desc = tool["description"].as_str().unwrap_or("").to_string();
                let params = tool["inputSchema"].clone();
                
                self.tools.push(ToolDefinition {
                    name: name.clone(),
                    description: desc,
                    parameters: params,
                    tool_source: server_name.clone(),
                });
                discovered_tools.push(name);
            }
        }
        
        self.mcp_clients.push(client);
        Ok(discovered_tools)
    }

    /// Disconnects an MCP server and removes its tools.
    pub fn remove_mcp_server(&mut self, name: &str) {
        self.mcp_clients.retain(|c| c.name != name);
        self.tools.retain(|t| t.tool_source != name);
    }

    /// Returns a list of names of currently connected MCP servers.
    pub fn list_mcp_servers(&self) -> Vec<String> {
        self.mcp_clients.iter().map(|c| c.name.clone()).collect()
    }

    pub fn register_tool<F, Fut>(&mut self, mut definition: ToolDefinition, handler: F)
    where
        F: Fn(Value) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<Value, String>> + Send + 'static,
    {
        definition.tool_source = "local".to_string();
        self.tool_impls.insert(
            definition.name.clone(),
            Box::new(move |args| Box::pin(handler(args))),
        );
        self.tools.push(definition);
    }

    /// Executes a single step of reasoning. Returns either a tool call request or a final answer.
    pub async fn step(&self, prompt: &str) -> Result<AgentEvent, String> {
        let tools_json = serde_json::to_string(&self.tools).map_err(|e| e.to_string())?;
        let input = format!(
            "Available Tools: {}\n\nContext: {}\n\nInstructions:\n1. Use tool_call JSON if a tool is needed.\n2. Use text JSON for direct answers.\n3. Do not repeat tool calls unnecessarily.",
            tools_json, prompt
        );

        let mut options = GenerationOptions::default();
        options.temperature = 0.0;

        let response: ToolCallResponse = self.session.respond_json_with_options(&input, options).await?;
        
        match response {
            ToolCallResponse::Plain(text) => Ok(AgentEvent::FinalAnswer(text)),
            ToolCallResponse::Call(call) => Ok(AgentEvent::ToolCallRequested(call)),
        }
    }

    pub async fn execute_tool(&self, call: &ToolCall) -> Result<String, String> {
        if let Some(handler) = self.tool_impls.get(&call.tool_name) {
            match handler(call.arguments.clone()).await {
                Ok(val) => Ok(val.to_string()),
                Err(e) => Ok(format!("Error: {}", e)),
            }
        } else {
            let mut mcp_res = None;
            for client in &self.mcp_clients {
                match client.call("tools/call", json!({
                    "name": call.tool_name,
                    "arguments": call.arguments
                })).await {
                    Ok(res) => {
                        mcp_res = Some(res.to_string());
                        break;
                    }
                    Err(_) => continue,
                }
            }
            Ok(mcp_res.unwrap_or_else(|| format!("Error: Tool '{}' not found.", call.tool_name)))
        }
    }

    pub async fn chat(&self, prompt: &str) -> Result<String, String> {
        let mut current_prompt = prompt.to_string();
        let mut iterations = 0;

        loop {
            if iterations >= self.max_iterations {
                return Err("Agent exceeded maximum tool call iterations.".to_string());
            }

            match self.step(&current_prompt).await? {
                AgentEvent::FinalAnswer(text) => return Ok(text),
                AgentEvent::ToolCallRequested(call) => {
                    println!("{} {} {}", "│".dimmed(), "⚙️ Calling tool:".bold().yellow(), call.tool_name.cyan());
                    let result = self.execute_tool(&call).await?;
                    println!("{} {} {}", "│".dimmed(), "📥 Tool result:".bold().blue(), "success".green());
                    current_prompt = format!("Tool '{}' result: {}. Proceed to next step or provide final answer.", call.tool_name, result);
                    iterations += 1;
                }
            }
        }
    }

    pub fn clear_history(&self) {
        self.session.clear();
    }
}
