// NATS message consumer routing tasks to executor pool
// Bridges NATS JetStream with executor channels

use std::collections::HashMap;

use tokio::sync::mpsc;
use tracing::{debug, warn};

use crate::messages::TaskAssignment;

/// NATS consumer routing tasks to agent executors
pub struct NatsConsumer {
    executor_pool: HashMap<String, mpsc::Sender<TaskAssignment>>,
}

impl NatsConsumer {
    /// Create new consumer
    pub fn new() -> Self {
        Self {
            executor_pool: HashMap::new(),
        }
    }

    /// Register executor for agent
    pub fn register_executor(&mut self, agent_id: String, sender: mpsc::Sender<TaskAssignment>) {
        debug!("Registered executor for agent: {}", agent_id);
        self.executor_pool.insert(agent_id, sender);
    }

    /// Unregister executor
    pub fn unregister_executor(&mut self, agent_id: &str) {
        self.executor_pool.remove(agent_id);
        debug!("Unregistered executor for agent: {}", agent_id);
    }

    /// Get executor sender for agent
    pub fn get_executor(&self, agent_id: &str) -> Option<&mpsc::Sender<TaskAssignment>> {
        self.executor_pool.get(agent_id)
    }

    /// Route task to agent executor
    pub async fn route_task(&self, task: TaskAssignment) -> Result<(), TaskRoutingError> {
        if let Some(tx) = self.executor_pool.get(&task.agent_id) {
            tx.send(task.clone())
                .await
                .map_err(|_| TaskRoutingError::ExecutorUnavailable(task.agent_id.clone()))?;
            Ok(())
        } else {
            warn!("No executor found for agent: {}", task.agent_id);
            Err(TaskRoutingError::AgentNotFound(task.agent_id.clone()))
        }
    }

    /// Get list of registered agents
    pub fn list_agents(&self) -> Vec<String> {
        self.executor_pool.keys().cloned().collect()
    }
}

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

/// Task routing errors
#[derive(Debug)]
pub enum TaskRoutingError {
    AgentNotFound(String),
    ExecutorUnavailable(String),
}

impl std::fmt::Display for TaskRoutingError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::AgentNotFound(agent_id) => write!(f, "Agent not found: {}", agent_id),
            Self::ExecutorUnavailable(agent_id) => {
                write!(f, "Executor unavailable for agent: {}", agent_id)
            }
        }
    }
}

impl std::error::Error for TaskRoutingError {}

#[cfg(test)]
mod tests {
    use chrono::Utc;

    use super::*;

    #[tokio::test]
    async fn test_consumer_registration() {
        let mut consumer = NatsConsumer::new();
        let (tx, _rx) = mpsc::channel(10);

        consumer.register_executor("agent-1".to_string(), tx);

        assert!(consumer.get_executor("agent-1").is_some());
        assert!(consumer.get_executor("agent-2").is_none());
    }

    #[tokio::test]
    async fn test_task_routing() {
        let mut consumer = NatsConsumer::new();
        let (tx, mut rx) = mpsc::channel(10);

        consumer.register_executor("agent-1".to_string(), tx);

        let task = TaskAssignment {
            id: "task-1".to_string(),
            agent_id: "agent-1".to_string(),
            required_role: "developer".to_string(),
            title: "Test task".to_string(),
            description: "Test description".to_string(),
            context: "{}".to_string(),
            priority: 1,
            deadline: None,
            assigned_at: Utc::now(),
        };

        let result = consumer.route_task(task.clone()).await;
        assert!(result.is_ok());

        // Verify task was received
        let received = rx.recv().await;
        assert!(received.is_some());
        assert_eq!(received.unwrap().id, "task-1");
    }

    #[tokio::test]
    async fn test_routing_to_nonexistent_agent() {
        let consumer = NatsConsumer::new();

        let task = TaskAssignment {
            id: "task-1".to_string(),
            agent_id: "nonexistent".to_string(),
            required_role: "developer".to_string(),
            title: "Test task".to_string(),
            description: "Test description".to_string(),
            context: "{}".to_string(),
            priority: 1,
            deadline: None,
            assigned_at: Utc::now(),
        };

        let result = consumer.route_task(task).await;
        assert!(result.is_err());
    }
}
