zc2 0.0.29

//! Worker registry for tracking available compute workers.

use chrono::{DateTime, Duration, Utc};
use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use std::collections::{HashSet, VecDeque};
use std::sync::Mutex;

/// Wildcard `served_models` entry meaning "serves any model" (used by general providers).
pub const MODEL_WILDCARD: &str = "*";

/// Whether a worker serves a fixed catalog of published models (`Specialized`,
/// indexed by model uuid) or can serve any model (`General`, e.g. a generic
/// inference backend).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ProviderType {
    /// Serves any requested model.
    General,
    /// Serves only the models listed in `served_models`.
    #[default]
    Specialized,
}

/// Worker status
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum WorkerStatus {
    /// Worker is healthy and accepting requests
    Healthy,
    /// Worker is busy but still accepting requests
    Busy,
    /// Worker is unhealthy or unreachable
    Unhealthy,
    /// Worker is draining (not accepting new requests)
    Draining,
}

/// Resource availability on a worker
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkerResources {
    /// Total CPU cores
    pub cpus_total: f64,
    /// Available CPU cores
    pub cpus_available: f64,
    /// Total memory in bytes
    pub memory_total: u64,
    /// Available memory in bytes
    pub memory_available: u64,
    /// Total GPUs
    pub gpus_total: u32,
    /// Available GPUs
    pub gpus_available: u32,
}

impl Default for WorkerResources {
    fn default() -> Self {
        Self {
            cpus_total: 1.0,
            cpus_available: 1.0,
            memory_total: 1024 * 1024 * 1024, // 1 GiB
            memory_available: 1024 * 1024 * 1024,
            gpus_total: 0,
            gpus_available: 0,
        }
    }
}

/// Aggregated free-resource snapshot across a broker's healthy local workers.
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct BrokerResources {
    /// Sum of available CPU cores across healthy workers (floored to whole cores).
    pub cpus_available: u64,
    /// Sum of available memory (bytes) across healthy workers.
    pub memory_available: u64,
    /// Sum of available GPUs across healthy workers.
    pub gpus_available: u32,
    /// Count of workers included in the aggregation.
    pub healthy_workers: u32,
}

pub(crate) fn default_price_per_hour() -> f64 {
    3.6
}
fn default_min_charge() -> f64 {
    0.001
}

/// Pricing information for a worker
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkerPricing {
    /// Flat hourly rate for the worker (in credits)
    #[serde(default = "default_price_per_hour")]
    pub price_per_hour: f64,
    /// Minimum charge per request (in credits)
    #[serde(default = "default_min_charge")]
    pub min_charge: f64,
}

impl Default for WorkerPricing {
    fn default() -> Self {
        Self {
            price_per_hour: 3.6, // 0.001 credits per second
            min_charge: 0.001,   // Minimum 0.001 credits per request
        }
    }
}

impl WorkerPricing {
    /// Calculate estimated cost for a request
    pub fn estimate_cost(&self, duration_secs: f64) -> f64 {
        (self.price_per_hour / 3600.0 * duration_secs).max(self.min_charge)
    }

    /// Calculate price score for worker selection (lower is better).
    /// Uses a 10-second reference job to weight min_charge fairly against hourly rate.
    pub fn price_score(&self) -> f64 {
        self.estimate_cost(10.0)
    }
}

/// Per-worker request quotas (max requests per time window)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkerQuotas {
    /// Max requests allowed in a rolling 5-hour window (0 = unlimited)
    pub per_5h: u64,
    /// Max requests allowed in a rolling 7-day window (0 = unlimited)
    pub per_week: u64,
    /// Max requests allowed in a rolling 30-day window (0 = unlimited)
    pub per_month: u64,
}

impl Default for WorkerQuotas {
    fn default() -> Self {
        Self {
            per_5h: std::env::var("ZAKURO_QUOTA_5H")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(0),
            per_week: std::env::var("ZAKURO_QUOTA_WEEK")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(0),
            per_month: std::env::var("ZAKURO_QUOTA_MONTH")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(0),
        }
    }
}

/// Hardware details reported by the worker's /info endpoint
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct HardwareInfo {
    /// GPU model name (e.g. "A100 80GB", "RTX 4090")
    #[serde(default)]
    pub gpu_model: Option<String>,
    /// GPU VRAM in GiB
    #[serde(default)]
    pub gpu_vram_gb: Option<u32>,
    /// CPU model name (e.g. "AMD EPYC 7543")
    #[serde(default)]
    pub cpu_model: Option<String>,
    /// Total storage in GiB
    #[serde(default)]
    pub storage_gb: Option<u32>,
}

/// Worker information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Worker {
    /// Unique worker ID
    pub id: String,
    /// Worker name/label
    pub name: String,
    /// Worker endpoint URI
    pub uri: String,
    /// Worker type (ray, dask, spark, zakuro)
    pub worker_type: String,
    /// Current status
    pub status: WorkerStatus,
    /// Available resources
    pub resources: WorkerResources,
    /// Pricing information
    pub pricing: WorkerPricing,
    /// Last heartbeat timestamp
    pub last_heartbeat: DateTime<Utc>,
    /// Number of active requests
    pub active_requests: u32,
    /// Total requests processed
    pub total_requests: u64,
    /// Average request duration in milliseconds
    pub avg_latency_ms: f64,
    /// Tags for filtering
    pub tags: Vec<String>,
    /// Maximum timeout this worker allows per request (seconds), 0 = unlimited
    #[serde(default)]
    pub max_timeout_secs: f64,
    /// Hardware details (GPU model, CPU model, etc.)
    #[serde(default)]
    pub hardware: HardwareInfo,
    /// WireGuard IP extracted from worker's URI
    #[serde(default)]
    pub wireguard_ip: Option<String>,
    /// Whether worker is running inside a Docker container
    #[serde(default)]
    pub is_docker: Option<bool>,
    /// Node that owns this worker (its broker's node_name), for peer workers.
    /// None for a locally-registered worker (it belongs to this node).
    #[serde(default)]
    pub source_node: Option<String>,
    /// True when the operator declared this address as one of THIS broker's own
    /// workers via ZAKURO_WORKERS. Needed when the worker is not reachable on
    /// loopback -- e.g. a broker in a WireGuard sidecar's network namespace
    /// reaching its own worker through the container gateway (172.17.0.1).
    #[serde(default)]
    pub explicit_local: bool,
    /// Node fingerprint this worker was registered under, for key-derived URIs.
    #[serde(default)]
    pub node_fp: String,
    /// Stable discriminator within a node (e.g. listen port), for key-derived URIs.
    #[serde(default)]
    pub slot: String,
    /// Whether this worker serves a fixed model catalog or any model.
    #[serde(default)]
    pub provider_type: ProviderType,
    /// Model uuids this worker serves (or `["*"]` for a general provider).
    #[serde(default)]
    pub served_models: Vec<String>,
    /// Price per 1,000,000 tokens (model-inference billing), in credits.
    #[serde(default)]
    pub price_per_mtok: f64,
    /// This worker's own price (from `/info` or the default) from before the
    /// hub first overrode it. `sync-workers` keeps reporting this, and the hub
    /// stores it as `reported_price_per_hour`. None until the hub sets a price.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub reported_price_per_hour: Option<f64>,
}

impl Worker {
    /// Opaque, IP-free, key-derived handle for this worker: `zc://worker-{node_fp}-{slot}`.
    pub fn zc_uri(&self) -> String {
        format!("zc://worker-{}-{}", self.node_fp, self.slot)
    }

