orchard/model/

registry.rs

//! Model registry for tracking loaded models and their state.

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

use hf_hub::api::tokio::ApiBuilder;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use tokio::sync::{oneshot, Mutex, Notify, RwLock};

use crate::error::Error;
use crate::formatter::ChatFormatter;
use crate::ipc::client::IPCClient;
use crate::model::resolver::{ModelResolver, ResolvedModel};

/// Model load state machine.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ModelLoadState {
    /// Model not yet requested
    Idle,
    /// Downloading from HuggingFace
    Downloading,
    /// Loading weights into engine
    Loading,
    /// Waiting for engine activation
    Activating,
    /// Ready for inference
    Ready,
    /// Failed to load
    Failed,
}

impl std::fmt::Display for ModelLoadState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Idle => write!(f, "IDLE"),
            Self::Downloading => write!(f, "DOWNLOADING"),
            Self::Loading => write!(f, "LOADING"),
            Self::Activating => write!(f, "ACTIVATING"),
            Self::Ready => write!(f, "READY"),
            Self::Failed => write!(f, "FAILED"),
        }
    }
}

/// Information about a loaded model.
#[derive(Debug, Clone)]
pub struct ModelInfo {
    pub model_id: String,
    pub model_path: String,
    pub formatter: Option<Arc<ChatFormatter>>,
    pub capabilities: Option<HashMap<String, Vec<i32>>>,
    pub minimum_memory_bytes: Option<u64>,
}

impl ModelInfo {
    pub fn require_formatter(&self) -> std::result::Result<&ChatFormatter, crate::error::Error> {
        self.formatter.as_deref().ok_or_else(|| {
            crate::error::Error::ModelNotReady(format!(
                "Model '{}' does not have a chat formatter",
                self.model_id
            ))
        })
    }
}

/// Entry in the model registry tracking a model's state.
pub struct ModelEntry {
    pub state: ModelLoadState,
    pub info: Option<ModelInfo>,
    pub error: Option<String>,
    pub notify: Arc<Notify>,
    pub resolved: Option<ResolvedModel>,
    pub bytes_downloaded: Option<u64>,
    pub bytes_total: Option<u64>,
    /// Activation completion waiters (set when ACTIVATING)
    pub activation_waiters: Vec<oneshot::Sender<Result<(), String>>>,
}

impl Default for ModelEntry {
    fn default() -> Self {
        Self {
            state: ModelLoadState::Idle,
            info: None,
            error: None,
            notify: Arc::new(Notify::new()),
            resolved: None,
            bytes_downloaded: None,
            bytes_total: None,
            activation_waiters: Vec::new(),
        }
    }
}

#[derive(Debug, Deserialize)]
struct RemoteRepoFile {
    rfilename: String,
    #[serde(default)]
    size: Option<u64>,
}

#[derive(Debug, Deserialize)]
struct RemoteRepoInfo {
    siblings: Vec<RemoteRepoFile>,
}

/// Registry of loaded models.
pub struct ModelRegistry {
    entries: Arc<RwLock<HashMap<String, ModelEntry>>>,
    resolver: Mutex<ModelResolver>,
    alias_cache: RwLock<HashMap<String, String>>,
    /// IPC client for sending management commands to PIE
    ipc_client: RwLock<Option<Arc<IPCClient>>>,
}

impl ModelRegistry {
    /// Create a new model registry.
    pub fn new() -> Result<Self, Error> {
        Ok(Self {
            entries: Arc::new(RwLock::new(HashMap::new())),
            resolver: Mutex::new(ModelResolver::new()?),
            alias_cache: RwLock::new(HashMap::new()),
            ipc_client: RwLock::new(None),
        })
    }

    /// Set the IPC client for sending management commands to PIE.
    pub async fn set_ipc_client(&self, client: Arc<IPCClient>) {
        let mut ipc = self.ipc_client.write().await;
        *ipc = Some(client);
    }

