zc2 0.0.12

//! Request routing logic for optimal worker selection.

use std::sync::atomic::{AtomicUsize, Ordering};

use serde::{Deserialize, Serialize};

use super::worker::{Worker, WorkerRegistry};
use super::credits::CreditManager;

/// Routing strategy for worker selection
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum RoutingStrategy {
    /// Select worker with lowest price (default)
    #[default]
    BestPrice,
    /// Select worker with lowest latency
    BestLatency,
    /// Select worker with most available resources
    BestAvailability,
    /// Round-robin across all healthy workers
    RoundRobin,
    /// Random worker selection
    Random,
    /// Weighted random based on capacity
    WeightedCapacity,
}

impl RoutingStrategy {
    pub fn as_str(&self) -> &'static str {
        match self {
            RoutingStrategy::BestPrice => "best_price",
            RoutingStrategy::BestLatency => "best_latency",
            RoutingStrategy::BestAvailability => "best_availability",
            RoutingStrategy::RoundRobin => "round_robin",
            RoutingStrategy::Random => "random",
            RoutingStrategy::WeightedCapacity => "weighted_capacity",
        }
    }

    pub fn description(&self) -> &'static str {
        match self {
            RoutingStrategy::BestPrice => "Lowest cost per compute",
            RoutingStrategy::BestLatency => "Fastest response time",
            RoutingStrategy::BestAvailability => "Most available resources",
            RoutingStrategy::RoundRobin => "Even distribution",
            RoutingStrategy::Random => "Random selection",
            RoutingStrategy::WeightedCapacity => "Weighted by capacity",
        }
    }
}

impl std::str::FromStr for RoutingStrategy {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "best_price" | "price" | "cheap" | "cheapest" => Ok(RoutingStrategy::BestPrice),
            "best_latency" | "latency" | "fast" | "fastest" => Ok(RoutingStrategy::BestLatency),
            "best_availability" | "availability" | "available" => Ok(RoutingStrategy::BestAvailability),
            "round_robin" | "robin" | "rr" => Ok(RoutingStrategy::RoundRobin),
            "random" | "rand" => Ok(RoutingStrategy::Random),
            "weighted_capacity" | "weighted" | "capacity" => Ok(RoutingStrategy::WeightedCapacity),
            _ => Err(format!("Unknown routing strategy: {}", s)),
        }
    }
}

/// Resource requirements for a compute request
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceRequirements {
    /// Required CPU cores
    #[serde(default = "default_cpus")]
    pub cpus: f64,
    /// Required memory in bytes
    #[serde(default = "default_memory")]
    pub memory_bytes: u64,
    /// Required GPUs
    #[serde(default)]
    pub gpus: u32,
    /// Preferred worker type (optional)
    #[serde(default)]
    pub worker_type: Option<String>,
    /// Required tags (optional)
    #[serde(default)]
    pub tags: Vec<String>,
    /// Estimated duration in seconds (for cost estimation)
    #[serde(default = "default_duration")]
    pub estimated_duration_secs: f64,
    /// Routing strategy (optional, defaults to BestPrice)
    #[serde(default)]
    pub strategy: RoutingStrategy,
    /// Requested timeout in seconds. The request will be aborted after this duration.
    /// Workers whose max_timeout_secs is > 0 and < this value will be excluded.
    #[serde(default)]
    pub timeout_secs: f64,
    /// If true, require remote execution only (no local fallback)
    #[serde(default)]
    pub remote_only: bool,
    /// Maximum credits to spend on this job. Auto-derives an effective timeout
    /// so the job stops when its credit budget is exhausted.
    #[serde(default)]
    pub budget_credits: Option<f64>,
}

fn default_cpus() -> f64 { 1.0 }
fn default_memory() -> u64 { 1024 * 1024 * 1024 } // 1 GiB
fn default_duration() -> f64 { 1.0 }

