//! Unified chat runner interface for agents
//!
//! Provides a common trait for both single agents and multi-agent quorums.

use crate::api::{Agent, AgentInfra};
use crate::config::{Config, ConfigSource, load_config};
use crate::events::EventEnvelope;
#[cfg(feature = "api")]
use crate::server::AgentServer;
use anyhow::Result;
use async_trait::async_trait;
use serde_json::Value;
use tokio::sync::broadcast;

/// Unified interface for chat operations
///
/// Implemented by `Agent` to provide a consistent API.
#[async_trait]
pub trait ChatRunner: Send + Sync {
    /// Send a chat message and get a response
    async fn chat(&self, prompt: &str) -> Result<String>;

    /// Create a new chat session for maintaining conversation history
    async fn chat_session(&self) -> Result<Box<dyn ChatSession>>;

    /// Subscribe to agent events (tool calls, messages, etc.)
    fn subscribe(&self) -> broadcast::Receiver<EventEnvelope>;

    /// Register a callback for tool call events (boxed version)
    fn on_tool_call_boxed(&self, callback: Box<dyn Fn(String, Value) + Send + Sync>);

    /// Register a callback for tool completion events (boxed version)
    fn on_tool_complete_boxed(&self, callback: Box<dyn Fn(String, String) + Send + Sync>);

    /// Register a callback for message events (boxed version)
    fn on_message_boxed(&self, callback: Box<dyn Fn(String, String) + Send + Sync>);

    /// Register a callback for delegation events (boxed version)
    fn on_delegation_boxed(&self, callback: Box<dyn Fn(String, String) + Send + Sync>);

    /// Register a callback for error events (boxed version)
    fn on_error_boxed(&self, callback: Box<dyn Fn(String) + Send + Sync>);

    /// Get the UI server
    #[cfg(feature = "api")]
    fn server(&self) -> AgentServer;
}

/// Extension trait for convenient callback registration  
pub trait ChatRunnerExt: ChatRunner {
    fn on_tool_call(&self, callback: impl Fn(String, Value) + Send + Sync + 'static) {
        self.on_tool_call_boxed(Box::new(callback));
    }

    fn on_tool_complete(&self, callback: impl Fn(String, String) + Send + Sync + 'static) {
        self.on_tool_complete_boxed(Box::new(callback));
    }

    fn on_message(&self, callback: impl Fn(String, String) + Send + Sync + 'static) {
        self.on_message_boxed(Box::new(callback));
    }

    fn on_delegation(&self, callback: impl Fn(String, String) + Send + Sync + 'static) {
        self.on_delegation_boxed(Box::new(callback));
    }

    fn on_error(&self, callback: impl Fn(String) + Send + Sync + 'static) {
        self.on_error_boxed(Box::new(callback));
    }
}

// Blanket implementation
impl<T: ChatRunner + ?Sized> ChatRunnerExt for T {}

/// A chat session for maintaining conversation context
#[async_trait]
pub trait ChatSession: Send + Sync {
    /// Get the session ID
    fn id(&self) -> &str;

    /// Send a chat message within this session
    async fn chat(&self, prompt: &str) -> Result<String>;

    /// Register a callback for tool call events in this session
    fn on_tool_call_boxed(&self, callback: Box<dyn Fn(String, Value) + Send + Sync>);

    /// Register a callback for tool completion events in this session
    fn on_tool_complete_boxed(&self, callback: Box<dyn Fn(String, String) + Send + Sync>);

    /// Register a callback for message events in this session
    fn on_message_boxed(&self, callback: Box<dyn Fn(String, String) + Send + Sync>);

    /// Register a callback for error events in this session
    fn on_error_boxed(&self, callback: Box<dyn Fn(String) + Send + Sync>);
}

/// Extension trait for convenient callback registration
pub trait ChatSessionExt: ChatSession {
    fn on_tool_call(&self, callback: impl Fn(String, Value) + Send + Sync + 'static) {
        self.on_tool_call_boxed(Box::new(callback));
    }

    fn on_tool_complete(&self, callback: impl Fn(String, String) + Send + Sync + 'static) {
        self.on_tool_complete_boxed(Box::new(callback));
    }

    fn on_message(&self, callback: impl Fn(String, String) + Send + Sync + 'static) {
        self.on_message_boxed(Box::new(callback));
    }

    fn on_error(&self, callback: impl Fn(String) + Send + Sync + 'static) {
        self.on_error_boxed(Box::new(callback));
    }
}

// Blanket implementation
impl<T: ChatSession + ?Sized> ChatSessionExt for T {}

/// Unified runner for both single agents and multi-agent quorums.
///
/// Since `Agent` now handles both single and multi-agent configurations,
/// `AgentRunner` is a thin wrapper around `Agent` for backward compatibility.
///
/// # Examples
///
/// ```no_run
/// use querymt_agent::prelude::*;
///
/// # #[tokio::main]
/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
/// // Load from config
/// let runner = from_config("agent.toml").await?;
///
/// // Use chat functionality
/// let response = runner.chat("Hello!").await?;
///
/// // Start ACP stdio server
/// runner.acp("stdio").await?;
/// # Ok(())
/// # }
/// ```
pub struct AgentRunner(Agent);

impl AgentRunner {
    /// Create a new runner from an `Agent`.
    pub fn new(agent: Agent) -> Self {
        Self(agent)
    }

    /// Send a chat message and get a response
    pub async fn chat(&self, prompt: &str) -> Result<String> {
        self.0.chat(prompt).await
    }

    /// Create a new chat session for maintaining conversation history
    pub async fn chat_session(&self) -> Result<Box<dyn ChatSession>> {
        ChatRunner::chat_session(&self.0).await
    }