    /// Create a new worker
    pub fn new(id: String, name: String, uri: String, worker_type: String) -> Self {
        Self {
            id,
            name,
            uri,
            worker_type,
            status: WorkerStatus::Healthy,
            resources: WorkerResources::default(),
            pricing: WorkerPricing::default(),
            last_heartbeat: Utc::now(),
            active_requests: 0,
            total_requests: 0,
            avg_latency_ms: 0.0,
            tags: Vec::new(),
            max_timeout_secs: 0.0, // 0 = unlimited
            hardware: HardwareInfo::default(),
            wireguard_ip: None,
            is_docker: None,
            source_node: None,
            explicit_local: false,
            node_fp: String::new(),
            slot: String::new(),
            provider_type: ProviderType::default(),
            served_models: Vec::new(),
            price_per_mtok: 0.0,
            reported_price_per_hour: None,
        }
    }

    /// Check if worker can handle the requested resources
    pub fn can_handle(&self, cpus: f64, memory_bytes: u64, gpus: u32) -> bool {
        self.status == WorkerStatus::Healthy
            && self.resources.cpus_available >= cpus
            && self.resources.memory_available >= memory_bytes
            && self.resources.gpus_available >= gpus
    }

    /// Update heartbeat timestamp
    pub fn heartbeat(&mut self) {
        self.last_heartbeat = Utc::now();
    }

    /// Check if worker is stale (no heartbeat within timeout)
    pub fn is_stale(&self, timeout_secs: i64) -> bool {
        let elapsed = Utc::now().signed_duration_since(self.last_heartbeat);
        elapsed.num_seconds() > timeout_secs
    }
}

/// Worker registration request
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkerRegistration {
    pub name: String,
    pub uri: String,
    pub worker_type: String,
    #[serde(default)]
    pub resources: WorkerResources,
    #[serde(default)]
    pub pricing: WorkerPricing,
    #[serde(default)]
    pub tags: Vec<String>,
    /// Maximum timeout this worker allows per request (seconds), 0 = unlimited
    #[serde(default)]
    pub max_timeout_secs: f64,
    /// Hardware details from the worker
    #[serde(default)]
    pub hardware: HardwareInfo,
    /// WireGuard IP (optional, extracted from URI if not provided)
    #[serde(default)]
    pub wireguard_ip: Option<String>,
    /// Whether worker is running inside a Docker container
    #[serde(default)]
    pub is_docker: Option<bool>,
    /// Owning node (for peer workers); None for a local registration.
    #[serde(default)]
    pub source_node: Option<String>,
    /// Declared as this broker's own worker via ZAKURO_WORKERS (see Worker).
    #[serde(default)]
    pub explicit_local: bool,
    /// Whether this worker serves a fixed model catalog or any model.
    #[serde(default)]
    pub provider_type: ProviderType,
    /// Model uuids this worker serves (or `["*"]` for a general provider).
    #[serde(default)]
    pub served_models: Vec<String>,
    /// Price per 1,000,000 tokens (model-inference billing), in credits.
    #[serde(default)]
    pub price_per_mtok: f64,
}

/// Build a `WorkerRegistration` with neutral defaults for tests, varying only
/// name, uri and price. Shared across broker test modules so each doesn't
/// hand-roll its own literal.
#[cfg(test)]
pub(crate) fn test_registration(name: &str, uri: String, price: f64) -> WorkerRegistration {
    WorkerRegistration {
        name: name.to_string(),
        uri,
        worker_type: "zakuro".to_string(),
        resources: WorkerResources::default(),
        pricing: WorkerPricing {
            price_per_hour: price,
            min_charge: 0.001,
        },
        tags: vec![],
        max_timeout_secs: 0.0,
        hardware: HardwareInfo::default(),
        wireguard_ip: None,
        is_docker: None,
        source_node: None,
        explicit_local: false,
        provider_type: ProviderType::Specialized,
        served_models: vec![],
        price_per_mtok: 0.0,
    }
}

/// Worker heartbeat request
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkerHeartbeat {
    pub worker_id: String,
    #[serde(default)]
    pub resources: Option<WorkerResources>,
    #[serde(default)]
    pub active_requests: Option<u32>,
    #[serde(default)]
    pub status: Option<WorkerStatus>,
    #[serde(default)]
    pub max_timeout_secs: Option<f64>,
}

/// Thread-safe worker registry
#[derive(Debug)]
pub struct WorkerRegistry {
    workers: DashMap<String, Worker>,
    /// Per-worker request timestamps for time-windowed counts.
    /// Each entry is a ring-buffer of UTC timestamps; entries older than
    /// 30 days are pruned on every write.
    request_history: DashMap<String, Mutex<VecDeque<DateTime<Utc>>>>,
    /// Per-worker quotas (max requests per time window)
    quotas: DashMap<String, WorkerQuotas>,
    /// Model uuid -> worker ids of SPECIALIZED workers serving that model.
    model_index: DashMap<String, HashSet<String>>,
    /// Worker ids of GENERAL providers (serve any model uuid).
    general_workers: DashMap<String, ()>,
    /// Hub-set effective prices, by worker NAME (stable across re-registration,
    /// unlike the per-registration `id`), so a worker's own restart/re-probe
    /// doesn't forget the hub's price for it (spec §6.6). Applied both when
    /// `apply_hub_prices` runs and when `register` builds a fresh `Worker` for
    /// a name this map already knows about.
    hub_prices: DashMap<String, f64>,
}

impl WorkerRegistry {
    /// Create a new empty registry
    pub fn new() -> Self {
        Self {
            workers: DashMap::new(),
            request_history: DashMap::new(),
            quotas: DashMap::new(),
            model_index: DashMap::new(),
            general_workers: DashMap::new(),
            hub_prices: DashMap::new(),
        }
    }

    /// Insert `worker` into the model index according to its `provider_type`/`served_models`.
    fn index_worker(&self, worker: &Worker) {
        match worker.provider_type {
            ProviderType::General => {
                self.general_workers.insert(worker.id.clone(), ());
            }
            ProviderType::Specialized => {
                for model_uuid in &worker.served_models {
                    if model_uuid == MODEL_WILDCARD {
                        // Specialized worker with a wildcard entry behaves like a general one.
                        self.general_workers.insert(worker.id.clone(), ());
                        continue;
                    }
                    self.model_index
                        .entry(model_uuid.clone())
                        .or_default()
                        .insert(worker.id.clone());
                }
            }
        }
    }

    /// Remove `worker_id` from every model-index entry and the general-worker set.
    fn deindex_worker(&self, worker_id: &str) {
        self.general_workers.remove(worker_id);
        self.model_index.retain(|_, ids| {
            ids.remove(worker_id);
            !ids.is_empty()
        });
    }

    /// Workers that can serve `model_uuid`: SPECIALIZED workers whose `served_models`
    /// contains it, UNION all GENERAL providers. Resolution/ordering among the
    /// candidates is left to the caller.
    pub fn workers_serving(&self, model_uuid: &str) -> Vec<Worker> {
        let mut ids: HashSet<String> = self
            .model_index
            .get(model_uuid)
            .map(|entry| entry.clone())
            .unwrap_or_default();
        for entry in self.general_workers.iter() {
            ids.insert(entry.key().clone());
        }
        ids.into_iter().filter_map(|id| self.get(&id)).collect()
    }

