pub mod card;
pub mod message;
pub mod protocol;

pub use card::{AgentCapabilities, AgentCard, AgentEndpoints};
pub use message::A2aMessageType;
pub use protocol::{A2aMessagePayload, A2aRequest, A2aResponse, A2aTaskStatus};

use anyhow::Result;
use enact_providers::{ModelProvider, OpenAICompatible};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tracing::info;

const DELEGATE_TIMEOUT_SECS: u64 = 120;

/// Configuration for a delegate agent
#[derive(Debug, Clone)]
pub struct DelegateAgentConfig {
    pub provider: String,
    pub model: String,
    pub system_prompt: Option<String>,
    pub api_key: Option<String>,
    pub temperature: Option<f32>,
    pub max_depth: u32,
}

impl Default for DelegateAgentConfig {
    fn default() -> Self {
        Self {
            provider: "openai".to_string(),
            model: "gpt-4o-mini".to_string(),
            system_prompt: None,
            api_key: None,
            temperature: Some(0.7),
            max_depth: 3,
        }
    }
}

/// Multi-Agent system that can delegate tasks to specialized agents
pub struct MultiAgentSystem {
    agents: HashMap<String, DelegateAgentConfig>,
    fallback_credential: Option<String>,
}

impl MultiAgentSystem {
    pub fn new() -> Self {
        Self {
            agents: HashMap::new(),
            fallback_credential: None,
        }
    }

    pub fn with_agents(agents: HashMap<String, DelegateAgentConfig>) -> Self {
        Self {
            agents,
            fallback_credential: None,
        }
    }

    pub fn set_fallback_credential(&mut self, credential: String) {
        self.fallback_credential = Some(credential);
    }

    pub fn register_agent(&mut self, name: impl Into<String>, config: DelegateAgentConfig) {
        self.agents.insert(name.into(), config);
    }

    pub fn get_agent_names(&self) -> Vec<&str> {
        self.agents.keys().map(|s| s.as_str()).collect()
    }

    /// Delegate a task to a specialized agent
    pub async fn delegate(
        &self,
        agent_name: &str,
        prompt: &str,
        context: Option<&str>,
        current_depth: u32,
    ) -> Result<DelegateResult> {
        let agent_name = agent_name.trim();
        if agent_name.is_empty() {
            anyhow::bail!("Agent name cannot be empty");
        }

        let prompt = prompt.trim();
        if prompt.is_empty() {
            anyhow::bail!("Prompt cannot be empty");
        }

        // Look up agent config
        let config = self.agents.get(agent_name).ok_or_else(|| {
            let available: Vec<&str> = self.agents.keys().map(|s| s.as_str()).collect();
            anyhow::anyhow!(
                "Unknown agent '{agent_name}'. Available agents: {}",
                if available.is_empty() {
                    "(none configured)".to_string()
                } else {
                    available.join(", ")
                }
            )
        })?;

        // Check recursion depth
        if current_depth >= config.max_depth {
            return Ok(DelegateResult {
                success: false,
                output: String::new(),
                error: Some(format!(
                    "Delegation depth limit reached ({current_depth}/{}). Cannot delegate further to prevent infinite loops.",
                    config.max_depth
                )),
                agent: agent_name.to_string(),
                provider: config.provider.clone(),
                model: config.model.clone(),
            });
        }

        info!(
            "Delegating to agent '{}' ({}/{})",
            agent_name, current_depth, config.max_depth
        );

        // Create provider for this agent
        let api_key = config
            .api_key
            .clone()
            .or_else(|| self.fallback_credential.clone());

        let provider: Arc<dyn ModelProvider> = match config.provider.as_str() {
            "openai" | "openai-compatible" => {
                let base_url = std::env::var("OPENAI_BASE_URL")
                    .unwrap_or_else(|_| "https://api.openai.com".to_string());
                Arc::new(OpenAICompatible::new(
                    base_url,
                    config.model.clone(),
                    api_key,
                ))
            }
            _ => {
                return Ok(DelegateResult {
                    success: false,
                    output: String::new(),
                    error: Some(format!(
                        "Unsupported provider '{}' for agent '{}'",
                        config.provider, agent_name
                    )),
                    agent: agent_name.to_string(),
                    provider: config.provider.clone(),
                    model: config.model.clone(),
                });
            }
        };

        // Build the full prompt
        let full_prompt = if let Some(ctx) = context {
            if ctx.trim().is_empty() {
                prompt.to_string()
            } else {
                format!("[Context]\n{}\n\n[Task]\n{}", ctx, prompt)
            }
        } else {
            prompt.to_string()
        };

        // Execute with timeout
        let temperature = config.temperature.unwrap_or(0.7);
        let result = tokio::time::timeout(
            Duration::from_secs(DELEGATE_TIMEOUT_SECS),
            execute_with_provider(
                provider,
                config.system_prompt.as_deref(),
                &full_prompt,
                temperature,
            ),
        )
        .await;

        match result {
            Ok(Ok(response)) => {
                let output = if response.trim().is_empty() {
                    "[Empty response]".to_string()
                } else {
                    response
                };

                Ok(DelegateResult {
                    success: true,
                    output: format!(
                        "[Agent '{}' ({}/{})]\n{}",
                        agent_name, config.provider, config.model, output
                    ),
                    error: None,
                    agent: agent_name.to_string(),
                    provider: config.provider.clone(),
                    model: config.model.clone(),
                })
            }
            Ok(Err(e)) => Ok(DelegateResult {
                success: false,
                output: String::new(),
                error: Some(format!("Agent '{}' failed: {}", agent_name, e)),
                agent: agent_name.to_string(),
                provider: config.provider.clone(),
                model: config.model.clone(),
            }),
            Err(_) => Ok(DelegateResult {
                success: false,
                output: String::new(),
                error: Some(format!(
                    "Agent '{}' timed out after {}s",
                    agent_name, DELEGATE_TIMEOUT_SECS
                )),
                agent: agent_name.to_string(),
                provider: config.provider.clone(),
                model: config.model.clone(),
            }),
        }
    }
}

