zc2 0.0.4

//! Worker registry for tracking available compute workers.

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

/// 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,
        }
    }
}

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,
    /// Tailscale IP extracted from worker's URI
    #[serde(default)]
    pub tailscale_ip: Option<String>,
    /// Whether worker is running inside a Docker container
    #[serde(default)]
    pub is_docker: Option<bool>,
}

impl Worker {
    /// 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(),
            tailscale_ip: None,
            is_docker: 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,
    /// Tailscale IP (optional, extracted from URI if not provided)
    #[serde(default)]
    pub tailscale_ip: Option<String>,
    /// Whether worker is running inside a Docker container
    #[serde(default)]
    pub is_docker: Option<bool>,
}

/// 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>,
}

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

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

        // Extract IP from URI if tailscale_ip not explicitly provided.
        // If the extracted IP is loopback (127.x or ::1), fall back to
        // the auto-detected Tailscale IP so Docker nodes show their real IP.
        let tailscale_ip = registration.tailscale_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_tailscale_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.tailscale_ip = tailscale_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()));

        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());
        worker
    }

    /// 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 {
                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.workers.remove(id).map(|(_, w)| w)
    }

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

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

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

    /// Mark a worker as unhealthy
    pub fn mark_unhealthy(&self, worker_id: &str) {
        if let Some(mut worker) = self.workers.get_mut(worker_id) {
            worker.status = WorkerStatus::Unhealthy;
        }
    }
}

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

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

    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(),
            tailscale_ip: None,
            is_docker: None,
        }
    }

    // --- 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 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 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, 1 * 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);
    }

    #[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);
    }
}