    /// Register a new worker
    pub fn register(&self, registration: WorkerRegistration) -> Worker {
        let id = uuid::Uuid::new_v4().to_string();

        // Extract IP from URI if wireguard_ip not explicitly provided.
        // If the extracted IP is loopback (127.x or ::1), fall back to
        // the auto-detected WireGuard IP so Docker nodes show their real IP.
        let wireguard_ip = registration.wireguard_ip.clone().or_else(|| {
            let uri_ip = registration
                .uri
                .strip_prefix("http://")
                .or_else(|| registration.uri.strip_prefix("https://"))
                .and_then(|rest| {
                    let host = rest.split('/').next().unwrap_or(rest);
                    let ip = host.split(':').next().unwrap_or(host);
                    if ip.is_empty() {
                        None
                    } else {
                        Some(ip.to_string())
                    }
                });
            match uri_ip.as_deref() {
                Some("127.0.0.1") | Some("::1") | Some("localhost") => {
                    super::discovery::get_mesh_ip().or(uri_ip)
                }
                _ => uri_ip,
            }
        });

        let mut worker = Worker::new(
            id.clone(),
            registration.name,
            registration.uri,
            registration.worker_type,
        );
        worker.resources = registration.resources;
        worker.pricing = registration.pricing;
        worker.tags = registration.tags;
        worker.max_timeout_secs = registration.max_timeout_secs;
        worker.hardware = registration.hardware;
        worker.wireguard_ip = wireguard_ip;
        // Auto-detect is_docker from broker environment if not explicitly provided by worker:
        // if /.dockerenv exists the broker itself is running in Docker, so its workers are too.
        worker.is_docker = registration
            .is_docker
            .or_else(|| Some(std::path::Path::new("/.dockerenv").exists()));
        worker.source_node = registration.source_node;
        worker.explicit_local = registration.explicit_local;
        worker.provider_type = registration.provider_type;
        worker.served_models = registration.served_models;
        worker.price_per_mtok = registration.price_per_mtok;
        // Stable discriminator within a node for key-derived worker uris:
        // the worker's listen port, extracted from its uri.
        worker.slot = worker
            .uri
            .strip_prefix("http://")
            .or_else(|| worker.uri.strip_prefix("https://"))
            .and_then(|rest| rest.split('/').next())
            .and_then(|host| host.rsplit(':').next())
            .unwrap_or_default()
            .to_string();

        // A hub price previously applied for THIS name survives re-registration
        // (a worker restart re-probes and rebuilds `pricing` from scratch).
        // Peer workers are never repriced by this broker's own hub_prices map.
        if worker.source_node.is_none() {
            if let Some(hub_price) = self.hub_prices.get(&worker.name).map(|p| *p) {
                worker.reported_price_per_hour = Some(worker.pricing.price_per_hour);
                worker.pricing.price_per_hour = hub_price;
            }
        }

        self.workers.insert(id.clone(), worker.clone());
        // Initialise time-series and quota buckets for the new worker
        self.request_history
            .insert(id.clone(), Mutex::new(VecDeque::new()));
        self.quotas.insert(id.clone(), WorkerQuotas::default());
        self.index_worker(&worker);
        worker
    }

    /// Set the owning node's fingerprint on a registered worker (key-derived uri
    /// identity). Called at local worker registration with this broker's own
    /// `NodeKey::fingerprint()`.
    pub fn set_node_fp(&self, worker_id: &str, node_fp: &str) -> Option<Worker> {
        self.workers.get_mut(worker_id).map(|mut w| {
            w.node_fp = node_fp.to_string();
            w.clone()
        })
    }

    /// Return the number of requests completed by a worker in the last `window` duration.
    pub fn requests_in_window(&self, worker_id: &str, window: Duration) -> u64 {
        let cutoff = Utc::now() - window;
        self.request_history
            .get(worker_id)
            .map(|entry| {
                let ring = entry.lock().unwrap();
                ring.iter().filter(|&&ts| ts >= cutoff).count() as u64
            })
            .unwrap_or(0)
    }

    /// Return the quota limits configured for a worker.
    pub fn get_quotas(&self, worker_id: &str) -> WorkerQuotas {
        self.quotas
            .get(worker_id)
            .map(|q| q.clone())
            .unwrap_or_default()
    }

    /// Update worker from heartbeat
    pub fn heartbeat(&self, heartbeat: WorkerHeartbeat) -> Option<Worker> {
        self.workers.get_mut(&heartbeat.worker_id).map(|mut w| {
            w.heartbeat();
            if let Some(resources) = heartbeat.resources {
                w.resources = resources;
            }
            if let Some(active) = heartbeat.active_requests {
                w.active_requests = active;
            }
            if let Some(status) = heartbeat.status {
                // A worker's own heartbeat must never undo a drain its owner
                // asked for: that would route new jobs to a worker being stopped.
                if w.status != WorkerStatus::Draining {
                    w.status = status;
                }
            }
            if let Some(max_timeout) = heartbeat.max_timeout_secs {
                w.max_timeout_secs = max_timeout;
            }
            w.clone()
        })
    }

    /// Refresh heartbeat for a worker by ID (used by discovery to keep workers alive)
    pub fn refresh_heartbeat(&self, id: &str) {
        if let Some(mut w) = self.workers.get_mut(id) {
            w.heartbeat();
            // Also mark as healthy if it was marked unhealthy due to timeout
            if w.status == WorkerStatus::Unhealthy {
                w.status = WorkerStatus::Healthy;
            }
        }
    }

    /// Update live resources and storage for an existing worker (called on each discovery scan).
    /// Refreshes heartbeat, updates all resource fields, and updates storage_gb (dynamic value).
    /// Static hardware fields (cpu_model, gpu_model, gpu_vram_gb) are left unchanged.
    pub fn update_resources(&self, id: &str, resources: WorkerResources, hardware: HardwareInfo) {
        if let Some(mut w) = self.workers.get_mut(id) {
            w.heartbeat();
            w.resources = resources;
            // storage_gb is dynamic (free disk space); update it on each probe
            if hardware.storage_gb.is_some() {
                w.hardware.storage_gb = hardware.storage_gb;
            }
            if w.status == WorkerStatus::Unhealthy {
                w.status = WorkerStatus::Healthy;
            }
        }
    }

    /// Get a worker by ID
    pub fn get(&self, id: &str) -> Option<Worker> {
        self.workers.get(id).map(|w| w.clone())
    }

    /// Remove a worker
    pub fn remove(&self, id: &str) -> Option<Worker> {
        self.deindex_worker(id);
        self.workers.remove(id).map(|(_, w)| w)
    }

    /// List all workers
    pub fn list(&self) -> Vec<Worker> {
        self.workers.iter().map(|w| w.clone()).collect()
    }

    /// Aggregate free-resource snapshot summed over all currently healthy workers.
    /// Uses saturating adds so pathological worker resource values can't overflow.
    pub fn aggregate_available(&self) -> BrokerResources {
        self.healthy()
            .into_iter()
            .fold(BrokerResources::default(), |mut acc, w| {
                let cpus = w.resources.cpus_available.max(0.0) as u64;
                acc.cpus_available = acc.cpus_available.saturating_add(cpus);
                acc.memory_available = acc
                    .memory_available
                    .saturating_add(w.resources.memory_available);
                acc.gpus_available = acc
                    .gpus_available
                    .saturating_add(w.resources.gpus_available);
                acc.healthy_workers = acc.healthy_workers.saturating_add(1);
                acc
            })
    }

    /// List healthy workers
    pub fn healthy(&self) -> Vec<Worker> {
        self.workers
            .iter()
            .filter(|w| w.status == WorkerStatus::Healthy)
            .map(|w| w.clone())
            .collect()
    }

    /// IDs of healthy workers — a cheap alternative to `healthy()` for the
    /// routing hot path, which only needs to clone the *selected* worker, not
    /// every candidate. Avoids N full `Worker` deep-clones per request.
    pub fn healthy_ids(&self) -> Vec<String> {
        self.workers
            .iter()
            .filter(|w| w.status == WorkerStatus::Healthy)
            .map(|w| w.key().clone())
            .collect()
    }

    /// Find workers that can handle the request
    pub fn find_capable(&self, cpus: f64, memory_bytes: u64, gpus: u32) -> Vec<Worker> {
        self.workers
            .iter()
            .filter(|w| w.can_handle(cpus, memory_bytes, gpus))
            .map(|w| w.clone())
            .collect()
    }

    /// Mark stale workers as unhealthy
    pub fn mark_stale(&self, timeout_secs: i64) {
        for mut entry in self.workers.iter_mut() {
            if entry.is_stale(timeout_secs) && entry.status == WorkerStatus::Healthy {
                entry.status = WorkerStatus::Unhealthy;
            }
        }
    }