    /// Ensure a model is loaded and ready.
    ///
    /// This will:
    /// 1. Resolve the model identifier
    /// 2. Download if needed
    /// 3. Load the formatter
    /// 4. Send load_model command to PIE
    /// 5. Wait for engine activation
    pub async fn ensure_loaded(&self, requested_model_id: &str) -> Result<ModelInfo, Error> {
        let (_state, canonical_id) = self
            .schedule_model(requested_model_id, false)
            .await
            .map_err(Error::ModelNotReady)?;

        // Wait for local readiness (download + formatter)
        let (state, info, error) = self
            .await_model(&canonical_id, None)
            .await
            .map_err(Error::ModelNotReady)?;

        if state == ModelLoadState::Failed {
            return Err(Error::ModelNotReady(error.unwrap_or_else(|| {
                format!("Model '{}' failed to load", canonical_id)
            })));
        }

        if state == ModelLoadState::Ready {
            return info
                .ok_or_else(|| Error::ModelNotReady("Model ready but info missing".to_string()));
        }

        // At this point we have LOADING state with info - need to activate on PIE
        let info = info.ok_or_else(|| {
            Error::ModelNotReady(format!("Model '{}' info missing", canonical_id))
        })?;

        // Check if already activating or ready
        {
            let entries = self.entries.read().await;
            if let Some(entry) = entries.get(&canonical_id) {
                if entry.state == ModelLoadState::Ready {
                    return entry
                        .info
                        .clone()
                        .ok_or_else(|| Error::ModelNotReady("Ready but no info".to_string()));
                }
            }
        }

        // Send load_model command and wait for activation
        let activation_rx = self
            .send_load_model_command(requested_model_id, &canonical_id, &info)
            .await
            .map_err(Error::ModelNotReady)?;

        // Wait for activation to complete
        match activation_rx.await {
            Ok(Ok(())) => {
                // Activation succeeded, get the ready info
                self.get_if_ready(&canonical_id).await.ok_or_else(|| {
                    Error::ModelNotReady(format!("Model '{}' failed to activate", canonical_id))
                })
            }
            Ok(Err(e)) => Err(Error::ModelNotReady(e)),
            Err(_) => Err(Error::ModelNotReady(format!(
                "Activation channel closed for '{}'",
                canonical_id
            ))),
        }
    }

    /// Resolve a model identifier and schedule any required local download work.
    pub async fn resolve_or_download(
        &self,
        requested_model_id: &str,
    ) -> Result<(ModelLoadState, String), String> {
        self.schedule_model(requested_model_id, false).await
    }

    /// Send the load_model command to PIE.
    async fn send_load_model_command(
        &self,
        requested_id: &str,
        canonical_id: &str,
        info: &ModelInfo,
    ) -> Result<oneshot::Receiver<Result<(), String>>, String> {
        // Create activation channel
        let (tx, rx) = oneshot::channel();

        // Set up activation state
        {
            let mut entries = self.entries.write().await;
            let entry = entries
                .get_mut(canonical_id)
                .ok_or_else(|| format!("Model '{}' not in registry", canonical_id))?;

            match entry.state {
                ModelLoadState::Ready => {
                    let _ = tx.send(Ok(()));
                    return Ok(rx);
                }
                ModelLoadState::Failed => {
                    let message = entry
                        .error
                        .clone()
                        .unwrap_or_else(|| format!("Model '{}' failed to load", canonical_id));
                    let _ = tx.send(Err(message));
                    return Ok(rx);
                }
                ModelLoadState::Activating => {
                    entry.activation_waiters.push(tx);
                    return Ok(rx);
                }
                _ => {
                    entry.state = ModelLoadState::Activating;
                    entry.activation_waiters.push(tx);
                }
            }
        }

        // Get IPC client
        let ipc = {
            let guard = self.ipc_client.read().await;
            guard
                .clone()
                .ok_or_else(|| "IPC client not set".to_string())?
        };

        // Build and send the command
        let command = json!({
            "type": "load_model",
            "requested_id": requested_id,
            "canonical_id": canonical_id,
            "model_path": info.model_path,
            "wait_for_completion": false,
        });

        let response = ipc
            .send_management_command_async(command, Duration::from_secs(30))
            .await
            .map_err(|e| format!("Failed to send load_model command: {}", e))?;

        // Check response status
        let status = response
            .get("status")
            .and_then(|v| v.as_str())
            .unwrap_or("");

        match status {
            "ok" => {
                // Immediate success - extract capabilities and mark ready
                let capabilities = self.parse_capabilities(&response);
                let minimum_memory_bytes = self.parse_minimum_memory_bytes(&response);
                self.complete_activation(canonical_id, capabilities, minimum_memory_bytes)
                    .await;
            }
            "accepted" => {
                // Async activation - wait for model_loaded event
                tracing::debug!(
                    "Model '{}' activation accepted, waiting for model_loaded event",
                    canonical_id
                );
            }
            _ => {
                let message = response
                    .get("message")
                    .and_then(|v| v.as_str())
                    .unwrap_or("unknown error");
                self.fail_activation(
                    canonical_id,
                    &format!("Engine rejected load_model: {}", message),
                )
                .await;
                return Err(format!(
                    "Engine rejected load_model for '{}': {}",
                    requested_id, message
                ));
            }
        }

        Ok(rx)
    }