    /// Start an ACP server with the specified transport
    pub async fn acp(&self, transport: &str) -> Result<()> {
        self.0.acp(transport).await
    }

    /// Get the UI server
    #[cfg(feature = "api")]
    pub fn server(&self) -> AgentServer {
        self.0.server()
    }

    /// Subscribe to agent events (tool calls, messages, etc.)
    pub fn subscribe(&self) -> broadcast::Receiver<EventEnvelope> {
        self.0.subscribe()
    }

    /// Register a callback for tool call events
    pub fn on_tool_call(&self, callback: impl Fn(String, Value) + Send + Sync + 'static) -> &Self {
        self.0.on_tool_call(callback);
        self
    }

    /// Register a callback for tool completion events
    pub fn on_tool_complete(
        &self,
        callback: impl Fn(String, String) + Send + Sync + 'static,
    ) -> &Self {
        self.0.on_tool_complete(callback);
        self
    }

    /// Register a callback for message events
    pub fn on_message(&self, callback: impl Fn(String, String) + Send + Sync + 'static) -> &Self {
        self.0.on_message(callback);
        self
    }

    /// Register a callback for delegation events
    pub fn on_delegation(
        &self,
        callback: impl Fn(String, String) + Send + Sync + 'static,
    ) -> &Self {
        self.0.on_delegation(callback);
        self
    }

    /// Register a callback for error events
    pub fn on_error(&self, callback: impl Fn(String) + Send + Sync + 'static) -> &Self {
        self.0.on_error(callback);
        self
    }

    /// Access the underlying `AgentHandle` for advanced use cases.
    pub fn handle(&self) -> std::sync::Arc<crate::agent::LocalAgentHandle> {
        self.0.handle()
    }

    /// Get a reference to the inner Agent.
    pub fn as_agent(&self) -> &Agent {
        &self.0
    }

    /// Convert into the inner Agent.
    pub fn into_agent(self) -> Agent {
        self.0
    }

    /// Returns true if this is a multi-agent (quorum) runner.
    pub fn is_multi(&self) -> bool {
        self.0.is_multi()
    }
}

/// Convert AgentRunner to a trait object for backwards compatibility.
impl From<AgentRunner> for Box<dyn ChatRunner> {
    fn from(runner: AgentRunner) -> Self {
        Box::new(runner.0)
    }
}

/// Load an agent or quorum from configuration source.
///
/// Automatically detects whether the config is for a single agent or multi-agent quorum.
/// Returns an `AgentRunner` which provides access to all agent functionality including
/// the `.acp()` method for starting ACP servers.
///
/// `source` accepts either a config file path or inline TOML (`ConfigSource::Toml`).
/// Inline TOML must contain fully inlined system prompts (no `{ file = ... }` entries).
///
/// ## Config-Driven Mesh Bootstrap
///
/// When the `remote` feature is enabled and the TOML config contains:
///
/// ```toml
/// [mesh]
/// enabled = true
/// listen = "/ip4/0.0.0.0/tcp/0"
/// discovery = "mdns"
///
/// [[remote_agents]]
/// id = "gpu-coder"
/// name = "GPU Coder"
/// peer = "dev-gpu"
/// capabilities = ["shell", "gpu"]
/// ```
///
/// `from_config` will automatically:
/// 1. Bootstrap the kameo libp2p swarm.
/// 2. Register this node as a `RemoteNodeManager` in the DHT.
/// 3. Pre-populate the agent registry with `AgentInfo` for each `[[remote_agents]]` entry,
///    backed by a `RemoteAgentHandle` that routes delegation calls via the mesh.
///
/// # Example
///
/// ```no_run
/// use querymt_agent::config::ConfigSource;
/// use querymt_agent::prelude::*;
///
/// # #[tokio::main]
/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
/// let runner = from_config("agent.toml").await?;
///
/// // Or from inline TOML
/// let inline = r#"[agent]
/// provider = \"openai\"
/// model = \"gpt-4.1-mini\"
/// tools = []
/// "#;
/// let _runner = from_config(ConfigSource::Toml(inline.to_string())).await?;
///
/// // Chat functionality
/// let response = runner.chat("Hello!").await?;
/// println!("{}", response);
///
/// // Start ACP stdio server
/// runner.acp("stdio").await?;
/// # Ok(())
/// # }
/// ```
pub async fn from_config(source: impl Into<ConfigSource>) -> Result<AgentRunner> {
    let config = load_config(source).await?;
    build_runner_from_config(config, None).await
}

/// Build a runner from config while reusing caller-provided infrastructure.
pub async fn from_config_with_infra(
    source: impl Into<ConfigSource>,
    infra: AgentInfra,
) -> Result<AgentRunner> {
    let config = load_config(source).await?;
    build_runner_from_config(config, Some(infra)).await
}

/// Build a runner from an already-loaded config while reusing caller-provided infrastructure.
pub async fn from_config_value_with_infra(
    config: Config,
    infra: AgentInfra,
) -> Result<AgentRunner> {
    build_runner_from_config(config, Some(infra)).await
}

async fn build_runner_from_config(
    config: Config,
    infra: Option<AgentInfra>,
) -> Result<AgentRunner> {
    match config {
        Config::Single(single_config) => {
            let single_config = *single_config;
            let agent = match infra {
                Some(infra) => Agent::from_single_config_with_infra(single_config, infra).await?,
                None => Agent::from_single_config(single_config).await?,
            };
            Ok(AgentRunner::new(agent))
        }
        Config::Multi(quorum_config) => {
            let quorum_config = *quorum_config;
            let agent = match infra {
                Some(infra) => Agent::from_quorum_config_with_infra(quorum_config, infra).await?,
                None => Agent::from_quorum_config(quorum_config).await?,
            };
            Ok(AgentRunner::new(agent))
        }
    }
}
