// vapora-agents: Agent registry - manages agent lifecycle and availability

use std::collections::HashMap;
use std::sync::{Arc, RwLock};

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use uuid::Uuid;

#[derive(Debug, Error)]
pub enum RegistryError {
    #[error("Agent not found: {0}")]
    AgentNotFound(String),

    #[error("Agent already registered: {0}")]
    AgentAlreadyRegistered(String),

    #[error("Maximum agents reached for role: {0}")]
    MaxAgentsReached(String),

    #[error("Invalid agent state transition: {0}")]
    InvalidStateTransition(String),
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum AgentStatus {
    Active,
    Inactive,
    Updating,
    Error(String),
    Scaling,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentMetadata {
    /// Ephemeral UUID generated at runtime — changes every startup.
    pub id: String,
    /// Deterministic identity keyed on role — survives hot-reloads.
    /// Defaults to `role` for data produced before this field existed.
    #[serde(default)]
    pub stable_id: String,
    pub role: String,
    pub name: String,
    pub version: String,
    pub status: AgentStatus,
    pub capabilities: Vec<String>,
    pub llm_provider: String,
    pub llm_model: String,
    pub max_concurrent_tasks: u32,
    pub current_tasks: u32,
    pub created_at: DateTime<Utc>,
    pub last_heartbeat: DateTime<Utc>,
    pub uptime_percentage: f64,
    pub total_tasks_completed: u64,
    /// Domain-optimized system prompt from capability package.
    /// None means the executor uses its default generic prompt.
    #[serde(default)]
    pub system_prompt: Option<String>,
}

impl AgentMetadata {
    pub fn new(
        role: String,
        name: String,
        llm_provider: String,
        llm_model: String,
        capabilities: Vec<String>,
    ) -> Self {
        let now = Utc::now();
        Self {
            id: Uuid::new_v4().to_string(),
            stable_id: role.clone(),
            role,
            name,
            version: "0.1.0".to_string(),
            status: AgentStatus::Active,
            capabilities,
            llm_provider,
            llm_model,
            max_concurrent_tasks: 5,
            current_tasks: 0,
            created_at: now,
            last_heartbeat: now,
            uptime_percentage: 100.0,
            total_tasks_completed: 0,
            system_prompt: None,
        }
    }

    /// Builder: attach a domain-optimized system prompt from a capability
    /// package.
    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
        self.system_prompt = Some(prompt.into());
        self
    }

    /// Stable profile identity: `stable_id` when non-empty, `role` otherwise.
    ///
    /// The fallback handles data deserialized from records produced before
    /// `stable_id` was introduced.
    pub fn stable_id_or_role(&self) -> &str {
        if self.stable_id.is_empty() {
            &self.role
        } else {
            &self.stable_id
        }
    }

    /// Check if agent can accept new tasks
    pub fn can_accept_task(&self) -> bool {
        self.status == AgentStatus::Active && self.current_tasks < self.max_concurrent_tasks
    }

    /// Increment task count
    pub fn assign_task(&mut self) {
        if self.current_tasks < self.max_concurrent_tasks {
            self.current_tasks += 1;
        }
    }

    /// Decrement task count
    pub fn complete_task(&mut self) {
        if self.current_tasks > 0 {
            self.current_tasks -= 1;
        }
        self.total_tasks_completed += 1;
    }
}

/// Thread-safe agent registry
#[derive(Clone)]
pub struct AgentRegistry {
    inner: Arc<RwLock<AgentRegistryInner>>,
}

struct AgentRegistryInner {
    agents: HashMap<String, AgentMetadata>,
    running_count: HashMap<String, u32>,
    max_agents_per_role: u32,
}

impl AgentRegistry {
    pub fn new(max_agents_per_role: u32) -> Self {
        Self {
            inner: Arc::new(RwLock::new(AgentRegistryInner {
                agents: HashMap::new(),
                running_count: HashMap::new(),
                max_agents_per_role,
            })),
        }
    }

    /// Register a new agent
    pub fn register_agent(&self, metadata: AgentMetadata) -> Result<String, RegistryError> {
        let mut inner = self.inner.write().expect("Failed to acquire write lock");

        // Check if agent already registered
        if inner.agents.contains_key(&metadata.id) {
            return Err(RegistryError::AgentAlreadyRegistered(metadata.id.clone()));
        }

        // Check if we've reached max agents for this role
        let count = inner.running_count.get(&metadata.role).unwrap_or(&0);
        if *count >= inner.max_agents_per_role {
            return Err(RegistryError::MaxAgentsReached(metadata.role.clone()));
        }

        let role = metadata.role.clone();
        let id = metadata.id.clone();

        inner.agents.insert(id.clone(), metadata);
        *inner.running_count.entry(role).or_insert(0) += 1;

        Ok(id)
    }