    /// Parse capabilities from management response.
    fn parse_capabilities(&self, response: &Value) -> Option<HashMap<String, Vec<i32>>> {
        response
            .get("data")
            .and_then(|d| d.get("load_model"))
            .and_then(|lm| lm.get("capabilities"))
            .and_then(|c| c.as_object())
            .map(|obj| {
                obj.iter()
                    .filter_map(|(k, v)| {
                        let vals: Vec<i32> = if let Some(arr) = v.as_array() {
                            arr.iter()
                                .filter_map(|x| x.as_i64().map(|n| n as i32))
                                .collect()
                        } else if let Some(n) = v.as_i64() {
                            vec![n as i32]
                        } else {
                            return None;
                        };
                        Some((k.clone(), vals))
                    })
                    .collect()
            })
    }

    fn parse_minimum_memory_bytes(&self, response: &Value) -> Option<u64> {
        response
            .get("data")
            .and_then(|d| d.get("load_model"))
            .and_then(|lm| lm.get("minimum_memory_bytes"))
            .and_then(|value| value.as_u64())
    }

    /// Complete activation successfully.
    async fn complete_activation(
        &self,
        model_id: &str,
        capabilities: Option<HashMap<String, Vec<i32>>>,
        minimum_memory_bytes: Option<u64>,
    ) {
        let mut entries = self.entries.write().await;
        if let Some(entry) = entries.get_mut(model_id) {
            if let Some(ref mut info) = entry.info {
                if let Some(capabilities) = capabilities {
                    info.capabilities = Some(capabilities);
                }
                if minimum_memory_bytes.is_some() {
                    info.minimum_memory_bytes = minimum_memory_bytes;
                }
            }
            entry.state = ModelLoadState::Ready;
            entry.notify.notify_waiters();

            // Signal activation complete
            for tx in entry.activation_waiters.drain(..) {
                let _ = tx.send(Ok(()));
            }
        }
    }

    /// Fail activation with error.
    async fn fail_activation(&self, model_id: &str, error: &str) {
        let mut entries = self.entries.write().await;
        if let Some(entry) = entries.get_mut(model_id) {
            entry.state = ModelLoadState::Failed;
            entry.error = Some(error.to_string());
            entry.notify.notify_waiters();

            // Signal activation failed
            for tx in entry.activation_waiters.drain(..) {
                let _ = tx.send(Err(error.to_string()));
            }
        }
    }