    /// Remove workers that have been stale (unhealthy) for longer than `timeout_secs`.
    /// Returns the IDs of removed workers.
    pub fn remove_stale(&self, timeout_secs: i64) -> Vec<String> {
        let to_remove: Vec<String> = self
            .workers
            .iter()
            .filter(|w| w.status == WorkerStatus::Unhealthy && w.is_stale(timeout_secs))
            .map(|w| w.id.clone())
            .collect();
        for id in &to_remove {
            self.deindex_worker(id);
            self.workers.remove(id);
            self.request_history.remove(id);
            self.quotas.remove(id);
        }
        to_remove
    }

    /// Atomically reserve one quota slot for the given worker.
    ///
    /// Checks all configured quota windows under the ring-buffer mutex, then
    /// pre-commits by pushing the current timestamp if every window has room.
    /// Returns `true` if the slot was reserved; `false` if any quota is full.
    ///
    /// On success the caller MUST eventually call either `record_request`
    /// (keeps the pre-committed slot) or `cancel_quota_reservation` (removes it).
    pub fn try_reserve_quota(&self, worker_id: &str) -> bool {
        let quotas = self.get_quotas(worker_id);
        let all_unlimited = quotas.per_5h == 0 && quotas.per_week == 0 && quotas.per_month == 0;

        match self.request_history.get(worker_id) {
            None => true, // new worker with no history — allow
            Some(entry) => {
                let mut ring = entry.lock().unwrap();
                let now = Utc::now();

                if !all_unlimited {
                    if quotas.per_5h > 0 {
                        let cutoff = now - Duration::hours(5);
                        let count = ring.iter().filter(|&&ts| ts >= cutoff).count() as u64;
                        if count >= quotas.per_5h {
                            return false;
                        }
                    }
                    if quotas.per_week > 0 {
                        let cutoff = now - Duration::weeks(1);
                        let count = ring.iter().filter(|&&ts| ts >= cutoff).count() as u64;
                        if count >= quotas.per_week {
                            return false;
                        }
                    }
                    if quotas.per_month > 0 {
                        let cutoff = now - Duration::days(30);
                        let count = ring.iter().filter(|&&ts| ts >= cutoff).count() as u64;
                        if count >= quotas.per_month {
                            return false;
                        }
                    }
                }

                // All checks passed — pre-commit the slot
                ring.push_back(now);
                // Prune entries older than 30 days while the mutex is held
                let prune_cutoff = now - Duration::days(30);
                while ring.front().map(|&ts| ts < prune_cutoff).unwrap_or(false) {
                    ring.pop_front();
                }
                true
            }
        }
    }

    /// Cancel a previously reserved quota slot (call on request failure).
    /// Removes the last pre-committed timestamp from the worker's ring buffer.
    pub fn cancel_quota_reservation(&self, worker_id: &str) {
        if let Some(entry) = self.request_history.get(worker_id) {
            let mut ring = entry.lock().unwrap();
            ring.pop_back();
        }
    }

    /// Update worker stats after request completion.
    ///
    /// NOTE: does NOT push a timestamp to the ring buffer — that was already
    /// done atomically by `try_reserve_quota` at dispatch time.
    pub fn record_request(&self, worker_id: &str, duration_ms: f64, _success: bool) {
        if let Some(mut worker) = self.workers.get_mut(worker_id) {
            worker.total_requests += 1;
            // Exponential moving average for latency
            let alpha = 0.1;
            worker.avg_latency_ms = alpha * duration_ms + (1.0 - alpha) * worker.avg_latency_ms;
            if worker.active_requests > 0 {
                worker.active_requests -= 1;
            }
        }
    }

    /// Increment active request count for a worker
    pub fn increment_active(&self, worker_id: &str) {
        if let Some(mut worker) = self.workers.get_mut(worker_id) {
            worker.active_requests += 1;
        }
    }

    /// Decrement a worker's in-flight counter, flooring at 0 (idempotent).
    pub fn decrement_active(&self, worker_id: &str) {
        if let Some(mut worker) = self.workers.get_mut(worker_id) {
            if worker.active_requests > 0 {
                worker.active_requests -= 1;
            }
        }
    }

    /// RAII in-flight counter: increments now, and decrements on drop UNLESS
    /// disarmed. The execute path arms one of these at dispatch so every
    /// early-return/error/timeout path releases the slot; the success path
    /// (where `record_request` already decremented) calls `disarm()` to avoid a
    /// double-decrement. Fixes the active_requests leak that biased
    /// best_availability routing against workers that ever failed (audit M2).
    pub fn active_guard(&self, worker_id: &str) -> ActiveGuard<'_> {
        self.increment_active(worker_id);
        ActiveGuard {
            registry: self,
            worker_id: worker_id.to_string(),
            armed: true,
        }
    }

    /// Get worker count
    pub fn count(&self) -> usize {
        self.workers.len()
    }

    /// Mark a worker as unhealthy. Leaves a `Draining` worker untouched:
    /// otherwise a failed forward to a draining worker would flip it to
    /// `Unhealthy`, and the next discovery scan would lift it back to
    /// `Healthy`, putting a drained worker back in rotation.
    pub fn mark_unhealthy(&self, worker_id: &str) {
        if let Some(mut worker) = self.workers.get_mut(worker_id) {
            if worker.status != WorkerStatus::Draining {
                worker.status = WorkerStatus::Unhealthy;
            }
        }
    }

    /// Stop routing new work to `id` (status `Draining`) so its owner can stop
    /// it without killing a job. Selection only ever picks `Healthy` workers,
    /// and neither discovery refreshes nor heartbeats lift `Draining`.
    /// Returns the updated worker, or None when `id` is unknown.
    pub fn drain(&self, id: &str) -> Option<Worker> {
        self.workers.get_mut(id).map(|mut w| {
            w.status = WorkerStatus::Draining;
            w.clone()
        })
    }

    /// Adopt the hub's effective prices (spec 2026-09-13-macos-widget §6.6) for
    /// THIS broker's own workers, matched by name (the `worker_id` sync sends).
    /// The first override remembers the worker's own price in
    /// `reported_price_per_hour`. Peer workers (`source_node` set) are never
    /// touched. Returns how many workers changed price.
    ///
    /// Every name in `prices` is also remembered in `hub_prices` (regardless of
    /// whether a live worker with that name exists right now), so a worker that
    /// re-registers later -- after a restart, or a fresh discovery probe -- has
    /// the hub's price reapplied by `register` immediately, with no sync round
    /// trip in between (fix round 1, Finding 1).
    pub fn apply_hub_prices(&self, prices: &std::collections::HashMap<String, f64>) -> usize {
        for (name, &price) in prices.iter() {
            self.hub_prices.insert(name.clone(), price);
        }
        let mut changed = 0;
        for mut w in self.workers.iter_mut() {
            if w.source_node.is_some() {
                continue;
            }
            if let Some(&price) = prices.get(&w.name) {
                if w.reported_price_per_hour.is_none() {
                    w.reported_price_per_hour = Some(w.pricing.price_per_hour);
                }
                if w.pricing.price_per_hour != price {
                    w.pricing.price_per_hour = price;
                    changed += 1;
                }
            }
        }
        changed
    }

    /// Resolve the best worker to serve `model_uuid`.
    ///
    /// Candidates come from [`workers_serving`], filtered to the same
    /// HEALTHY criterion the registry already applies elsewhere (`status ==
    /// WorkerStatus::Healthy`, as used by [`healthy`]/[`healthy_ids`]).
    /// Specialized providers pinned to this uuid are preferred over general
    /// providers; within the chosen tier the cheapest `price_per_mtok` wins,
    /// ties broken deterministically by worker `name`.
    pub fn resolve_model(&self, model_uuid: &str) -> Result<Worker, ModelResolveError> {
        let candidates: Vec<Worker> = self
            .workers_serving(model_uuid)
            .into_iter()
            .filter(|w| w.status == WorkerStatus::Healthy)
            .collect();

        let (specialized, general): (Vec<Worker>, Vec<Worker>) = candidates
            .into_iter()
            .partition(|w| w.provider_type == ProviderType::Specialized);

        let pick_cheapest = |mut tier: Vec<Worker>| -> Option<Worker> {
            tier.sort_by(|a, b| {
                a.price_per_mtok
                    .partial_cmp(&b.price_per_mtok)
                    .unwrap_or(std::cmp::Ordering::Equal)
                    .then_with(|| a.name.cmp(&b.name))
            });
            tier.into_iter().next()
        };

        if !specialized.is_empty() {
            return Ok(pick_cheapest(specialized).expect("non-empty specialized tier"));
        }
        if !general.is_empty() {
            return Ok(pick_cheapest(general).expect("non-empty general tier"));
        }
        Err(ModelResolveError::NoProvider {
            model_uuid: model_uuid.to_string(),
        })
    }
}