impl Default for ResourceRequirements {
    fn default() -> Self {
        Self {
            cpus: 1.0,
            memory_bytes: 1024 * 1024 * 1024,
            gpus: 0,
            worker_type: None,
            tags: Vec::new(),
            estimated_duration_secs: 1.0,
            strategy: RoutingStrategy::default(),
            timeout_secs: 0.0, // 0 = no explicit timeout
            remote_only: false,
            budget_credits: None,
        }
    }
}

/// Routing decision result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RoutingDecision {
    /// Selected worker
    pub worker: Worker,
    /// Estimated cost in credits
    pub estimated_cost: f64,
    /// Reason for selection
    pub reason: String,
    /// Alternative workers considered
    pub alternatives_count: usize,
}

/// Routing error
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RoutingError {
    /// Error code
    pub code: String,
    /// Error message
    pub message: String,
}

impl RoutingError {
    pub fn no_workers() -> Self {
        Self {
            code: "NO_WORKERS".to_string(),
            message: "No workers available".to_string(),
        }
    }

    pub fn no_capacity(requirements: &ResourceRequirements) -> Self {
        Self {
            code: "NO_CAPACITY".to_string(),
            message: format!(
                "No worker has capacity for {} CPUs, {} bytes memory, {} GPUs",
                requirements.cpus, requirements.memory_bytes, requirements.gpus
            ),
        }
    }

    pub fn insufficient_credits(required: f64, available: f64) -> Self {
        Self {
            code: "INSUFFICIENT_CREDITS".to_string(),
            message: format!(
                "Insufficient credits: need {:.4}, have {:.4}",
                required, available
            ),
        }
    }

    pub fn rate_limited() -> Self {
        Self {
            code: "RATE_LIMITED".to_string(),
            message: "Rate limit exceeded".to_string(),
        }
    }

    pub fn worker_type_unavailable(worker_type: &str) -> Self {
        Self {
            code: "WORKER_TYPE_UNAVAILABLE".to_string(),
            message: format!("No workers of type '{}' available", worker_type),
        }
    }

    pub fn timeout_incompatible(timeout: f64) -> Self {
        Self {
            code: "TIMEOUT_INCOMPATIBLE".to_string(),
            message: format!(
                "No worker accepts timeout of {:.1}s. Reduce timeout or wait for a compatible worker.",
                timeout
            ),
        }
    }

    pub fn quota_exceeded() -> Self {
        Self {
            code: "QUOTA_EXCEEDED".to_string(),
            message: "All workers have exceeded their request quota for this time window".to_string(),
        }
    }
}

/// Router for selecting optimal workers
pub struct Router {
    /// Round-robin counter for even distribution
    round_robin_counter: AtomicUsize,
}

impl Router {
    /// Create a new router
    pub fn new() -> Self {
        Self {
            round_robin_counter: AtomicUsize::new(0),
        }
    }

    /// Get all available routing strategies
    pub fn available_strategies() -> Vec<RoutingStrategy> {
        vec![
            RoutingStrategy::BestPrice,
            RoutingStrategy::BestLatency,
            RoutingStrategy::BestAvailability,
            RoutingStrategy::RoundRobin,
            RoutingStrategy::Random,
            RoutingStrategy::WeightedCapacity,
        ]
    }

    /// Select a worker without credit/rate checks (for local mode).
    ///
    /// Uses round-robin for fair distribution across equal-priority workers.
    /// Atomically reserves a quota slot via `try_reserve_quota` — the caller
    /// must call `cancel_quota_reservation` if the request ultimately fails.
    pub fn select_worker_no_checks(
        &self,
        registry: &WorkerRegistry,
        requirements: &ResourceRequirements,
    ) -> Result<RoutingDecision, RoutingError> {
        let workers = registry.healthy();
        if workers.is_empty() {
            return Err(RoutingError::no_workers());
        }

        // Round-robin starting offset for fair distribution
        let n = workers.len();
        let start = self.round_robin_counter.fetch_add(1, Ordering::Relaxed);

        for i in 0..n {
            let worker = workers[(start + i) % n].clone();
            if registry.try_reserve_quota(&worker.id) {
                return Ok(RoutingDecision {
                    worker,
                    estimated_cost: 0.0,
                    reason: format!("Local round-robin: slot {}", (start + i) % n),
                    alternatives_count: n - 1,
                });
            }
        }

        Err(RoutingError::quota_exceeded())
    }