    /// Schedule a model for loading.
    ///
    /// Returns the current state and canonical ID.
    pub async fn schedule_model(
        &self,
        requested_model_id: &str,
        force_reload: bool,
    ) -> Result<(ModelLoadState, String), String> {
        let resolved = {
            let mut resolver = self.resolver.lock().await;
            resolver
                .resolve(requested_model_id)
                .await
                .map_err(|e| e.to_string())?
        };

        let canonical_id = resolved.canonical_id.clone();

        {
            let mut alias_cache = self.alias_cache.write().await;
            alias_cache.insert(requested_model_id.to_lowercase(), canonical_id.clone());
            alias_cache
                .entry(canonical_id.to_lowercase())
                .or_insert_with(|| canonical_id.clone());
        }

        let mut entries = self.entries.write().await;
        let entry = entries
            .entry(canonical_id.clone())
            .or_insert_with(ModelEntry::default);

        if entry.state == ModelLoadState::Ready && !force_reload {
            return Ok((ModelLoadState::Ready, canonical_id));
        }

        if matches!(
            entry.state,
            ModelLoadState::Loading | ModelLoadState::Downloading | ModelLoadState::Activating
        ) && !force_reload
        {
            return Ok((entry.state, canonical_id));
        }

        if entry.state == ModelLoadState::Failed && !force_reload {
            return Ok((ModelLoadState::Failed, canonical_id));
        }

        entry.error = None;
        entry.info = None;
        entry.resolved = Some(resolved.clone());
        entry.bytes_downloaded = None;
        entry.bytes_total = None;
        entry.notify = Arc::new(Notify::new());
        entry.activation_waiters.clear();

        if resolved.source == "local" || resolved.source == "hf_cache" {
            let formatter = ChatFormatter::new(&resolved.model_path).ok().map(Arc::new);
            entry.info = Some(ModelInfo {
                model_id: canonical_id.clone(),
                model_path: resolved.model_path.to_string_lossy().to_string(),
                formatter,
                capabilities: None,
                minimum_memory_bytes: None,
            });
            entry.state = ModelLoadState::Loading;
            entry.notify.notify_waiters();
            return Ok((ModelLoadState::Loading, canonical_id));
        }

        // Model needs to be downloaded from HuggingFace
        entry.state = ModelLoadState::Downloading;
        let notify = entry.notify.clone();

        // Drop entries lock before spawning to avoid deadlock
        drop(entries);

        // Spawn download task
        let hf_repo = resolved
            .hf_repo
            .clone()
            .unwrap_or_else(|| resolved.canonical_id.clone());
        let canonical_id_for_task = canonical_id.clone();
        let entries_ref = self.entries.clone();

        tokio::spawn(async move {
            let result = Self::download_model(
                Arc::clone(&entries_ref),
                canonical_id_for_task.as_str(),
                hf_repo.as_str(),
            )
            .await;

            let mut entries: tokio::sync::RwLockWriteGuard<'_, HashMap<String, ModelEntry>> =
                entries_ref.write().await;
            if let Some(entry) = entries.get_mut(&canonical_id_for_task) {
                match result {
                    Ok(download_path) => {
                        // Update resolved path
                        if let Some(ref mut resolved) = entry.resolved {
                            resolved.model_path = download_path.clone();
                            resolved.source = "hf_cache".to_string();
                        }

                        let formatter = ChatFormatter::new(&download_path).ok().map(Arc::new);
                        entry.info = Some(ModelInfo {
                            model_id: canonical_id_for_task.clone(),
                            model_path: download_path.to_string_lossy().to_string(),
                            formatter,
                            capabilities: None,
                            minimum_memory_bytes: None,
                        });
                        entry.state = ModelLoadState::Loading;
                    }
                    Err(e) => {
                        entry.error = Some(format!("Download failed: {}", e));
                        entry.state = ModelLoadState::Failed;
                    }
                }
                notify.notify_waiters();
            }
        });

        Ok((ModelLoadState::Downloading, canonical_id))
    }