    /// Unregister an agent
    pub fn unregister_agent(&self, id: &str) -> Result<(), RegistryError> {
        let mut inner = self.inner.write().expect("Failed to acquire write lock");

        let agent = inner
            .agents
            .remove(id)
            .ok_or_else(|| RegistryError::AgentNotFound(id.to_string()))?;

        if let Some(count) = inner.running_count.get_mut(&agent.role) {
            if *count > 0 {
                *count -= 1;
            }
        }

        Ok(())
    }

    /// Get agent metadata
    pub fn get_agent(&self, id: &str) -> Option<AgentMetadata> {
        let inner = self.inner.read().expect("Failed to acquire read lock");
        inner.agents.get(id).cloned()
    }

    /// Get all agents for a specific role
    pub fn get_agents_by_role(&self, role: &str) -> Vec<AgentMetadata> {
        let inner = self.inner.read().expect("Failed to acquire read lock");
        inner
            .agents
            .values()
            .filter(|a| a.role == role && a.status == AgentStatus::Active)
            .cloned()
            .collect()
    }

    /// List all agents
    pub fn list_all(&self) -> Vec<AgentMetadata> {
        let inner = self.inner.read().expect("Failed to acquire read lock");
        inner.agents.values().cloned().collect()
    }

    /// Update agent status
    pub fn update_agent_status(&self, id: &str, status: AgentStatus) -> Result<(), RegistryError> {
        let mut inner = self.inner.write().expect("Failed to acquire write lock");

        let agent = inner
            .agents
            .get_mut(id)
            .ok_or_else(|| RegistryError::AgentNotFound(id.to_string()))?;

        agent.status = status;
        agent.last_heartbeat = Utc::now();

        Ok(())
    }

    /// Update agent heartbeat
    pub fn heartbeat(&self, id: &str) -> Result<(), RegistryError> {
        let mut inner = self.inner.write().expect("Failed to acquire write lock");

        let agent = inner
            .agents
            .get_mut(id)
            .ok_or_else(|| RegistryError::AgentNotFound(id.to_string()))?;

        agent.last_heartbeat = Utc::now();

        Ok(())
    }

    /// Get an available agent for a specific role
    pub fn get_available_agent(&self, role: &str) -> Option<AgentMetadata> {
        let agents = self.get_agents_by_role(role);

        agents
            .into_iter()
            .filter(|a| a.can_accept_task())
            .min_by_key(|a| a.current_tasks)
    }

    /// Assign task to agent
    pub fn assign_task(&self, agent_id: &str) -> Result<(), RegistryError> {
        let mut inner = self.inner.write().expect("Failed to acquire write lock");

        let agent = inner
            .agents
            .get_mut(agent_id)
            .ok_or_else(|| RegistryError::AgentNotFound(agent_id.to_string()))?;

        if !agent.can_accept_task() {
            return Err(RegistryError::InvalidStateTransition(
                "Agent cannot accept more tasks".to_string(),
            ));
        }

        agent.assign_task();

        Ok(())
    }

    /// Complete task for agent
    pub fn complete_task(&self, agent_id: &str) -> Result<(), RegistryError> {
        let mut inner = self.inner.write().expect("Failed to acquire write lock");

        let agent = inner
            .agents
            .get_mut(agent_id)
            .ok_or_else(|| RegistryError::AgentNotFound(agent_id.to_string()))?;

        agent.complete_task();

        Ok(())
    }

    /// Get count of agents by role
    pub fn count_by_role(&self, role: &str) -> u32 {
        let inner = self.inner.read().expect("Failed to acquire read lock");
        *inner.running_count.get(role).unwrap_or(&0)
    }

    /// Get total agent count
    pub fn total_count(&self) -> usize {
        let inner = self.inner.read().expect("Failed to acquire read lock");
        inner.agents.len()
    }

    /// Remove all agents for `role` from the registry and return their
    /// ephemeral IDs. The `running_count` entry for the role is also cleared.
    ///
    /// Called during hot-reload to drain a role before re-spawning executors.
    /// Learning profiles (keyed by `stable_id`) are unaffected.
    pub fn drain_role(&self, role: &str) -> Vec<String> {
        let mut inner = self.inner.write().expect("registry write lock");
        let ids: Vec<String> = inner
            .agents
            .values()
            .filter(|a| a.role == role)
            .map(|a| a.id.clone())
            .collect();
        for id in &ids {
            inner.agents.remove(id);
        }
        inner.running_count.remove(role);
        ids
    }

    /// Return the set of roles currently present in the registry.
    pub fn list_roles(&self) -> Vec<String> {
        let inner = self.inner.read().expect("registry read lock");
        let mut roles: Vec<String> = inner
            .agents
            .values()
            .map(|a| a.role.clone())
            .collect::<std::collections::HashSet<_>>()
            .into_iter()
            .collect();
        roles.sort();
        roles
    }
}

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

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