    /// Select a LOCAL worker only — filters to workers whose URI matches own_ip.
    /// Used in the publish-lock model: the broker first tries local, then offers to peers.
    pub fn select_local_worker(
        &self,
        registry: &WorkerRegistry,
        own_ip: Option<&str>,
        _requirements: &ResourceRequirements,
    ) -> Result<RoutingDecision, RoutingError> {
        let all = registry.healthy();
        // Local workers are always at 127.0.0.1/localhost; additionally accept own Tailscale IP.
        let local: Vec<Worker> = all.into_iter().filter(|w| {
            w.uri.contains("127.0.0.1") || w.uri.contains("localhost")
            || own_ip.map_or(false, |ip| w.uri.contains(ip))
        }).collect();

        if local.is_empty() {
            return Err(RoutingError::no_workers());
        }

        let n = local.len();
        let start = self.round_robin_counter.fetch_add(1, Ordering::Relaxed);

        for i in 0..n {
            let worker = local[(start + i) % n].clone();
            if registry.try_reserve_quota(&worker.id) {
                return Ok(RoutingDecision {
                    worker,
                    estimated_cost: 0.0,
                    reason: format!("Local worker (publish-lock): slot {}", (start + i) % n),
                    alternatives_count: n - 1,
                });
            }
        }

        Err(RoutingError::quota_exceeded())
    }

    /// Select the best worker for the given requirements.
    ///
    /// `balance` is the caller's pre-fetched credit balance for the `user_id`.
    /// The `credits` parameter is only used for rate-limit enforcement.
    pub fn select_worker(
        &self,
        registry: &WorkerRegistry,
        credits: &CreditManager,
        user_id: &str,
        balance: f64,
        requirements: &ResourceRequirements,
    ) -> Result<RoutingDecision, RoutingError> {
        // Check rate limit first
        if !credits.check_rate_limit(user_id) {
            return Err(RoutingError::rate_limited());
        }

        // Get all healthy workers
        let workers = registry.healthy();
        if workers.is_empty() {
            return Err(RoutingError::no_workers());
        }

        // Filter by worker type if specified
        let workers: Vec<Worker> = if let Some(ref wt) = requirements.worker_type {
            let filtered: Vec<Worker> = workers
                .into_iter()
                .filter(|w| &w.worker_type == wt)
                .collect();
            if filtered.is_empty() {
                return Err(RoutingError::worker_type_unavailable(wt));
            }
            filtered
        } else {
            workers
        };

        // Filter by required tags
        let workers: Vec<Worker> = if !requirements.tags.is_empty() {
            workers
                .into_iter()
                .filter(|w| requirements.tags.iter().all(|t| w.tags.contains(t)))
                .collect()
        } else {
            workers
        };

        // Filter by capacity
        let capable: Vec<Worker> = workers
            .iter()
            .filter(|w| w.can_handle(requirements.cpus, requirements.memory_bytes, requirements.gpus))
            .cloned()
            .collect();

        if capable.is_empty() {
            return Err(RoutingError::no_capacity(requirements));
        }

        // Filter by timeout compatibility: if the request has a timeout, exclude workers
        // whose max_timeout_secs > 0 (has a limit) and < requested timeout.
        let capable: Vec<Worker> = if requirements.timeout_secs > 0.0 {
            let filtered: Vec<Worker> = capable.into_iter()
                .filter(|w| w.max_timeout_secs <= 0.0 || w.max_timeout_secs >= requirements.timeout_secs)
                .collect();
            if filtered.is_empty() {
                return Err(RoutingError::timeout_incompatible(requirements.timeout_secs));
            }
            filtered
        } else {
            capable
        };

        // Atomically reserve a quota slot on the selected worker.
        // If a race causes the selected worker to be over-quota, remove it
        // from candidates and retry with the next-best worker.
        let mut candidates = capable;
        loop {
            if candidates.is_empty() {
                return Err(RoutingError::quota_exceeded());
            }

            let alternatives_count = candidates.len();
            let (best_worker, strategy_reason) = self.select_by_strategy(&candidates, requirements);

            // Credit check before reserving quota (cheap — avoids wasted slot)
            let estimated_cost = best_worker.pricing.estimate_cost(
                requirements.estimated_duration_secs,
            );
            if balance < estimated_cost {
                return Err(RoutingError::insufficient_credits(estimated_cost, balance));
            }

            // Atomic check+commit — only one goroutine wins for each quota slot
            if registry.try_reserve_quota(&best_worker.id) {
                let reason = format!("{} (checked {} workers)", strategy_reason, alternatives_count);
                return Ok(RoutingDecision {
                    worker: best_worker,
                    estimated_cost,
                    reason,
                    alternatives_count,
                });
            }

            // Lost the race — this worker just hit its quota; try the next best
            let loser_id = best_worker.id.clone();
            candidates.retain(|w| w.id != loser_id);
        }
    }