/// Errors from [`WorkerRegistry::resolve_model`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ModelResolveError {
    /// No healthy worker (specialized or general) serves this model uuid.
    NoProvider { model_uuid: String },
}

impl std::fmt::Display for ModelResolveError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::NoProvider { model_uuid } => {
                write!(f, "no provider serving zc://{}", model_uuid)
            }
        }
    }
}

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

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

/// RAII guard returned by [`WorkerRegistry::active_guard`]. Decrements the
/// worker's in-flight counter on drop unless [`ActiveGuard::disarm`] was called.
pub struct ActiveGuard<'a> {
    registry: &'a WorkerRegistry,
    worker_id: String,
    armed: bool,
}

impl<'a> ActiveGuard<'a> {
    /// Disarm so `Drop` will NOT decrement — call on the success path where
    /// `record_request` has already released the slot.
    pub fn disarm(&mut self) {
        self.armed = false;
    }
}

impl<'a> Drop for ActiveGuard<'a> {
    fn drop(&mut self) {
        if self.armed {
            self.registry.decrement_active(&self.worker_id);
        }
    }
}

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

    #[test]
    fn worker_zc_uri_and_source_node_default() {
        let mut w = Worker::new(
            "id1".into(),
            "worker-abc123".into(),
            "http://10.13.13.21:3960".into(),
            "zakuro".into(),
        );
        w.node_fp = "abc123".into();
        w.slot = "3960".into();
        assert_eq!(w.zc_uri(), "zc://worker-abc123-3960");
        assert!(w.source_node.is_none());
        w.source_node = Some("node-i9".into());
        assert_eq!(w.source_node.as_deref(), Some("node-i9"));
    }

    #[test]
    fn worker_zc_uri_is_key_derived_and_collision_free() {
        let mut w1 = Worker::new(
            "id1".into(),
            "worker-lxd".into(),
            "http://127.0.0.1:3960".into(),
            "zakuro".into(),
        );
        let mut w2 = Worker::new(
            "id2".into(),
            "worker-lxd".into(),
            "http://127.0.0.1:3960".into(),
            "zakuro".into(),
        );
        w1.node_fp = "aaaaaaaaaaaaaaaa".into();
        w1.slot = "3960".into();
        w2.node_fp = "bbbbbbbbbbbbbbbb".into();
        w2.slot = "3960".into();
        assert_eq!(w1.zc_uri(), "zc://worker-aaaaaaaaaaaaaaaa-3960");
        assert_ne!(w1.zc_uri(), w2.zc_uri()); // same name "worker-lxd", different nodes -> distinct
    }

    fn registration(name: &str, uri: &str) -> WorkerRegistration {
        WorkerRegistration {
            name: name.to_string(),
            uri: uri.to_string(),
            worker_type: "zakuro".to_string(),
            resources: WorkerResources::default(),
            pricing: WorkerPricing::default(),
            tags: vec![],
            max_timeout_secs: 0.0,
            hardware: HardwareInfo::default(),
            wireguard_ip: None,
            is_docker: None,
            source_node: None,
            explicit_local: false,
            provider_type: ProviderType::default(),
            served_models: vec![],
            price_per_mtok: 0.0,
        }
    }

    /// Registration helper with explicit provider_type/served_models, for model-index tests.
    fn model_registration(
        name: &str,
        uri: &str,
        provider_type: ProviderType,
        served_models: Vec<&str>,
    ) -> WorkerRegistration {
        let mut r = registration(name, uri);
        r.provider_type = provider_type;
        r.served_models = served_models.into_iter().map(String::from).collect();
        r
    }

    // --- Model index / workers_serving ---

    #[test]
    fn test_workers_serving_specialized_matches_only_its_model() {
        let reg = WorkerRegistry::new();
        let w = reg.register(model_registration(
            "specialist",
            "http://127.0.0.1:4001",
            ProviderType::Specialized,
            vec!["uuid-a"],
        ));

        let matches = reg.workers_serving("uuid-a");
        assert_eq!(matches.len(), 1);
        assert_eq!(matches[0].id, w.id);

        let no_matches = reg.workers_serving("uuid-b");
        assert!(no_matches.is_empty());
    }

    #[test]
    fn test_workers_serving_general_matches_any_model() {
        let reg = WorkerRegistry::new();
        let w = reg.register(model_registration(
            "generalist",
            "http://127.0.0.1:4002",
            ProviderType::General,
            vec!["*"],
        ));

        for uuid in ["uuid-a", "uuid-b", "any-random-uuid"] {
            let matches = reg.workers_serving(uuid);
            assert_eq!(matches.len(), 1, "general worker should serve {uuid}");
            assert_eq!(matches[0].id, w.id);
        }
    }

    #[test]
    fn test_workers_serving_unions_specialized_and_general() {
        let reg = WorkerRegistry::new();
        let specialist = reg.register(model_registration(
            "specialist",
            "http://127.0.0.1:4003",
            ProviderType::Specialized,
            vec!["uuid-a"],
        ));
        let generalist = reg.register(model_registration(
            "generalist",
            "http://127.0.0.1:4004",
            ProviderType::General,
            vec![],
        ));

        let mut ids: Vec<String> = reg
            .workers_serving("uuid-a")
            .into_iter()
            .map(|w| w.id)
            .collect();
        ids.sort();
        let mut expected = vec![specialist.id.clone(), generalist.id.clone()];
        expected.sort();
        assert_eq!(ids, expected);
    }

    #[test]
    fn test_reregister_updates_model_index() {
        let reg = WorkerRegistry::new();
        let w1 = reg.register(model_registration(
            "flip",
            "http://127.0.0.1:4005",
            ProviderType::Specialized,
            vec!["uuid-old"],
        ));
        assert_eq!(reg.workers_serving("uuid-old").len(), 1);

        // Simulate re-registration: remove the old registration then register
        // the same worker (by name) with a different served_models set.
        reg.remove(&w1.id);
        assert!(
            reg.workers_serving("uuid-old").is_empty(),
            "removed worker must drop out of the old model's index entry"
        );

        let w2 = reg.register(model_registration(
            "flip",
            "http://127.0.0.1:4005",
            ProviderType::Specialized,
            vec!["uuid-new"],
        ));
        assert!(reg.workers_serving("uuid-old").is_empty());
        let matches = reg.workers_serving("uuid-new");
        assert_eq!(matches.len(), 1);
        assert_eq!(matches[0].id, w2.id);
    }

    #[test]
    fn test_remove_drops_worker_from_general_index() {
        let reg = WorkerRegistry::new();
        let w = reg.register(model_registration(
            "generalist",
            "http://127.0.0.1:4006",
            ProviderType::General,
            vec![],
        ));
        assert_eq!(reg.workers_serving("any-uuid").len(), 1);
        reg.remove(&w.id);
        assert!(reg.workers_serving("any-uuid").is_empty());
    }

    #[test]
    fn test_legacy_worker_json_without_new_fields_defaults_back_compat() {
        // Legacy payload predating provider_type/served_models/price_per_mtok.
        let legacy_json = r#"{
            "id": "legacy-id",
            "name": "legacy-worker",
            "uri": "http://127.0.0.1:3960",
            "worker_type": "zakuro",
            "status": "healthy",
            "resources": {
                "cpus_total": 1.0,
                "cpus_available": 1.0,
                "memory_total": 1073741824,
                "memory_available": 1073741824,
                "gpus_total": 0,
                "gpus_available": 0
            },
            "pricing": {
                "price_per_hour": 3.6,
                "min_charge": 0.001
            },
            "last_heartbeat": "2026-01-01T00:00:00Z",
            "active_requests": 0,
            "total_requests": 0,
            "avg_latency_ms": 0.0,
            "tags": []
        }"#;

        let w: Worker = serde_json::from_str(legacy_json).expect("legacy JSON must deserialize");
        assert_eq!(w.provider_type, ProviderType::Specialized);
        assert!(w.served_models.is_empty());
        assert_eq!(w.price_per_mtok, 0.0);
    }

    #[test]
    fn test_legacy_worker_registration_json_without_new_fields_defaults_back_compat() {
        let legacy_json = r#"{
            "name": "legacy-worker",
            "uri": "http://127.0.0.1:3960",
            "worker_type": "zakuro"
        }"#;

        let r: WorkerRegistration =
            serde_json::from_str(legacy_json).expect("legacy JSON must deserialize");
        assert_eq!(r.provider_type, ProviderType::Specialized);
        assert!(r.served_models.is_empty());
        assert_eq!(r.price_per_mtok, 0.0);
    }

    // --- WorkerRegistry ---

    #[test]
    fn test_register_assigns_unique_id() {
        let reg = WorkerRegistry::new();
        let w1 = reg.register(registration("w1", "http://127.0.0.1:3960"));
        let w2 = reg.register(registration("w2", "http://127.0.0.1:3961"));
        assert!(!w1.id.is_empty());
        assert_ne!(w1.id, w2.id);
        assert_eq!(reg.count(), 2);
    }

    #[test]
    fn test_register_stores_fields() {
        let reg = WorkerRegistry::new();
        let worker = reg.register(registration("my-worker", "http://10.0.0.1:3960"));
        let got = reg.get(&worker.id).unwrap();
        assert_eq!(got.name, "my-worker");
        assert_eq!(got.uri, "http://10.0.0.1:3960");
        assert_eq!(got.status, WorkerStatus::Healthy);
        assert_eq!(got.active_requests, 0);
    }

    #[test]
    fn test_remove_worker() {
        let reg = WorkerRegistry::new();
        let w = reg.register(registration("w", "http://127.0.0.1:3960"));
        assert!(reg.get(&w.id).is_some());
        let removed = reg.remove(&w.id);
        assert!(removed.is_some());
        assert!(reg.get(&w.id).is_none());
        assert_eq!(reg.count(), 0);
    }

    #[test]
    fn test_heartbeat_updates_resources_and_status() {
        let reg = WorkerRegistry::new();
        let w = reg.register(registration("w", "http://127.0.0.1:3960"));

        let hb = WorkerHeartbeat {
            worker_id: w.id.clone(),
            resources: Some(WorkerResources {
                cpus_total: 8.0,
                cpus_available: 4.0,
                memory_total: 16 * 1024 * 1024 * 1024,
                memory_available: 8 * 1024 * 1024 * 1024,
                gpus_total: 2,
                gpus_available: 1,
            }),
            active_requests: Some(5),
            status: Some(WorkerStatus::Busy),
            max_timeout_secs: Some(120.0),
        };

        let updated = reg.heartbeat(hb).unwrap();
        assert_eq!(updated.resources.cpus_available, 4.0);
        assert_eq!(updated.resources.gpus_available, 1);
        assert_eq!(updated.active_requests, 5);
        assert_eq!(updated.status, WorkerStatus::Busy);
        assert_eq!(updated.max_timeout_secs, 120.0);
    }

    #[test]
    fn test_heartbeat_unknown_worker_returns_none() {
        let reg = WorkerRegistry::new();
        let hb = WorkerHeartbeat {
            worker_id: "nonexistent".to_string(),
            resources: None,
            active_requests: None,
            status: None,
            max_timeout_secs: None,
        };
        assert!(reg.heartbeat(hb).is_none());
    }

    #[test]
    fn test_mark_stale_sets_unhealthy() {
        let reg = WorkerRegistry::new();
        let w = reg.register(registration("w", "http://127.0.0.1:3960"));
        // timeout=-1 means every worker is immediately stale (elapsed >= 0 > -1)
        reg.mark_stale(-1);
        let got = reg.get(&w.id).unwrap();
        assert_eq!(got.status, WorkerStatus::Unhealthy);
    }

    #[test]
    fn test_refresh_heartbeat_restores_healthy() {
        let reg = WorkerRegistry::new();
        let w = reg.register(registration("w", "http://127.0.0.1:3960"));
        reg.mark_stale(-1);
        assert_eq!(reg.get(&w.id).unwrap().status, WorkerStatus::Unhealthy);
        reg.refresh_heartbeat(&w.id);
        assert_eq!(reg.get(&w.id).unwrap().status, WorkerStatus::Healthy);
    }

    #[test]
    fn aggregate_available_sums_healthy_workers() {
        let reg = WorkerRegistry::new();
        let mut r1 = registration("w1", "http://127.0.0.1:3960");
        r1.resources = WorkerResources {
            cpus_total: 4.0,
            cpus_available: 4.0,
            memory_total: 8_000_000_000,
            memory_available: 8_000_000_000,
            gpus_total: 1,
            gpus_available: 1,
        };
        let mut r2 = registration("w2", "http://127.0.0.1:3961");
        r2.resources = WorkerResources {
            cpus_total: 2.0,
            cpus_available: 2.0,
            memory_total: 4_000_000_000,
            memory_available: 4_000_000_000,
            gpus_total: 0,
            gpus_available: 0,
        };
        reg.register(r1);
        reg.register(r2);

        let agg = reg.aggregate_available();
        assert_eq!(agg.cpus_available, 6);
        assert_eq!(agg.memory_available, 12_000_000_000);
        assert_eq!(agg.gpus_available, 1);
        assert_eq!(agg.healthy_workers, 2);
    }

    #[test]
    fn aggregate_available_excludes_unhealthy_and_ignores_empty() {
        let reg = WorkerRegistry::new();
        let empty = reg.aggregate_available();
        assert_eq!(empty, BrokerResources::default());

        let w = reg.register(registration("sick", "http://127.0.0.1:3962"));
        reg.mark_unhealthy(&w.id);
        let agg = reg.aggregate_available();
        assert_eq!(agg.healthy_workers, 0);
        assert_eq!(agg.cpus_available, 0);
    }

    #[test]
    fn test_healthy_list_excludes_unhealthy() {
        let reg = WorkerRegistry::new();
        let w1 = reg.register(registration("healthy", "http://127.0.0.1:3960"));
        let w2 = reg.register(registration("sick", "http://127.0.0.1:3961"));
        reg.mark_unhealthy(&w2.id);

        let healthy = reg.healthy();
        assert_eq!(healthy.len(), 1);
        assert_eq!(healthy[0].id, w1.id);
    }

    #[test]
    fn test_increment_active_and_decrement_on_record() {
        let reg = WorkerRegistry::new();
        let w = reg.register(registration("w", "http://127.0.0.1:3960"));
        reg.increment_active(&w.id);
        reg.increment_active(&w.id);
        assert_eq!(reg.get(&w.id).unwrap().active_requests, 2);
        reg.record_request(&w.id, 100.0, true);
        assert_eq!(reg.get(&w.id).unwrap().active_requests, 1);
    }

    #[test]
    fn active_guard_decrements_on_drop_even_without_record() {
        let reg = WorkerRegistry::new();
        let w = reg.register(registration("w1", "http://127.0.0.1:3960"));
        {
            let _g = reg.active_guard(&w.id);
            assert_eq!(reg.get(&w.id).unwrap().active_requests, 1);
        } // guard dropped here (simulates an error path with no record_request)
        assert_eq!(
            reg.get(&w.id).unwrap().active_requests,
            0,
            "active_requests must return to 0 when the guard drops on an error path"
        );
    }

    #[test]
    fn active_guard_disarm_skips_decrement() {
        let reg = WorkerRegistry::new();
        let w = reg.register(registration("w2", "http://127.0.0.1:3960"));
        {
            let mut g = reg.active_guard(&w.id);
            assert_eq!(reg.get(&w.id).unwrap().active_requests, 1);
            // Success path: record_request decremented, so disarm the guard.
            reg.record_request(&w.id, 10.0, true);
            g.disarm();
        }
        assert_eq!(
            reg.get(&w.id).unwrap().active_requests,
            0,
            "disarmed guard must not double-decrement after record_request"
        );
    }

    #[test]
    fn test_record_request_latency_ema() {
        let reg = WorkerRegistry::new();
        let w = reg.register(registration("w", "http://127.0.0.1:3960"));
        // First recording: EMA = 0.1 * 200 + 0.9 * 0 = 20
        reg.record_request(&w.id, 200.0, true);
        let got = reg.get(&w.id).unwrap();
        assert!((got.avg_latency_ms - 20.0).abs() < 0.001);
        assert_eq!(got.total_requests, 1);
    }

    #[test]
    fn test_find_capable_filters_by_resources() {
        let reg = WorkerRegistry::new();
        let mut r_small = registration("small", "http://127.0.0.1:3960");
        r_small.resources = WorkerResources {
            cpus_total: 2.0,
            cpus_available: 2.0,
            memory_total: 2 * 1024 * 1024 * 1024,
            memory_available: 2 * 1024 * 1024 * 1024,
            gpus_total: 0,
            gpus_available: 0,
        };
        let mut r_large = registration("large", "http://127.0.0.1:3961");
        r_large.resources = WorkerResources {
            cpus_total: 32.0,
            cpus_available: 32.0,
            memory_total: 128 * 1024 * 1024 * 1024,
            memory_available: 128 * 1024 * 1024 * 1024,
            gpus_total: 4,
            gpus_available: 4,
        };

        reg.register(r_small);
        reg.register(r_large);

        // Only large can serve 16 CPUs
        let capable = reg.find_capable(16.0, 1024 * 1024 * 1024, 0);
        assert_eq!(capable.len(), 1);
        assert_eq!(capable[0].name, "large");

        // Both can serve 1 CPU
        assert_eq!(reg.find_capable(1.0, 512 * 1024 * 1024, 0).len(), 2);

        // None can serve 8 GPUs
        assert_eq!(reg.find_capable(1.0, 512 * 1024 * 1024, 8).len(), 0);
    }

    // --- Worker::can_handle ---

    #[test]
    fn test_can_handle_exact_match() {
        let w = Worker::new(
            "id".to_string(),
            "w".to_string(),
            "http://x".to_string(),
            "zakuro".to_string(),
        );
        // default: 1 CPU, 1 GiB, 0 GPU
        assert!(w.can_handle(1.0, 1024 * 1024 * 1024, 0));
    }

    #[test]
    fn test_can_handle_over_cpu_fails() {
        let w = Worker::new(
            "id".to_string(),
            "w".to_string(),
            "http://x".to_string(),
            "zakuro".to_string(),
        );
        assert!(!w.can_handle(2.0, 512 * 1024 * 1024, 0));
    }

    #[test]
    fn test_can_handle_over_memory_fails() {
        let w = Worker::new(
            "id".to_string(),
            "w".to_string(),
            "http://x".to_string(),
            "zakuro".to_string(),
        );
        assert!(!w.can_handle(0.5, 2 * 1024 * 1024 * 1024, 0));
    }

    #[test]
    fn test_can_handle_gpu_required_but_none_fails() {
        let w = Worker::new(
            "id".to_string(),
            "w".to_string(),
            "http://x".to_string(),
            "zakuro".to_string(),
        );
        assert!(!w.can_handle(0.5, 512 * 1024 * 1024, 1));
    }

    #[test]
    fn test_can_handle_unhealthy_worker_fails() {
        let mut w = Worker::new(
            "id".to_string(),
            "w".to_string(),
            "http://x".to_string(),
            "zakuro".to_string(),
        );
        w.status = WorkerStatus::Unhealthy;
        assert!(!w.can_handle(0.1, 1024, 0));
    }

    // --- WorkerPricing ---

    #[test]
    fn test_pricing_cost_formula() {
        let p = WorkerPricing {
            price_per_hour: 3.6,
            min_charge: 0.001,
        };
        // 10s at 3.6 credits/hour = 3.6/3600 * 10 = 0.010 > min_charge
        let cost = p.estimate_cost(10.0);
        assert!((cost - 0.010).abs() < 0.0001);
    }

    #[test]
    fn test_pricing_min_charge_enforced() {
        let p = WorkerPricing {
            price_per_hour: 0.0,
            min_charge: 0.005,
        };
        let cost = p.estimate_cost(0.001);
        assert_eq!(cost, 0.005);
    }

    // --- resolve_model ---

    fn priced_model_registration(
        name: &str,
        uri: &str,
        provider_type: ProviderType,
        served_models: Vec<&str>,
        price_per_mtok: f64,
    ) -> WorkerRegistration {
        let mut r = model_registration(name, uri, provider_type, served_models);
        r.price_per_mtok = price_per_mtok;
        r
    }

    #[test]
    fn test_resolve_model_prefers_specialized_over_general() {
        let reg = WorkerRegistry::new();
        let specialized = reg.register(priced_model_registration(
            "specialist",
            "http://127.0.0.1:5001",
            ProviderType::Specialized,
            vec!["uuid-a"],
            10.0,
        ));
        reg.register(priced_model_registration(
            "generalist",
            "http://127.0.0.1:5002",
            ProviderType::General,
            vec![],
            1.0, // cheaper, but general should lose to specialized
        ));

        let resolved = reg.resolve_model("uuid-a").unwrap();
        assert_eq!(resolved.id, specialized.id);
    }

    #[test]
    fn test_resolve_model_cheapest_among_specialized() {
        let reg = WorkerRegistry::new();
        reg.register(priced_model_registration(
            "expensive",
            "http://127.0.0.1:5003",
            ProviderType::Specialized,
            vec!["uuid-a"],
            20.0,
        ));
        let cheap = reg.register(priced_model_registration(
            "cheap",
            "http://127.0.0.1:5004",
            ProviderType::Specialized,
            vec!["uuid-a"],
            5.0,
        ));

        let resolved = reg.resolve_model("uuid-a").unwrap();
        assert_eq!(resolved.id, cheap.id);
    }

    #[test]
    fn test_resolve_model_general_fallback_cheapest() {
        let reg = WorkerRegistry::new();
        reg.register(priced_model_registration(
            "general-expensive",
            "http://127.0.0.1:5005",
            ProviderType::General,
            vec![],
            8.0,
        ));
        let cheap_general = reg.register(priced_model_registration(
            "general-cheap",
            "http://127.0.0.1:5006",
            ProviderType::General,
            vec![],
            2.0,
        ));

        let resolved = reg.resolve_model("uuid-z").unwrap();
        assert_eq!(resolved.id, cheap_general.id);
    }

    #[test]
    fn test_resolve_model_excludes_unhealthy_specialized_falls_back_to_general() {
        let reg = WorkerRegistry::new();
        let stale_specialist = reg.register(priced_model_registration(
            "stale-specialist",
            "http://127.0.0.1:5007",
            ProviderType::Specialized,
            vec!["uuid-a"],
            1.0, // cheapest, but stale/unhealthy so must be excluded
        ));
        reg.mark_unhealthy(&stale_specialist.id);

        let healthy_general = reg.register(priced_model_registration(
            "healthy-general",
            "http://127.0.0.1:5008",
            ProviderType::General,
            vec![],
            5.0,
        ));

        let resolved = reg.resolve_model("uuid-a").unwrap();
        assert_eq!(resolved.id, healthy_general.id);
    }

    #[test]
    fn test_resolve_model_no_provider_when_no_healthy_candidate() {
        let reg = WorkerRegistry::new();
        let w = reg.register(priced_model_registration(
            "only-specialist",
            "http://127.0.0.1:5009",
            ProviderType::Specialized,
            vec!["uuid-a"],
            1.0,
        ));
        reg.mark_unhealthy(&w.id);

        let err = reg.resolve_model("uuid-a").unwrap_err();
        assert_eq!(
            err,
            ModelResolveError::NoProvider {
                model_uuid: "uuid-a".to_string()
            }
        );
        assert_eq!(err.to_string(), "no provider serving zc://uuid-a");
    }

    #[test]
    fn test_pricing_price_score_weighted() {
        // 10s at 3.6 credits/hr = 0.010 > min_charge (0.001) → score = 0.010
        let p = WorkerPricing::default();
        let score = p.price_score();
        assert!((score - 0.010).abs() < 0.0001);
    }
}

