use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{mpsc, RwLock};
use tracing::{error, info};

use crate::agent::{AgentHandle, AgentId, AgentStatus};
use crate::config::SwarmConfig;
use crate::error::{SwarmError, SwarmResult};
use crate::memory::MemoryStore;
use crate::message::{Message, MessageKind};

pub struct Orchestrator {
    config: SwarmConfig,
    agents: Arc<RwLock<HashMap<AgentId, AgentHandle>>>,
    bus_tx: mpsc::Sender<Message>,
    bus_rx: Option<mpsc::Receiver<Message>>,
    memory: Arc<MemoryStore>,
}

impl Orchestrator {
    pub fn new(config: SwarmConfig, memory: MemoryStore) -> Self {
        let (bus_tx, bus_rx) = mpsc::channel(1024);
        Self {
            config,
            agents: Arc::new(RwLock::new(HashMap::new())),
            bus_tx,
            bus_rx: Some(bus_rx),
            memory: Arc::new(memory),
        }
    }

    pub fn bus_sender(&self) -> mpsc::Sender<Message> {
        self.bus_tx.clone()
    }

    pub fn memory(&self) -> &Arc<MemoryStore> {
        &self.memory
    }

    pub fn config(&self) -> &SwarmConfig {
        &self.config
    }

    pub fn agents(&self) -> &Arc<RwLock<HashMap<AgentId, AgentHandle>>> {
        &self.agents
    }

    pub async fn register_agent(&self, handle: AgentHandle) -> SwarmResult<()> {
        let id = handle.id.clone();
        let mut agents = self.agents.write().await;
        if agents.contains_key(&id) {
            return Err(SwarmError::AgentAlreadyExists(id.to_string()));
        }
        info!(agent = %id, "Registered agent");
        agents.insert(id, handle);
        Ok(())
    }

    pub async fn unregister_agent(&self, id: &AgentId) -> SwarmResult<AgentHandle> {
        let mut agents = self.agents.write().await;
        agents
            .remove(id)
            .ok_or_else(|| SwarmError::AgentNotFound(id.to_string()))
    }

    pub async fn send_to_agent(&self, id: &AgentId, msg: Message) -> SwarmResult<()> {
        let agents = self.agents.read().await;
        let handle = agents
            .get(id)
            .ok_or_else(|| SwarmError::AgentNotFound(id.to_string()))?;
        handle.send(msg).await
    }

    pub async fn broadcast(&self, msg: Message) {
        let agents = self.agents.read().await;
        for (id, handle) in agents.iter() {
            if let Err(e) = handle.send(msg.clone()).await {
                error!(agent = %id, error = %e, "Failed to broadcast to agent");
            }
        }
    }

    pub async fn list_agents(&self) -> Vec<(AgentId, AgentStatus)> {
        let agents = self.agents.read().await;
        agents
            .iter()
            .map(|(id, handle)| (id.clone(), handle.status.clone()))
            .collect()
    }

    pub async fn run(mut self) -> SwarmResult<()> {
        let mut bus_rx = self
            .bus_rx
            .take()
            .expect("Orchestrator::run called more than once");
        let agents = self.agents.clone();
        let memory = self.memory.clone();

        info!(
            name = %self.config.name,
            "Swarm orchestrator starting"
        );

        while let Some(msg) = bus_rx.recv().await {
            // Store every message
            let _ = memory.store_message(&msg.from, &msg);

            match &msg.kind {
                MessageKind::AgentToAgent => {
                    if let Some(ref target) = msg.to {
                        let agents_read = agents.read().await;
                        if let Some(handle) = agents_read.get(target) {
                            if let Err(e) = handle.send(msg.clone()).await {
                                error!(
                                    from = %msg.from,
                                    to = %target,
                                    error = %e,
                                    "Agent-to-agent delivery failed"
                                );
                            }
                        } else {
                            error!(to = %target, "Target agent not found");
                        }
                    }
                }
                MessageKind::Broadcast => {
                    let agents_read = agents.read().await;
                    for (id, handle) in agents_read.iter() {
                        if *id != msg.from {
                            let _ = handle.send(msg.clone()).await;
                        }
                    }
                }
                _ => {
                    // Route to target if specified, otherwise log
                    if let Some(ref target) = msg.to {
                        let agents_read = agents.read().await;
                        if let Some(handle) = agents_read.get(target) {
                            let _ = handle.send(msg.clone()).await;
                        }
                    }
                }
            }
        }

        info!("Swarm orchestrator stopped");
        Ok(())
    }
}
