use crate::config::DispatchStrategy;
use crate::node_registry::{NodeInfo, NodeRegistry};
use crate::task_queue::{DispatchTask, TaskId, TaskQueue};
use anyhow::Result;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

pub struct Dispatcher {
    strategy: DispatchStrategy,
    node_registry: Arc<dyn NodeRegistry>,
    task_queue: Arc<dyn TaskQueue>,
    rr_counter: AtomicUsize,
}

impl Dispatcher {
    pub fn new(
        strategy: DispatchStrategy,
        node_registry: Arc<dyn NodeRegistry>,
        task_queue: Arc<dyn TaskQueue>,
    ) -> Self {
        Self {
            strategy,
            node_registry,
            task_queue,
            rr_counter: AtomicUsize::new(0),
        }
    }

    pub async fn dispatch(
        &self,
        task: DispatchTask,
        required_capability: Option<&str>,
    ) -> Result<TaskId> {
        let healthy = self.node_registry.list_healthy().await?;
        let _selected = self.select_worker(&healthy, required_capability);

        // Current implementation enqueues to shared stream; worker group claims tasks.
        // Selection is computed for policy/observability and future per-worker routing.
        self.task_queue.enqueue(task).await
    }

    fn select_worker<'a>(
        &self,
        healthy: &'a [NodeInfo],
        required_capability: Option<&str>,
    ) -> Option<&'a NodeInfo> {
        let candidates: Vec<&NodeInfo> = if let Some(capability) = required_capability {
            healthy
                .iter()
                .filter(|n| n.capabilities.iter().any(|c| c == capability))
                .collect()
        } else {
            healthy.iter().collect()
        };

        if candidates.is_empty() {
            return None;
        }

        match self.strategy {
            DispatchStrategy::LeastLoaded => candidates.into_iter().min_by(|a, b| {
                a.load
                    .partial_cmp(&b.load)
                    .unwrap_or(std::cmp::Ordering::Equal)
            }),
            DispatchStrategy::RoundRobin => {
                let idx = self.rr_counter.fetch_add(1, Ordering::Relaxed);
                Some(candidates[idx % candidates.len()])
            }
            DispatchStrategy::CapabilityMatch => candidates.into_iter().min_by(|a, b| {
                a.load
                    .partial_cmp(&b.load)
                    .unwrap_or(std::cmp::Ordering::Equal)
            }),
        }
    }
}