/// Drain and hub-price behaviour the zc agent relies on
/// (spec 2026-09-13-macos-widget §6.5-6.6).
#[cfg(test)]
mod drain_and_price_tests {
    use super::*;

    #[test]
    fn drain_marks_draining_and_leaves_the_routable_set() {
        let r = WorkerRegistry::new();
        let w = r.register(test_registration(
            "fp-w0",
            "http://127.0.0.1:3960".to_string(),
            3.6,
        ));
        assert_eq!(r.drain(&w.id).unwrap().status, WorkerStatus::Draining);
        assert!(
            r.healthy_ids().is_empty(),
            "a draining worker is never selected"
        );
        assert!(r.healthy().is_empty());
    }

    #[test]
    fn discovery_refresh_does_not_undo_a_drain() {
        let r = WorkerRegistry::new();
        let w = r.register(test_registration(
            "fp-w0",
            "http://127.0.0.1:3960".to_string(),
            3.6,
        ));
        r.drain(&w.id);
        r.refresh_heartbeat(&w.id);
        r.update_resources(&w.id, WorkerResources::default(), HardwareInfo::default());
        assert_eq!(r.get(&w.id).unwrap().status, WorkerStatus::Draining);
    }

    #[test]
    fn a_worker_heartbeat_does_not_undo_a_drain() {
        let r = WorkerRegistry::new();
        let w = r.register(test_registration(
            "fp-w0",
            "http://127.0.0.1:3960".to_string(),
            3.6,
        ));
        r.drain(&w.id);
        r.heartbeat(WorkerHeartbeat {
            worker_id: w.id.clone(),
            resources: None,
            active_requests: Some(1),
            status: Some(WorkerStatus::Healthy),
            max_timeout_secs: None,
        });
        let got = r.get(&w.id).unwrap();
        assert_eq!(got.status, WorkerStatus::Draining);
        assert_eq!(
            got.active_requests, 1,
            "the rest of the heartbeat still applies"
        );
    }