    /// Select worker based on routing strategy
    fn select_by_strategy(
        &self,
        workers: &[Worker],
        requirements: &ResourceRequirements,
    ) -> (Worker, String) {
        match requirements.strategy {
            RoutingStrategy::BestPrice => {
                // Sort by estimated cost (lowest first)
                let mut scored: Vec<(Worker, f64)> = workers
                    .iter()
                    .map(|w| {
                        let cost = w.pricing.estimate_cost(
                            requirements.estimated_duration_secs,
                        );
                        (w.clone(), cost)
                    })
                    .collect();
                scored.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
                let (worker, cost) = scored.into_iter().next().unwrap();
                (worker, format!("Best price: {:.6} credits", cost))
            }

            RoutingStrategy::BestLatency => {
                // Sort by average latency (lowest first)
                let mut scored: Vec<(Worker, f64)> = workers
                    .iter()
                    .map(|w| (w.clone(), w.avg_latency_ms))
                    .collect();
                scored.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
                let (worker, latency) = scored.into_iter().next().unwrap();
                (worker, format!("Best latency: {:.1}ms", latency))
            }

            RoutingStrategy::BestAvailability => {
                // Sort by available resources (highest first)
                let mut scored: Vec<(Worker, f64)> = workers
                    .iter()
                    .map(|w| {
                        // Score based on available CPUs, memory, and low active requests
                        let cpu_score = w.resources.cpus_available;
                        let mem_score = w.resources.memory_available as f64 / (1024.0 * 1024.0 * 1024.0);
                        let load_penalty = w.active_requests as f64 * 0.5;
                        let score = cpu_score + mem_score - load_penalty;
                        (w.clone(), score)
                    })
                    .collect();
                // Sort descending (highest availability first)
                scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
                let (worker, score) = scored.into_iter().next().unwrap();
                (worker, format!("Best availability: score {:.1}", score))
            }

            RoutingStrategy::RoundRobin => {
                // Simple round-robin across workers
                let idx = self.round_robin_counter.fetch_add(1, Ordering::Relaxed);
                let worker = workers[idx % workers.len()].clone();
                (worker, format!("Round-robin: index {}", idx % workers.len()))
            }

            RoutingStrategy::Random => {
                // Random selection using simple PRNG
                use std::time::{SystemTime, UNIX_EPOCH};
                let seed = SystemTime::now()
                    .duration_since(UNIX_EPOCH)
                    .unwrap_or_default()
                    .subsec_nanos() as usize;
                let idx = seed % workers.len();
                let worker = workers[idx].clone();
                (worker, format!("Random: index {}", idx))
            }

            RoutingStrategy::WeightedCapacity => {
                // Weighted random based on available capacity
                use std::time::{SystemTime, UNIX_EPOCH};

                let weights: Vec<f64> = workers
                    .iter()
                    .map(|w| {
                        // Weight by available CPUs and memory
                        let cpu_weight = w.resources.cpus_available.max(0.1);
                        let mem_weight = (w.resources.memory_available as f64 / (1024.0 * 1024.0 * 1024.0)).max(0.1);
                        cpu_weight * mem_weight
                    })
                    .collect();

                let total_weight: f64 = weights.iter().sum();

                // Simple random selection
                let seed = SystemTime::now()
                    .duration_since(UNIX_EPOCH)
                    .unwrap_or_default()
                    .subsec_nanos() as f64;
                let target = (seed / 1_000_000_000.0) * total_weight;

                let mut cumulative = 0.0;
                for (i, weight) in weights.iter().enumerate() {
                    cumulative += weight;
                    if cumulative >= target {
                        let worker = workers[i].clone();
                        return (worker, format!("Weighted capacity: weight {:.2}", weight));
                    }
                }

                // Fallback to first worker
                let worker = workers[0].clone();
                (worker, "Weighted capacity: fallback".to_string())
            }
        }
    }