    /// Download a model from HuggingFace Hub.
    async fn download_model(
        entries_ref: Arc<RwLock<HashMap<String, ModelEntry>>>,
        canonical_id: &str,
        repo_id: &str,
    ) -> Result<std::path::PathBuf, String> {
        tracing::info!("Downloading model from HuggingFace: {}", repo_id);

        let api = ApiBuilder::from_env()
            .with_progress(false)
            .build()
            .map_err(|e| format!("Failed to create HF API: {}", e))?;
        let repo = api.model(repo_id.to_string());

        let mut repo_info: RemoteRepoInfo = repo
            .info_request()
            .query(&[("blobs", "true")])
            .send()
            .await
            .map_err(|e| format!("Failed to query repo info: {}", e))?
            .json()
            .await
            .map_err(|e| format!("Failed to decode repo info: {}", e))?;

        if repo_info.siblings.is_empty() {
            return Err("Repository has no downloadable files".to_string());
        }

        if !repo_info
            .siblings
            .iter()
            .any(|file| file.rfilename == "config.json")
        {
            return Err("Repository is missing config.json".to_string());
        }

        repo_info.siblings.sort_by(|left, right| {
            let left_priority = u8::from(left.rfilename != "config.json");
            let right_priority = u8::from(right.rfilename != "config.json");
            left_priority
                .cmp(&right_priority)
                .then_with(|| left.rfilename.cmp(&right.rfilename))
        });

        let total_bytes = repo_info
            .siblings
            .iter()
            .filter_map(|file| file.size)
            .sum::<u64>();
        let mut downloaded_bytes = 0u64;

        {
            let mut entries = entries_ref.write().await;
            if let Some(entry) = entries.get_mut(canonical_id) {
                entry.bytes_downloaded = Some(0);
                entry.bytes_total = Some(total_bytes);
            }
        }

        let mut model_dir: Option<std::path::PathBuf> = None;

        for file in repo_info.siblings {
            if Self::download_cancelled(&entries_ref, canonical_id).await {
                return Err("Cancelled".to_string());
            }

            let path = repo
                .get(file.rfilename.as_str())
                .await
                .map_err(|e| format!("Failed to download {}: {}", file.rfilename, e))?;

            if file.rfilename == "config.json" {
                model_dir = path.parent().map(|parent| parent.to_path_buf());
            }

            let file_size = tokio::fs::metadata(&path)
                .await
                .map(|metadata| metadata.len())
                .unwrap_or_else(|_| file.size.unwrap_or(0));
            downloaded_bytes = downloaded_bytes.saturating_add(file_size);

            let completed_bytes = if total_bytes == 0 {
                downloaded_bytes
            } else {
                downloaded_bytes.min(total_bytes)
            };

            let mut entries = entries_ref.write().await;
            if let Some(entry) = entries.get_mut(canonical_id) {
                entry.bytes_downloaded = Some(completed_bytes);
                entry.bytes_total = Some(total_bytes.max(completed_bytes));
            }
        }

        if Self::download_cancelled(&entries_ref, canonical_id).await {
            return Err("Cancelled".to_string());
        }

        let model_dir =
            model_dir.ok_or_else(|| "Downloaded repo is missing config.json".to_string())?;

        tracing::info!("Model downloaded to {:?}", model_dir);
        Ok(model_dir)
    }

    /// Cancel an in-progress download.
    pub async fn cancel_download(&self, model_id: &str) -> Result<(), String> {
        let canonical_id = self.canonicalize(model_id).await?;
        let mut entries = self.entries.write().await;
        let entry = entries
            .get_mut(&canonical_id)
            .ok_or_else(|| format!("Model '{}' has not been scheduled", canonical_id))?;

        if entry.state != ModelLoadState::Downloading {
            return Err(format!(
                "Model '{}' is not downloading (current state: {})",
                canonical_id, entry.state
            ));
        }

        entry.state = ModelLoadState::Failed;
        entry.error = Some("Cancelled".to_string());
        entry.notify.notify_waiters();
        Ok(())
    }

    /// Wait for a model to finish loading.
    ///
    /// Returns immediately if state is already terminal (Loading/Ready/Failed).
    /// Only blocks if state is Downloading or Activating.
    pub async fn await_model(
        &self,
        model_id: &str,
        timeout: Option<std::time::Duration>,
    ) -> Result<(ModelLoadState, Option<ModelInfo>, Option<String>), String> {
        let canonical_id = self.canonicalize(model_id).await?;

        let _notify: Arc<Notify>;
        let notified = {
            let entries = self.entries.read().await;
            let entry = entries
                .get(&canonical_id)
                .ok_or_else(|| format!("Model '{}' has not been scheduled", model_id))?;

            // Notify is edge-triggered: arm the waiter while holding the lock.
            if !matches!(
                entry.state,
                ModelLoadState::Downloading | ModelLoadState::Activating
            ) {
                return Ok((entry.state, entry.info.clone(), entry.error.clone()));
            }

            _notify = entry.notify.clone();
            _notify.notified()
        };

        // Wait for state transition
        match timeout {
            Some(d) => {
                let _ = tokio::time::timeout(d, notified).await;
            }
            None => notified.await,
        }

        // Re-read final state
        let entries = self.entries.read().await;
        let entry = entries
            .get(&canonical_id)
            .ok_or_else(|| format!("Model '{}' not found", canonical_id))?;
        Ok((entry.state, entry.info.clone(), entry.error.clone()))
    }

    /// Get model info if ready.
    pub async fn get_if_ready(&self, model_id: &str) -> Option<ModelInfo> {
        let canonical_id = self.canonicalize(model_id).await.ok()?;
        let entries = self.entries.read().await;
        let entry = entries.get(&canonical_id)?;

        if entry.state == ModelLoadState::Ready {
            entry.info.clone()
        } else {
            None
        }
    }