impl Default for MultiAgentSystem {
    fn default() -> Self {
        Self::new()
    }
}

async fn execute_with_provider(
    provider: Arc<dyn ModelProvider>,
    system_prompt: Option<&str>,
    prompt: &str,
    temperature: f32,
) -> Result<String> {
    use enact_core::providers::{ChatMessage, ChatRequest};

    let messages = if let Some(system) = system_prompt {
        vec![ChatMessage::system(system), ChatMessage::user(prompt)]
    } else {
        vec![ChatMessage::user(prompt)]
    };

    let request = ChatRequest {
        messages,
        max_tokens: Some(4096),
        temperature: Some(temperature),
    };

    let response = provider.chat(request).await?;

    Ok(response
        .choices
        .first()
        .map(|c| c.message.content.clone())
        .unwrap_or_default())
}

/// Result of a delegation operation
#[derive(Debug, Clone)]
pub struct DelegateResult {
    pub success: bool,
    pub output: String,
    pub error: Option<String>,
    pub agent: String,
    pub provider: String,
    pub model: String,
}

/// Agent registry for discovering and managing agents
pub struct AgentRegistry {
    agents: HashMap<String, AgentInfo>,
}

#[derive(Debug, Clone)]
pub struct AgentInfo {
    pub name: String,
    pub description: String,
    pub capabilities: Vec<String>,
    pub endpoint: Option<String>,
}

impl AgentRegistry {
    pub fn new() -> Self {
        Self {
            agents: HashMap::new(),
        }
    }

    pub fn register(
        &mut self,
        name: impl Into<String>,
        description: impl Into<String>,
        capabilities: Vec<String>,
        endpoint: Option<String>,
    ) {
        let name = name.into();
        self.agents.insert(
            name.clone(),
            AgentInfo {
                name,
                description: description.into(),
                capabilities,
                endpoint,
            },
        );
    }

    pub fn find_by_capability(&self, capability: &str) -> Vec<&AgentInfo> {
        self.agents
            .values()
            .filter(|agent| agent.capabilities.contains(&capability.to_string()))
            .collect()
    }

    pub fn get(&self, name: &str) -> Option<&AgentInfo> {
        self.agents.get(name)
    }

    pub fn list(&self) -> Vec<&AgentInfo> {
        self.agents.values().collect()
    }
}

impl Default for AgentRegistry {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_multi_agent_system_empty() {
        let system = MultiAgentSystem::new();
        assert!(system.get_agent_names().is_empty());
    }

    #[test]
    fn test_register_agent() {
        let mut system = MultiAgentSystem::new();
        system.register_agent(
            "researcher",
            DelegateAgentConfig {
                provider: "openai".to_string(),
                model: "gpt-4".to_string(),
                system_prompt: Some("You are a research assistant.".to_string()),
                api_key: None,
                temperature: Some(0.3),
                max_depth: 3,
            },
        );

        let names = system.get_agent_names();
        assert_eq!(names.len(), 1);
        assert_eq!(names[0], "researcher");
    }

    #[test]
    fn test_agent_registry() {
        let mut registry = AgentRegistry::new();
        registry.register(
            "coder",
            "Writes code",
            vec!["coding".to_string(), "rust".to_string()],
            None,
        );

        let coders = registry.find_by_capability("coding");
        assert_eq!(coders.len(), 1);
        assert_eq!(coders[0].name, "coder");
    }

    #[tokio::test]
    async fn test_delegate_unknown_agent() {
        let system = MultiAgentSystem::new();
        let result = system.delegate("unknown", "test", None, 0).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_delegate_depth_limit() {
        let mut system = MultiAgentSystem::new();
        system.register_agent(
            "test",
            DelegateAgentConfig {
                max_depth: 2,
                ..Default::default()
            },
        );

        let result = system.delegate("test", "prompt", None, 2).await.unwrap();
        assert!(!result.success);
        assert!(result.error.unwrap().contains("depth limit"));
    }
}