    /// Get cost estimate without selecting a worker
    pub fn estimate_cost(
        &self,
        registry: &WorkerRegistry,
        requirements: &ResourceRequirements,
    ) -> Option<(f64, f64)> {
        let workers = registry.healthy();
        if workers.is_empty() {
            return None;
        }

        // Filter capable workers
        let capable: Vec<&Worker> = workers
            .iter()
            .filter(|w| w.can_handle(requirements.cpus, requirements.memory_bytes, requirements.gpus))
            .collect();

        if capable.is_empty() {
            return None;
        }

        // Calculate min and max costs
        let costs: Vec<f64> = capable
            .iter()
            .map(|w| {
                w.pricing.estimate_cost(
                    requirements.estimated_duration_secs,
                )
            })
            .collect();

        let min_cost = costs.iter().cloned().fold(f64::INFINITY, f64::min);
        let max_cost = costs.iter().cloned().fold(0.0, f64::max);

        Some((min_cost, max_cost))
    }

    /// List workers matching requirements (for preview)
    pub fn list_matching(
        &self,
        registry: &WorkerRegistry,
        requirements: &ResourceRequirements,
    ) -> Vec<(Worker, f64)> {
        let workers = registry.healthy();

        let mut matching: Vec<(Worker, f64)> = workers
            .iter()
            .filter(|w| {
                // Filter by type if specified
                if let Some(ref wt) = requirements.worker_type {
                    if &w.worker_type != wt {
                        return false;
                    }
                }
                // Filter by capacity
                w.can_handle(requirements.cpus, requirements.memory_bytes, requirements.gpus)
            })
            .map(|w| {
                let cost = w.pricing.estimate_cost(
                    requirements.estimated_duration_secs,
                );
                (w.clone(), cost)
            })
            .collect();

        // Sort by cost (lowest first)
        matching.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
        matching
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::broker::credits::CreditManager;
    use crate::broker::worker::{HardwareInfo, WorkerRegistration, WorkerResources, WorkerPricing, WorkerRegistry};

    // --- Helpers ---

    fn make_reg(name: &str, port: usize, cpus: f64, memory_gib: u64, gpus: u32, price_per_hour: f64) -> WorkerRegistration {
        WorkerRegistration {
            name: name.to_string(),
            uri: format!("http://127.0.0.1:{}", 3960 + port),
            worker_type: "zakuro".to_string(),
            resources: WorkerResources {
                cpus_total: cpus,
                cpus_available: cpus,
                memory_total: memory_gib * 1024 * 1024 * 1024,
                memory_available: memory_gib * 1024 * 1024 * 1024,
                gpus_total: gpus,
                gpus_available: gpus,
            },
            pricing: WorkerPricing {
                price_per_hour,
                min_charge: 0.001,
            },
            tags: vec![],
            max_timeout_secs: 0.0,
            hardware: HardwareInfo::default(),
            tailscale_ip: None,
            is_docker: None,
        }
    }

    fn credits_for(user: &str, balance: f64) -> CreditManager {
        let mgr = CreditManager::new();
        mgr.get_or_create(user, balance);
        mgr
    }

    fn req_default() -> ResourceRequirements {
        ResourceRequirements::default()
    }

    // --- Error cases ---

    #[test]
    fn test_no_healthy_workers_returns_no_workers() {
        let router = Router::new();
        let registry = WorkerRegistry::new();
        let credits = credits_for("alice", 1000.0);

        let err = router.select_worker(&registry, &credits, "alice", 1000.0, &req_default()).unwrap_err();
        assert_eq!(err.code, "NO_WORKERS");
    }

    #[test]
    fn test_no_capacity_returns_no_capacity() {
        let router = Router::new();
        let registry = WorkerRegistry::new();
        registry.register(make_reg("small", 0, 1.0, 1, 0, 0.001));
        let credits = credits_for("alice", 1000.0);

        let req = ResourceRequirements { cpus: 16.0, ..req_default() };
        let err = router.select_worker(&registry, &credits, "alice", 1000.0, &req).unwrap_err();
        assert_eq!(err.code, "NO_CAPACITY");
    }

    #[test]
    fn test_insufficient_credits_returns_error() {
        let router = Router::new();
        let registry = WorkerRegistry::new();
        // Expensive worker: 1.0 credit/CPU/s
        registry.register(make_reg("expensive", 0, 4.0, 4, 0, 1.0));
        let credits = credits_for("alice", 0.001); // near-zero balance

        let req = ResourceRequirements {
            cpus: 1.0,
            estimated_duration_secs: 3600.0, // 1 hour → very expensive
            ..req_default()
        };
        let err = router.select_worker(&registry, &credits, "alice", 0.001, &req).unwrap_err();
        assert_eq!(err.code, "INSUFFICIENT_CREDITS");
    }

    #[test]
    fn test_rate_limited_returns_error() {
        let router = Router::new();
        let registry = WorkerRegistry::new();
        registry.register(make_reg("w", 0, 4.0, 4, 0, 0.001));

        let mgr = CreditManager::new();
        mgr.get_or_create("alice", 1000.0);
        mgr.set_rate_limits("alice", Some(0), None, None); // 0/sec → immediate block

        let err = router.select_worker(&registry, &mgr, "alice", 1000.0, &req_default()).unwrap_err();
        assert_eq!(err.code, "RATE_LIMITED");
    }

    #[test]
    fn test_worker_type_mismatch_returns_error() {
        let router = Router::new();
        let registry = WorkerRegistry::new();
        let mut reg = make_reg("ray-w", 0, 4.0, 4, 0, 0.001);
        reg.worker_type = "ray".to_string();
        registry.register(reg);
        let credits = credits_for("alice", 1000.0);

        let req = ResourceRequirements { worker_type: Some("spark".to_string()), ..req_default() };
        let err = router.select_worker(&registry, &credits, "alice", 1000.0, &req).unwrap_err();
        assert_eq!(err.code, "WORKER_TYPE_UNAVAILABLE");
    }

    #[test]
    fn test_timeout_incompatible_returns_error() {
        let router = Router::new();
        let registry = WorkerRegistry::new();
        let mut reg = make_reg("limited", 0, 4.0, 4, 0, 0.001);
        reg.max_timeout_secs = 30.0; // only accepts ≤ 30s requests
        registry.register(reg);
        let credits = credits_for("alice", 1000.0);

        let req = ResourceRequirements { timeout_secs: 60.0, ..req_default() };
        let err = router.select_worker(&registry, &credits, "alice", 1000.0, &req).unwrap_err();
        assert_eq!(err.code, "TIMEOUT_INCOMPATIBLE");
    }

    #[test]
    fn test_timeout_unlimited_worker_accepts_any_timeout() {
        let router = Router::new();
        let registry = WorkerRegistry::new();
        let mut reg = make_reg("unlimited", 0, 4.0, 4, 0, 0.001);
        reg.max_timeout_secs = 0.0; // 0 = unlimited
        registry.register(reg);
        let credits = credits_for("alice", 1000.0);

        let req = ResourceRequirements { timeout_secs: 3600.0, ..req_default() };
        assert!(router.select_worker(&registry, &credits, "alice", 1000.0, &req).is_ok());
    }

    // --- Tag filtering ---

    #[test]
    fn test_tag_filter_requires_all_tags() {
        let router = Router::new();
        let registry = WorkerRegistry::new();
        let mut reg = make_reg("gpu-w", 0, 8.0, 32, 2, 0.001);
        reg.tags = vec!["gpu".to_string(), "a100".to_string()];
        registry.register(reg);
        let credits = credits_for("alice", 1000.0);

        // Matching both tags
        let req_match = ResourceRequirements {
            tags: vec!["gpu".to_string(), "a100".to_string()],
            ..req_default()
        };
        assert!(router.select_worker(&registry, &credits, "alice", 1000.0, &req_match).is_ok());

        // Requesting a tag that doesn't exist on any worker
        let req_no_match = ResourceRequirements {
            tags: vec!["gpu".to_string(), "h100".to_string()],
            ..req_default()
        };
        assert!(router.select_worker(&registry, &credits, "alice", 1000.0, &req_no_match).is_err());
    }

    // --- Strategy tests ---

    #[test]
    fn test_best_price_selects_cheapest_worker() {
        let router = Router::new();
        let registry = WorkerRegistry::new();
        registry.register(make_reg("cheap", 0, 4.0, 4, 0, 0.001));
        registry.register(make_reg("pricey", 1, 4.0, 4, 0, 1.0));
        let credits = credits_for("alice", 10000.0);

        // Use 1-hour duration so that price_per_hour values are directly comparable
        let req = ResourceRequirements {
            strategy: RoutingStrategy::BestPrice,
            estimated_duration_secs: 3600.0,
            ..req_default()
        };
        let decision = router.select_worker(&registry, &credits, "alice", 10000.0, &req).unwrap();
        assert_eq!(decision.worker.name, "cheap");
    }

    #[test]
    fn test_best_latency_selects_fastest_worker() {
        let router = Router::new();
        let registry = WorkerRegistry::new();
        let w_fast = registry.register(make_reg("fast", 0, 4.0, 4, 0, 0.001));
        let w_slow = registry.register(make_reg("slow", 1, 4.0, 4, 0, 0.001));

        // Simulate latency via record_request
        // fast: EMA = 0.1 * 5 = 0.5ms; slow: EMA = 0.1 * 500 = 50ms
        registry.record_request(&w_fast.id, 5.0, true);
        registry.record_request(&w_slow.id, 500.0, true);

        let credits = credits_for("alice", 1000.0);
        let req = ResourceRequirements { strategy: RoutingStrategy::BestLatency, ..req_default() };
        let decision = router.select_worker(&registry, &credits, "alice", 1000.0, &req).unwrap();
        assert_eq!(decision.worker.name, "fast");
    }

    #[test]
    fn test_round_robin_cycles_through_workers() {
        let router = Router::new();
        let registry = WorkerRegistry::new();
        registry.register(make_reg("w1", 0, 4.0, 4, 0, 0.001));
        registry.register(make_reg("w2", 1, 4.0, 4, 0, 0.001));
        let credits = credits_for("alice", 10000.0);

        let req = ResourceRequirements { strategy: RoutingStrategy::RoundRobin, ..req_default() };
        let d1 = router.select_worker(&registry, &credits, "alice", 10000.0, &req).unwrap();
        let d2 = router.select_worker(&registry, &credits, "alice", 10000.0, &req).unwrap();

        // Two consecutive calls must hit different workers
        assert_ne!(d1.worker.id, d2.worker.id);
    }

    #[test]
    fn test_local_mode_returns_first_worker_free() {
        let router = Router::new();
        let registry = WorkerRegistry::new();
        registry.register(make_reg("local-w", 0, 4.0, 4, 0, 0.001));

        let req = req_default();
        let decision = router.select_worker_no_checks(&registry, &req).unwrap();
        assert_eq!(decision.estimated_cost, 0.0);
    }

    #[test]
    fn test_local_mode_no_workers_returns_error() {
        let router = Router::new();
        let registry = WorkerRegistry::new();
        let err = router.select_worker_no_checks(&registry, &req_default()).unwrap_err();
        assert_eq!(err.code, "NO_WORKERS");
    }

    // --- Cost estimation ---

    #[test]
    fn test_estimate_cost_returns_min_and_max() {
        let router = Router::new();
        let registry = WorkerRegistry::new();
        registry.register(make_reg("cheap", 0, 4.0, 4, 0, 0.001));
        registry.register(make_reg("pricey", 1, 4.0, 4, 0, 0.1));

        // Use 1-hour duration: cheap=0.001 credits, pricey=0.1 credits
        let req = ResourceRequirements { cpus: 1.0, estimated_duration_secs: 3600.0, ..req_default() };
        let (min, max) = router.estimate_cost(&registry, &req).unwrap();
        assert!(min < max);
        // min: 0.001 credits/hour * 1 hour = 0.001 credits
        assert!((min - 0.001).abs() < 0.00001);
    }

    #[test]
    fn test_estimate_cost_no_workers_returns_none() {
        let router = Router::new();
        let registry = WorkerRegistry::new();
        assert!(router.estimate_cost(&registry, &req_default()).is_none());
    }

    // --- Worker type match ---

    #[test]
    fn test_worker_type_match_routes_correctly() {
        let router = Router::new();
        let registry = WorkerRegistry::new();
        let mut reg = make_reg("ray-w", 0, 4.0, 4, 0, 0.001);
        reg.worker_type = "ray".to_string();
        registry.register(reg);
        let credits = credits_for("alice", 1000.0);

        let req = ResourceRequirements { worker_type: Some("ray".to_string()), ..req_default() };
        let decision = router.select_worker(&registry, &credits, "alice", 1000.0, &req).unwrap();
        assert_eq!(decision.worker.worker_type, "ray");
    }

    // --- Price change scenario ---

    #[test]
    fn test_price_change_affects_routing_decision() {
        let router = Router::new();
        let registry = WorkerRegistry::new();

        // Two workers: w1 initially cheap, w2 initially expensive
        let w1 = registry.register(make_reg("w1", 0, 4.0, 4, 0, 0.001));
        let w2 = registry.register(make_reg("w2", 1, 4.0, 4, 0, 0.01));

        let credits = credits_for("alice", 10000.0);
        // Use 1-hour duration so price_per_hour values are directly comparable
        let req = ResourceRequirements {
            strategy: RoutingStrategy::BestPrice,
            estimated_duration_secs: 3600.0,
            ..req_default()
        };

        // w1 should win (cheaper)
        let d1 = router.select_worker(&registry, &credits, "alice", 10000.0, &req).unwrap();
        assert_eq!(d1.worker.name, "w1");

        // Verify min cost matches w1's price_per_hour (0.001 credits for 1 hour)
        let (min, _max) = router.estimate_cost(&registry, &req).unwrap();
        assert!((min - 0.001).abs() < 0.00001);
    }
}