    /// Get model status.
    pub async fn get_status(
        &self,
        model_id: &str,
    ) -> (ModelLoadState, Option<String>, Option<(u64, u64)>) {
        let canonical_id = match self.canonicalize(model_id).await {
            Ok(id) => id,
            Err(_) => return (ModelLoadState::Idle, None, None),
        };

        let entries = self.entries.read().await;
        let entry = match entries.get(&canonical_id) {
            Some(e) => e,
            None => return (ModelLoadState::Idle, None, None),
        };

        let progress = match (entry.bytes_downloaded, entry.bytes_total) {
            (Some(d), Some(t)) => Some((d, t)),
            _ => None,
        };

        (entry.state, entry.error.clone(), progress)
    }

    /// Update capabilities for a loaded model.
    pub async fn update_capabilities(
        &self,
        model_id: &str,
        capabilities: HashMap<String, Vec<i32>>,
    ) {
        let canonical_id = match self.canonicalize(model_id).await {
            Ok(id) => id,
            Err(_) => {
                tracing::warn!("Received capabilities for unknown model '{}'", model_id);
                return;
            }
        };

        let mut entries = self.entries.write().await;
        if let Some(entry) = entries.get_mut(&canonical_id) {
            if let Some(ref mut info) = entry.info {
                info.capabilities = Some(capabilities);
            }
        }
    }

    /// Mark a model as ready (called when engine confirms activation).
    pub async fn mark_ready(&self, model_id: &str) {
        let canonical_id = match self.canonicalize(model_id).await {
            Ok(id) => id,
            Err(_) => return,
        };

        let mut entries = self.entries.write().await;
        if let Some(entry) = entries.get_mut(&canonical_id) {
            entry.state = ModelLoadState::Ready;
            entry.notify.notify_waiters();
            for tx in entry.activation_waiters.drain(..) {
                let _ = tx.send(Ok(()));
            }
        }
    }

    /// Mark a model as failed.
    pub async fn mark_failed(&self, model_id: &str, error: String) {
        let canonical_id = match self.canonicalize(model_id).await {
            Ok(id) => id,
            Err(_) => return,
        };

        let mut entries = self.entries.write().await;
        if let Some(entry) = entries.get_mut(&canonical_id) {
            entry.state = ModelLoadState::Failed;
            entry.error = Some(error.clone());
            entry.notify.notify_waiters();

            // Signal activation failed if waiting
            for tx in entry.activation_waiters.drain(..) {
                let _ = tx.send(Err(error.clone()));
            }
        }
    }

    /// Handle model_loaded event from PIE.
    ///
    /// Called by the event callback when a model_loaded event is received.
    pub async fn handle_model_loaded(&self, payload: &Value) {
        let model_id = match payload.get("model_id").and_then(|v| v.as_str()) {
            Some(id) => id,
            None => {
                tracing::warn!("Received model_loaded event without model_id");
                return;
            }
        };

        // Extract and update capabilities
        if let Some(caps) = payload.get("capabilities").and_then(|c| c.as_object()) {
            let capabilities: HashMap<String, Vec<i32>> = caps
                .iter()
                .filter_map(|(k, v)| {
                    let vals: Vec<i32> = if let Some(arr) = v.as_array() {
                        arr.iter()
                            .filter_map(|x| x.as_i64().map(|n| n as i32))
                            .collect()
                    } else if let Some(n) = v.as_i64() {
                        vec![n as i32]
                    } else {
                        return None;
                    };
                    Some((k.clone(), vals))
                })
                .collect();

            if !capabilities.is_empty() {
                self.update_capabilities(model_id, capabilities).await;
            }
        }

        // Canonicalize the model_id first (fix for issue #4)
        let canonical_id = match self.canonicalize(model_id).await {
            Ok(id) => id,
            Err(_) => {
                // Try the raw model_id as fallback
                tracing::debug!("Model '{}' not found in alias cache, using as-is", model_id);
                model_id.to_string()
            }
        };

        // Complete the activation
        let minimum_memory_bytes = payload
            .get("minimum_memory_bytes")
            .and_then(|value| value.as_u64());
        self.complete_activation(&canonical_id, None, minimum_memory_bytes)
            .await;
    }