    #[test]
    fn draining_an_unknown_worker_is_none() {
        assert!(WorkerRegistry::new().drain("nope").is_none());
    }

    /// A failed forward must not undo a drain — otherwise the next
    /// discovery scan lifts a merely-Unhealthy worker back to Healthy and a
    /// drained-then-stopping worker is put back in rotation.
    #[test]
    fn mark_unhealthy_does_not_undo_a_drain() {
        let r = WorkerRegistry::new();
        let w = r.register(test_registration(
            "fp-w0",
            "http://127.0.0.1:3960".to_string(),
            3.6,
        ));
        r.drain(&w.id);
        r.mark_unhealthy(&w.id);
        assert_eq!(r.get(&w.id).unwrap().status, WorkerStatus::Draining);
    }

    #[test]
    fn hub_price_overrides_billing_and_remembers_the_workers_own_price() {
        use std::collections::HashMap;
        let r = WorkerRegistry::new();
        let w = r.register(test_registration(
            "fp-w0",
            "http://127.0.0.1:3960".to_string(),
            3.6,
        ));

        assert_eq!(
            r.apply_hub_prices(&HashMap::from([("fp-w0".to_string(), 18.0)])),
            1
        );
        let got = r.get(&w.id).unwrap();
        assert_eq!(got.pricing.price_per_hour, 18.0);
        assert_eq!(got.reported_price_per_hour, Some(3.6));

        // A later override keeps the ORIGINAL reported price.
        r.apply_hub_prices(&HashMap::from([("fp-w0".to_string(), 25.0)]));
        let got = r.get(&w.id).unwrap();
        assert_eq!(got.pricing.price_per_hour, 25.0);
        assert_eq!(got.reported_price_per_hour, Some(3.6));

        // Re-applying the same price is a no-op.
        assert_eq!(
            r.apply_hub_prices(&HashMap::from([("fp-w0".to_string(), 25.0)])),
            0
        );
    }