    #[test]
    fn test_agent_registration() {
        let registry = AgentRegistry::new(5);

        let agent = AgentMetadata::new(
            "developer".to_string(),
            "Developer Agent 1".to_string(),
            "claude".to_string(),
            "claude-sonnet-4".to_string(),
            vec!["coding".to_string()],
        );

        let id = registry.register_agent(agent).unwrap();
        assert!(registry.get_agent(&id).is_some());
        assert_eq!(registry.total_count(), 1);
    }

    #[test]
    fn test_max_agents_per_role() {
        let registry = AgentRegistry::new(2);

        for i in 0..2 {
            let agent = AgentMetadata::new(
                "developer".to_string(),
                format!("Developer {}", i),
                "claude".to_string(),
                "claude-sonnet-4".to_string(),
                vec![],
            );
            registry.register_agent(agent).unwrap();
        }

        // Third agent should fail
        let agent = AgentMetadata::new(
            "developer".to_string(),
            "Developer 3".to_string(),
            "claude".to_string(),
            "claude-sonnet-4".to_string(),
            vec![],
        );

        let result = registry.register_agent(agent);
        assert!(result.is_err());
    }

    #[test]
    fn test_agent_task_assignment() {
        let _registry = AgentRegistry::new(5);

        let mut agent = AgentMetadata::new(
            "developer".to_string(),
            "Developer Agent".to_string(),
            "claude".to_string(),
            "claude-sonnet-4".to_string(),
            vec![],
        );

        assert_eq!(agent.current_tasks, 0);
        assert!(agent.can_accept_task());

        agent.assign_task();
        assert_eq!(agent.current_tasks, 1);

        agent.complete_task();
        assert_eq!(agent.current_tasks, 0);
        assert_eq!(agent.total_tasks_completed, 1);
    }

    #[test]
    fn test_stable_id_deterministic() {
        let a = AgentMetadata::new(
            "developer".to_string(),
            "Dev 1".to_string(),
            "claude".to_string(),
            "claude-sonnet-4".to_string(),
            vec![],
        );
        let b = AgentMetadata::new(
            "developer".to_string(),
            "Dev 2".to_string(),
            "claude".to_string(),
            "claude-sonnet-4".to_string(),
            vec![],
        );

        assert_ne!(a.id, b.id, "ephemeral IDs must differ");
        assert_eq!(a.stable_id, "developer");
        assert_eq!(b.stable_id, "developer");
        assert_eq!(a.stable_id_or_role(), "developer");
    }

    #[test]
    fn test_stable_id_or_role_fallback() {
        let mut agent = AgentMetadata::new(
            "reviewer".to_string(),
            "Rev".to_string(),
            "claude".to_string(),
            "claude-sonnet-4".to_string(),
            vec![],
        );
        agent.stable_id = String::new();
        assert_eq!(agent.stable_id_or_role(), "reviewer");
    }

    #[test]
    fn test_drain_role() {
        let registry = AgentRegistry::new(5);

        for i in 0..2 {
            let agent = AgentMetadata::new(
                "developer".to_string(),
                format!("Dev {}", i),
                "claude".to_string(),
                "claude-sonnet-4".to_string(),
                vec![],
            );
            registry.register_agent(agent).unwrap();
        }
        let reviewer = AgentMetadata::new(
            "reviewer".to_string(),
            "Rev".to_string(),
            "claude".to_string(),
            "claude-sonnet-4".to_string(),
            vec![],
        );
        registry.register_agent(reviewer).unwrap();

        let drained = registry.drain_role("developer");
        assert_eq!(drained.len(), 2);
        assert_eq!(registry.count_by_role("developer"), 0);
        assert_eq!(registry.count_by_role("reviewer"), 1);
        assert_eq!(registry.total_count(), 1);
    }

    #[test]
    fn test_list_roles() {
        let registry = AgentRegistry::new(5);

        let a = AgentMetadata::new(
            "developer".to_string(),
            "Dev".to_string(),
            "claude".to_string(),
            "claude-sonnet-4".to_string(),
            vec![],
        );
        let b = AgentMetadata::new(
            "reviewer".to_string(),
            "Rev".to_string(),
            "claude".to_string(),
            "claude-sonnet-4".to_string(),
            vec![],
        );
        registry.register_agent(a).unwrap();
        registry.register_agent(b).unwrap();

        let mut roles = registry.list_roles();
        roles.sort();
        assert_eq!(roles, vec!["developer", "reviewer"]);
    }

    #[test]
    fn test_get_available_agent() {
        let registry = AgentRegistry::new(5);

        let agent1 = AgentMetadata::new(
            "developer".to_string(),
            "Developer 1".to_string(),
            "claude".to_string(),
            "claude-sonnet-4".to_string(),
            vec![],
        );

        let id1 = registry.register_agent(agent1).unwrap();

        let available = registry.get_available_agent("developer");
        assert!(available.is_some());

        // Assign tasks to fill capacity
        for _ in 0..5 {
            registry.assign_task(&id1).unwrap();
        }

        // Should no longer be available
        let available = registry.get_available_agent("developer");
        assert!(available.is_none());
    }
}