    /// Handle model_load_failed event from PIE.
    ///
    /// Called by the event callback when a model_load_failed event is received.
    pub async fn handle_model_load_failed(&self, payload: &Value) {
        let model_id = match payload.get("model_id").and_then(|v| v.as_str()) {
            Some(id) => id,
            None => {
                tracing::warn!("Received model_load_failed event without model_id");
                return;
            }
        };

        let error = match payload.get("error").and_then(|v| v.as_str()) {
            Some(message) => message,
            None => {
                tracing::warn!(
                    "Received model_load_failed event without error for '{}'",
                    model_id
                );
                "unknown error"
            }
        };

        let canonical_id = match self.canonicalize(model_id).await {
            Ok(id) => id,
            Err(_) => {
                tracing::debug!("Model '{}' not found in alias cache, using as-is", model_id);
                model_id.to_string()
            }
        };

        self.fail_activation(&canonical_id, error).await;
    }

    /// List all registered models.
    pub async fn list_models(&self) -> Vec<HashMap<String, String>> {
        let entries = self.entries.read().await;
        let mut catalog = Vec::new();

        for (canonical_id, entry) in entries.iter() {
            if let Some(ref resolved) = entry.resolved {
                let mut payload: HashMap<String, String> = resolved.metadata.clone();
                payload.insert("canonical_id".to_string(), canonical_id.clone());
                payload.insert(
                    "model_path".to_string(),
                    resolved.model_path.to_string_lossy().to_string(),
                );
                payload.insert("source".to_string(), resolved.source.clone());
                payload.insert(
                    "hf_repo".to_string(),
                    resolved.hf_repo.clone().unwrap_or_default(),
                );
                payload.insert("state".to_string(), entry.state.to_string());
                catalog.push(payload);
            }
        }

        catalog
    }

    async fn canonicalize(&self, model_id: &str) -> Result<String, String> {
        // Check if it's already a canonical ID
        {
            let entries = self.entries.read().await;
            if entries.contains_key(model_id) {
                return Ok(model_id.to_string());
            }
        }

        // Check alias cache
        {
            let alias_cache = self.alias_cache.read().await;
            if let Some(canonical) = alias_cache.get(&model_id.to_lowercase()) {
                return Ok(canonical.clone());
            }
        }

        Err(format!("Model '{}' not found in registry", model_id))
    }

    async fn download_cancelled(
        entries_ref: &Arc<RwLock<HashMap<String, ModelEntry>>>,
        canonical_id: &str,
    ) -> bool {
        let entries = entries_ref.read().await;
        matches!(
            entries.get(canonical_id),
            Some(entry)
                if entry.state == ModelLoadState::Failed
                    && entry.error.as_deref() == Some("Cancelled")
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;
    use tokio::sync::oneshot;

    #[tokio::test]
    async fn test_registry_creation() {
        let registry = ModelRegistry::new().unwrap();
        let models = registry.list_models().await;
        assert!(models.is_empty());
    }

    #[tokio::test]
    async fn test_handle_model_load_failed_fails_activation_waiters() {
        let registry = ModelRegistry::new().unwrap();
        let canonical_id = "moondream/moondream3-preview".to_string();
        let requested_id = canonical_id.clone();
        let error = "Weight shard file not found".to_string();
        let (tx, rx) = oneshot::channel();

        {
            let mut alias_cache = registry.alias_cache.write().await;
            alias_cache.insert(requested_id.to_lowercase(), canonical_id.clone());
        }

        {
            let mut entries = registry.entries.write().await;
            let entry = entries.entry(canonical_id.clone()).or_default();
            entry.state = ModelLoadState::Activating;
            entry.activation_waiters.push(tx);
        }

        registry
            .handle_model_load_failed(&json!({
                "model_id": requested_id,
                "error": error,
            }))
            .await;

        assert_eq!(rx.await.unwrap(), Err(error.clone()));

        let entries = registry.entries.read().await;
        let entry = entries.get(&canonical_id).unwrap();
        assert_eq!(entry.state, ModelLoadState::Failed);
        assert_eq!(entry.error.as_deref(), Some(error.as_str()));
        assert!(entry.activation_waiters.is_empty());
    }

    #[test]
    fn test_model_load_state_display() {
        assert_eq!(ModelLoadState::Idle.to_string(), "IDLE");
        assert_eq!(ModelLoadState::Ready.to_string(), "READY");
        assert_eq!(ModelLoadState::Failed.to_string(), "FAILED");
    }
}