    #[test]
    fn discovery_refresh_keeps_the_hub_price() {
        use std::collections::HashMap;
        let r = WorkerRegistry::new();
        let w = r.register(test_registration(
            "fp-w0",
            "http://127.0.0.1:3960".to_string(),
            3.6,
        ));
        r.apply_hub_prices(&HashMap::from([("fp-w0".to_string(), 18.0)]));
        r.update_resources(&w.id, WorkerResources::default(), HardwareInfo::default());
        r.refresh_heartbeat(&w.id);
        assert_eq!(r.get(&w.id).unwrap().pricing.price_per_hour, 18.0);
    }

    #[test]
    fn peer_workers_are_never_repriced() {
        use std::collections::HashMap;
        let r = WorkerRegistry::new();
        let mut peer = test_registration("fp-w0", "http://127.0.0.1:3960".to_string(), 3.6);
        peer.source_node = Some("node-other".to_string());
        let w = r.register(peer);
        assert_eq!(
            r.apply_hub_prices(&HashMap::from([("fp-w0".to_string(), 18.0)])),
            0
        );
        assert_eq!(r.get(&w.id).unwrap().pricing.price_per_hour, 3.6);
    }

    /// Fix round 1, Finding 1: a re-registration under the same name (a worker
    /// restart re-probing and rebuilding `pricing` from scratch) must not lose
    /// the hub's price, even with no sync in between the two registrations.
    #[test]
    fn hub_price_survives_re_registration_with_no_sync_in_between() {
        use std::collections::HashMap;
        let r = WorkerRegistry::new();
        let w1 = r.register(test_registration(
            "fp-w0",
            "http://127.0.0.1:3960".to_string(),
            3.6,
        ));
        r.apply_hub_prices(&HashMap::from([("fp-w0".to_string(), 18.0)]));
        assert_eq!(r.get(&w1.id).unwrap().pricing.price_per_hour, 18.0);

        // Re-register the SAME name: a fresh record (new id), freshly-probed
        // pricing (a different probed price, to prove it is the one recorded
        // as `reported_price_per_hour`, not just echoed back).
        let w2 = r.register(test_registration(
            "fp-w0",
            "http://127.0.0.1:3961".to_string(),
            4.2,
        ));
        assert_ne!(w1.id, w2.id, "re-registration creates a new record");

        let got = r.get(&w2.id).unwrap();
        assert_eq!(
            got.pricing.price_per_hour, 18.0,
            "billing must use the hub price immediately, no sync needed"
        );
        assert_eq!(
            got.reported_price_per_hour,
            Some(4.2),
            "sync must keep reporting the freshly-probed price, not the hub price"
        );
    }
}