IdeA/crates/infrastructure/src/model_server/mod.rs

//! Infrastructure adapters for local model servers.

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use async_trait::async_trait;
use futures_util::StreamExt;
use serde::{Deserialize, Serialize};
use tokio::io::AsyncWriteExt;
use tokio::process::{Child, Command};
use tokio::sync::Mutex as AsyncMutex;

use domain::model_server::{
    ExecutablePath, HfModelRef, LlamaCppOptions, LocalModelRef, LocalModelServerConfig,
    LocalModelServerKind, ModelPath, ModelServerEndpoint, ModelSource,
};
use domain::ports::{
    FileSystem, ManagedProcess, ManagedProcessHandle, ModelArtifactCancel, ModelArtifactDownloader,
    ModelArtifactProgress, ModelArtifactResolution, ModelArtifactState, ModelServerArgv,
    ModelServerError, ModelServerProbe, ModelServerRegistry, ModelServerRuntime, ProcessStatus,
    RemotePath, SpawnSpec,
};
use domain::{LocalModelServerId, ProjectPath, StopPolicy};

/// HTTP readiness probe for OpenAI-compatible servers.
#[derive(Clone)]
pub struct HttpOpenAiCompatibleProbe {
    client: reqwest::Client,
}

impl Default for HttpOpenAiCompatibleProbe {
    fn default() -> Self {
        Self::new(Duration::from_secs(2))
    }
}

impl HttpOpenAiCompatibleProbe {
    /// Builds the probe with a request timeout.
    #[must_use]
    pub fn new(timeout: Duration) -> Self {
        let client = reqwest::Client::builder()
            .timeout(timeout)
            .build()
            .expect("reqwest client with timeout builds");
        Self { client }
    }
}

#[async_trait]
impl ModelServerProbe for HttpOpenAiCompatibleProbe {
    async fn probe(
        &self,
        endpoint: &ModelServerEndpoint,
    ) -> Result<domain::ModelServerStatus, ModelServerError> {
        let models = format!("{}/models", endpoint.base_url.trim_end_matches('/'));
        if is_ready(self.client.get(models).send().await) {
            return Ok(domain::ModelServerStatus::ReadyReused);
        }
        let root = endpoint.base_url.trim_end_matches("/v1").to_owned();
        if is_ready(self.client.get(root).send().await) {
            return Ok(domain::ModelServerStatus::ReadyReused);
        }
        Ok(domain::ModelServerStatus::Unreachable)
    }
}

fn is_ready(result: Result<reqwest::Response, reqwest::Error>) -> bool {
    result
        .map(|response| response.status().is_success())
        .unwrap_or(false)
}

/// Hugging Face artifact resolver/downloader backed by an IdeA-owned cache dir.
#[derive(Clone)]
pub struct HfModelArtifactDownloader {
    cache_dir: PathBuf,
    client: reqwest::Client,
    repo_locks: Arc<Mutex<HashMap<String, Arc<AsyncMutex<()>>>>>,
}

impl HfModelArtifactDownloader {
    /// Builds a downloader that stores artifacts under `cache_dir`.
    #[must_use]
    pub fn new(cache_dir: impl Into<PathBuf>) -> Self {
        Self {
            cache_dir: cache_dir.into(),
            client: reqwest::Client::new(),
            repo_locks: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Deterministic cache path for a Hugging Face model reference.
    ///
    /// This is the local path used for a single **merged** (non-sharded) GGUF file, renamed to a
    /// synthetic `{quant}.gguf` name. Sharded files are never renamed this way: llama.cpp
    /// discovers multi-part siblings by pattern-matching the *original* HF filename in the same
    /// directory, so shard files keep their remote name (see `local_path_for_filename`).
    #[must_use]
    pub fn cache_path_for(&self, repo: &HfModelRef) -> PathBuf {
        let (base, quant) = split_hf_ref(repo);
        let file_stem = quant.unwrap_or("model");
        self.cache_dir
            .join(base.replace('/', "--"))
            .join(format!("{}.gguf", safe_cache_component(file_stem)))
    }

    fn cache_dir_for(&self, repo: &HfModelRef) -> PathBuf {
        let (base, _) = split_hf_ref(repo);
        self.cache_dir.join(base.replace('/', "--"))
    }

    /// Local path for a remote filename that is part of a multi-shard set: the original name is
    /// preserved (sanitized) so llama.cpp's shard autoload can find siblings by pattern.
    fn local_path_for_shard(&self, repo: &HfModelRef, remote_filename: &str) -> PathBuf {
        self.cache_dir_for(repo)
            .join(safe_cache_component(remote_filename))
    }

    fn manifest_path_for(&self, repo: &HfModelRef) -> PathBuf {
        let (_, quant) = split_hf_ref(repo);
        let file_stem = quant.unwrap_or("model");
        self.cache_dir_for(repo)
            .join(format!("{}.manifest.json", safe_cache_component(file_stem)))
    }

    /// Reads a manifest of a previously-downloaded shard set, if present and complete on disk.
    fn cached_shard_set(&self, repo: &HfModelRef) -> Option<Vec<PathBuf>> {
        let manifest_path = self.manifest_path_for(repo);
        let raw = std::fs::read(&manifest_path).ok()?;
        let manifest: ShardManifest = serde_json::from_slice(&raw).ok()?;
        let dir = self.cache_dir_for(repo);
        let paths: Vec<PathBuf> = manifest
            .files
            .iter()
            .map(|filename| dir.join(filename))
            .collect();
        if paths.iter().all(|path| path.is_file()) {
            Some(paths)
        } else {
            None
        }
    }

    fn write_shard_manifest(&self, repo: &HfModelRef, filenames: &[String]) -> std::io::Result<()> {
        let manifest_path = self.manifest_path_for(repo);
        let manifest = ShardManifest {
            files: filenames.to_vec(),
        };
        let json = serde_json::to_vec(&manifest)?;
        std::fs::write(manifest_path, json)
    }

    fn lock_for(&self, repo: &HfModelRef) -> Arc<AsyncMutex<()>> {
        let mut locks = self.repo_locks.lock().expect("repo locks mutex poisoned");
        Arc::clone(
            locks
                .entry(repo.as_str().to_owned())
                .or_insert_with(|| Arc::new(AsyncMutex::new(()))),
        )
    }

    fn cached_state(&self, repo: &HfModelRef) -> Result<ModelArtifactState, ModelServerError> {
        let merged_path = self.cache_path_for(repo);
        if merged_path.is_file() {
            return Ok(ModelArtifactState::Downloaded {
                size_bytes: Some(file_size(&merged_path)?),
                path: model_path_from_pathbuf(merged_path)?,
            });
        }
        if let Some(paths) = self.cached_shard_set(repo) {
            if let Some(first) = paths.first() {
                return Ok(ModelArtifactState::Downloaded {
                    size_bytes: Some(paths_size(&paths)?),
                    path: model_path_from_pathbuf(first.clone())?,
                });
            }
        }
        Ok(ModelArtifactState::Missing)
    }

    fn delete_cached(&self, repo: &HfModelRef) -> Result<(), ModelServerError> {
        let merged_path = self.cache_path_for(repo);
        if merged_path.is_file() {
            remove_file_if_exists(&merged_path)?;
        }

        let manifest_path = self.manifest_path_for(repo);
        if let Ok(raw) = std::fs::read(&manifest_path) {
            let manifest: ShardManifest =
                serde_json::from_slice(&raw).map_err(|e| ModelServerError::Store(e.to_string()))?;
            let dir = self.cache_dir_for(repo);
            for filename in manifest.files {
                remove_file_if_exists(&dir.join(filename))?;
            }
            remove_file_if_exists(&manifest_path)?;
        }

        let repo_dir = self.cache_dir_for(repo);
        if repo_dir.is_dir()
            && std::fs::read_dir(&repo_dir)
                .map(is_empty_dir)
                .unwrap_or(false)
        {
            std::fs::remove_dir(&repo_dir).map_err(|e| ModelServerError::Store(e.to_string()))?;
        }
        Ok(())
    }

    async fn resolve_remote_filenames(
        &self,
        repo: &HfModelRef,
    ) -> Result<Vec<String>, ModelServerError> {
        let (base, quant) = split_hf_ref(repo);
        let url = format!("https://huggingface.co/api/models/{base}");
        let response = self
            .client
            .get(url)
            .send()
            .await
            .map_err(|e| ModelServerError::Probe(e.to_string()))?;
        if !response.status().is_success() {
            return Err(ModelServerError::Probe(format!(
                "huggingface model metadata returned {}",
                response.status()
            )));
        }
        let metadata: HfModelMetadata = response
            .json()
            .await
            .map_err(|e| ModelServerError::Probe(e.to_string()))?;
        let files = select_gguf_files(&metadata.siblings, quant);
        if files.is_empty() {
            return Err(ModelServerError::PathNotAccessible(format!(
                "no matching .gguf artifact found for {}",
                repo.as_str()
            )));
        }
        Ok(files)
    }
}

#[derive(Debug, Default, Serialize, Deserialize)]
struct ShardManifest {
    files: Vec<String>,
}

#[async_trait]
impl ModelArtifactDownloader for HfModelArtifactDownloader {
    async fn hf_model_state(
        &self,
        repo: &HfModelRef,
    ) -> Result<ModelArtifactState, ModelServerError> {
        let lock = self.lock_for(repo);
        let _guard = lock.lock().await;
        self.cached_state(repo)
    }

    async fn resolve_hf_model(
        &self,
        repo: &HfModelRef,
        progress: std::sync::Arc<dyn Fn(ModelArtifactProgress) + Send + Sync>,
        cancel: ModelArtifactCancel,
    ) -> Result<ModelArtifactResolution, ModelServerError> {
        let lock = self.lock_for(repo);
        let _guard = lock.lock().await;
        if cancel.is_cancelled() {
            return Err(ModelServerError::Cancelled);
        }

        // Fast path 1: a single merged file was already downloaded (synthetic name), no network needed.
        let merged_path = self.cache_path_for(repo);
        if merged_path.is_file() {
            return Ok(ModelArtifactResolution {
                path: model_path_from_pathbuf(merged_path)?,
                cache_hit: true,
            });
        }
        // Fast path 2: a previously-downloaded shard set is fully present on disk, no network needed.
        if let Some(paths) = self.cached_shard_set(repo) {
            if let Some(first) = paths.into_iter().next() {
                return Ok(ModelArtifactResolution {
                    path: model_path_from_pathbuf(first)?,
                    cache_hit: true,
                });
            }
        }

        let filenames = self.resolve_remote_filenames(repo).await?;
        if cancel.is_cancelled() {
            return Err(ModelServerError::Cancelled);
        }
        let is_shard_set = filenames.len() > 1;
        let targets: Vec<PathBuf> = if is_shard_set {
            filenames
                .iter()
                .map(|filename| self.local_path_for_shard(repo, filename))
                .collect()
        } else {
            vec![merged_path.clone()]
        };

        let (base, _) = split_hf_ref(repo);
        let total_files = filenames.len() as u64;
        for (index, (filename, target)) in filenames.iter().zip(targets.iter()).enumerate() {
            if cancel.is_cancelled() {
                return Err(ModelServerError::Cancelled);
            }
            let url = format!(
                "https://huggingface.co/{base}/resolve/main/{}",
                url_path_segment(filename)
            );
            let response = self
                .client
                .get(url)
                .send()
                .await
                .map_err(|e| ModelServerError::Probe(e.to_string()))?;
            if !response.status().is_success() {
                return Err(ModelServerError::Probe(format!(
                    "huggingface artifact download returned {}",
                    response.status()
                )));
            }
            if let Some(parent) = target.parent() {
                tokio::fs::create_dir_all(parent)
                    .await
                    .map_err(|e| ModelServerError::Store(e.to_string()))?;
            }
            let tmp_path = PathBuf::from(format!("{}.part", target.to_string_lossy()));
            let mut file = tokio::fs::File::create(&tmp_path)
                .await
                .map_err(|e| ModelServerError::Store(e.to_string()))?;
            let total = response.content_length();
            let mut downloaded = 0_u64;
            let mut stream = response.bytes_stream();
            while let Some(chunk) = stream.next().await {
                if cancel.is_cancelled() {
                    let _ = tokio::fs::remove_file(&tmp_path).await;
                    return Err(ModelServerError::Cancelled);
                }
                let chunk = chunk.map_err(|e| ModelServerError::Probe(e.to_string()))?;
                file.write_all(&chunk)
                    .await
                    .map_err(|e| ModelServerError::Store(e.to_string()))?;
                downloaded += chunk.len() as u64;
                progress(ModelArtifactProgress {
                    downloaded_bytes: Some(downloaded + index as u64 * total.unwrap_or(0)),
                    total_bytes: total.map(|t| t * total_files),
                    source: Some(repo.as_str().to_owned()),
                });
            }
            file.flush()
                .await
                .map_err(|e| ModelServerError::Store(e.to_string()))?;
            drop(file);
            tokio::fs::rename(&tmp_path, target)
                .await
                .map_err(|e| ModelServerError::Store(e.to_string()))?;
        }

        if is_shard_set {
            self.write_shard_manifest(repo, &filenames)
                .map_err(|e| ModelServerError::Store(e.to_string()))?;
        }

        Ok(ModelArtifactResolution {
            path: model_path_from_pathbuf(targets[0].clone())?,
            cache_hit: false,
        })
    }

    async fn delete_hf_model(&self, repo: &HfModelRef) -> Result<(), ModelServerError> {
        let lock = self.lock_for(repo);
        let _guard = lock.lock().await;
        self.delete_cached(repo)
    }
}

#[derive(Debug, Deserialize)]
struct HfModelMetadata {
    #[serde(default)]
    siblings: Vec<HfSibling>,
}

#[derive(Debug, Deserialize)]
struct HfSibling {
    #[serde(rename = "rfilename")]
    filename: String,
}

fn split_hf_ref(repo: &HfModelRef) -> (&str, Option<&str>) {
    repo.as_str()
        .split_once(':')
        .map_or((repo.as_str(), None), |(base, quant)| (base, Some(quant)))
}

/// Matches the llama.cpp / HF multi-part GGUF naming convention, e.g.
/// `model-q5_k_m-00001-of-00002.gguf`. Captures the shard index and total shard count.
fn shard_suffix() -> &'static regex::Regex {
    static RE: std::sync::OnceLock<regex::Regex> = std::sync::OnceLock::new();
    RE.get_or_init(|| regex::Regex::new(r"-(\d{5})-of-(\d{5})\.gguf$").expect("valid regex"))
}

fn shard_index(filename: &str) -> Option<u32> {
    shard_suffix()
        .captures(filename)
        .and_then(|caps| caps.get(1))
        .and_then(|m| m.as_str().parse().ok())
}

/// Resolves the ordered set of remote filenames to download for a quantization: the merged
/// (non-sharded) file if one matches, otherwise all matching shards sorted by index. A merged
/// file always takes priority over a shard set when both exist in the repo, since it is what
/// `llama-server --model` expects and a shard set would otherwise be picked first by naive
/// alphabetical sort ('-' sorts before '.' in ASCII).
fn select_gguf_files(siblings: &[HfSibling], quant: Option<&str>) -> Vec<String> {
    let ggufs: Vec<&str> = siblings
        .iter()
        .map(|sibling| sibling.filename.as_str())
        .filter(|filename| filename.ends_with(".gguf"))
        .collect();

    let matching: Vec<&str> = match quant {
        None => ggufs,
        Some(quant) => {
            let quant = quant.to_ascii_lowercase();
            ggufs
                .into_iter()
                .filter(|filename| filename.to_ascii_lowercase().contains(&quant))
                .collect()
        }
    };

    let (mut merged, mut shards): (Vec<&str>, Vec<&str>) = (Vec::new(), Vec::new());
    for filename in matching {
        if shard_index(filename).is_some() {
            shards.push(filename);
        } else {
            merged.push(filename);
        }
    }

    if let Some(first_merged) = {
        merged.sort_unstable();
        merged.first()
    } {
        return vec![(*first_merged).to_owned()];
    }

    shards.sort_unstable_by_key(|filename| shard_index(filename).unwrap_or(u32::MAX));
    shards.into_iter().map(str::to_owned).collect()
}

fn safe_cache_component(raw: &str) -> String {
    raw.chars()
        .map(|ch| match ch {
            'A'..='Z' | 'a'..='z' | '0'..='9' | '.' | '_' | '-' => ch,
            _ => '_',
        })
        .collect()
}

fn url_path_segment(raw: &str) -> String {
    raw.split('/')
        .map(url_component)
        .collect::<Vec<_>>()
        .join("/")
}

fn url_component(raw: &str) -> String {
    raw.bytes()
        .flat_map(|byte| match byte {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'.' | b'_' | b'-' => {
                vec![byte as char]
            }
            other => format!("%{other:02X}").chars().collect(),
        })
        .collect()
}

fn model_path_from_pathbuf(path: PathBuf) -> Result<ModelPath, ModelServerError> {
    ModelPath::new(path.to_string_lossy().into_owned())
        .map_err(|e| ModelServerError::Invalid(e.to_string()))
}

fn file_size(path: &Path) -> Result<u64, ModelServerError> {
    std::fs::metadata(path)
        .map(|metadata| metadata.len())
        .map_err(|e| ModelServerError::Store(e.to_string()))
}

fn paths_size(paths: &[PathBuf]) -> Result<u64, ModelServerError> {
    paths
        .iter()
        .map(|path| file_size(path))
        .try_fold(0_u64, |acc, size| size.map(|size| acc.saturating_add(size)))
}

fn remove_file_if_exists(path: &Path) -> Result<(), ModelServerError> {
    match std::fs::remove_file(path) {
        Ok(()) => Ok(()),
        Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
        Err(err) => Err(ModelServerError::Store(err.to_string())),
    }
}

fn is_empty_dir(entries: std::fs::ReadDir) -> bool {
    entries.into_iter().next().is_none()
}

/// Builds `llama-server` argv without shell interpolation.
#[derive(Debug, Default, Clone, Copy)]
pub struct LlamaCppRuntime;

impl LlamaCppRuntime {
    /// Creates the runtime.
    #[must_use]
    pub const fn new() -> Self {
        Self
    }
}

impl ModelServerRuntime for LlamaCppRuntime {
    fn build_argv(
        &self,
        config: &LocalModelServerConfig,
    ) -> Result<ModelServerArgv, ModelServerError> {
        if config.kind != LocalModelServerKind::LlamaCpp {
            return Err(ModelServerError::Invalid(
                "only LlamaCpp local model servers are supported".to_owned(),
            ));
        }
        let command = resolve_binary(
            config
                .binary
                .as_ref()
                .map(|binary| binary.as_str())
                .unwrap_or("llama-server"),
        )?;
        let mut args = Vec::new();
        match config
            .model
            .source
            .as_ref()
            .ok_or_else(|| ModelServerError::PathNotAccessible("model.source missing".to_owned()))?
        {
            ModelSource::LocalPath { path } => {
                args.push("--model".to_owned());
                args.push(path.as_str().to_owned());
            }
            ModelSource::HuggingFace { repo } => {
                args.push("-hf".to_owned());
                args.push(repo.as_str().to_owned());
            }
        }
        args.extend([
            "--port".to_owned(),
            config.endpoint.port.to_string(),
            "--host".to_owned(),
            config.options.host.clone(),
        ]);
        if let Some(gpu_layers) = config.options.gpu_layers {
            args.extend(["-ngl".to_owned(), gpu_layers.to_string()]);
        }
        if let Some(context_size) = config.options.context_size {
            args.extend(["-c".to_owned(), context_size.to_string()]);
        }
        if config.options.jinja {
            args.push("--jinja".to_owned());
        }
        args.extend(config.args.clone());
        Ok(ModelServerArgv { command, args })
    }

    fn build_spawn_spec(
        &self,
        config: &LocalModelServerConfig,
    ) -> Result<SpawnSpec, ModelServerError> {
        let argv = self.build_argv(config)?;
        Ok(SpawnSpec {
            command: argv.command,
            args: argv.args,
            cwd: ProjectPath::new("/").map_err(|e| ModelServerError::Invalid(e.to_string()))?,
            env: Vec::new(),
            context_plan: None,
            sandbox: None,
        })
    }
}

fn resolve_binary(raw: &str) -> Result<String, ModelServerError> {
    if is_path_like(raw) {
        let path = Path::new(raw);
        if path.is_file() {
            return Ok(raw.to_owned());
        }
        return Err(ModelServerError::PathNotAccessible(raw.to_owned()));
    }
    if let Some(path) = find_in_path(raw) {
        return Ok(path.to_string_lossy().into_owned());
    }
    Err(ModelServerError::PathNotAccessible(format!(
        "{raw} not found in PATH"
    )))
}

fn is_path_like(raw: &str) -> bool {
    raw.contains('/') || raw.contains('\\') || Path::new(raw).is_absolute()
}

fn find_in_path(command: &str) -> Option<PathBuf> {
    std::env::var_os("PATH").and_then(|path| {
        std::env::split_paths(&path)
            .map(|dir| dir.join(command))
            .find(|candidate| candidate.is_file())
    })
}

/// Local child-process manager for long-lived model servers.
#[derive(Default)]
pub struct LocalManagedProcess {
    children: Mutex<HashMap<String, Child>>,
}

impl LocalManagedProcess {
    /// Creates an empty process manager.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }
}

#[async_trait]
impl ManagedProcess for LocalManagedProcess {
    async fn spawn(&self, spec: SpawnSpec) -> Result<ManagedProcessHandle, ModelServerError> {
        let mut command = Command::new(&spec.command);
        command.args(&spec.args);
        if spec.cwd.as_str() != "/" {
            command.current_dir(spec.cwd.as_str());
        }
        for (key, value) in &spec.env {
            command.env(key, value);
        }
        let child = command
            .spawn()
            .map_err(|e| ModelServerError::Process(format!("{}: {e}", spec.command)))?;
        let id = uuid::Uuid::new_v4().to_string();
        self.children.lock().unwrap().insert(id.clone(), child);
        Ok(ManagedProcessHandle { id })
    }

    async fn kill(&self, handle: &ManagedProcessHandle) -> Result<(), ModelServerError> {
        let Some(mut child) = self.children.lock().unwrap().remove(&handle.id) else {
            return Ok(());
        };
        child
            .start_kill()
            .map_err(|e| ModelServerError::Process(e.to_string()))
    }

    async fn status(
        &self,
        handle: &ManagedProcessHandle,
    ) -> Result<ProcessStatus, ModelServerError> {
        let mut children = self.children.lock().unwrap();
        let Some(child) = children.get_mut(&handle.id) else {
            return Ok(ProcessStatus::Unknown);
        };
        match child
            .try_wait()
            .map_err(|e| ModelServerError::Process(e.to_string()))?
        {
            Some(status) => Ok(ProcessStatus::Exited {
                code: status.code(),
            }),
            None => Ok(ProcessStatus::Running),
        }
    }
}

/// File name of the global local-model-server registry.
const MODEL_SERVERS_FILE: &str = "model-servers.json";
const MODEL_SERVERS_VERSION: u32 = 2;

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct ModelServersDoc {
    version: u32,
    servers: Vec<LocalModelServerConfig>,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
enum PersistedModelServersDoc {
    V2(ModelServersDoc),
    V1(ModelServersDocV1),
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct ModelServersDocV1 {
    servers: Vec<LocalModelServerConfigV1>,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct LocalModelServerConfigV1 {
    id: LocalModelServerId,
    kind: LocalModelServerKind,
    name: String,
    endpoint: ModelServerEndpoint,
    model: LocalModelRefV1,
    #[serde(default)]
    binary: Option<ExecutablePath>,
    #[serde(default)]
    args: Vec<String>,
    auto_start: bool,
    #[serde(default = "default_stop_policy")]
    stop_policy: StopPolicy,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct LocalModelRefV1 {
    id: String,
    label: String,
    #[serde(default)]
    path: Option<ModelPath>,
    served_name: String,
}

fn default_stop_policy() -> StopPolicy {
    StopPolicy::StopOnAppExit
}

impl From<ModelServersDocV1> for ModelServersDoc {
    fn from(doc: ModelServersDocV1) -> Self {
        Self {
            version: MODEL_SERVERS_VERSION,
            servers: doc
                .servers
                .into_iter()
                .map(|server| LocalModelServerConfig {
                    id: server.id,
                    kind: server.kind,
                    name: server.name,
                    endpoint: server.endpoint,
                    model: LocalModelRef {
                        id: server.model.id,
                        label: server.model.label,
                        source: server
                            .model
                            .path
                            .map(|path| ModelSource::LocalPath { path }),
                        served_name: server.model.served_name,
                    },
                    binary: server.binary,
                    options: LlamaCppOptions::default(),
                    args: server.args,
                    auto_start: server.auto_start,
                    stop_policy: server.stop_policy,
                    warmup_deadline_secs: None,
                })
                .collect(),
        }
    }
}

impl Default for ModelServersDoc {
    fn default() -> Self {
        Self {
            version: MODEL_SERVERS_VERSION,
            servers: Vec::new(),
        }
    }
}

/// Filesystem-backed global model-server registry.
#[derive(Clone)]
pub struct FsModelServerRegistry {
    fs: std::sync::Arc<dyn FileSystem>,
    app_data_dir: String,
}

impl FsModelServerRegistry {
    /// Builds the registry from a filesystem and global app-data dir.
    #[must_use]
    pub fn new(fs: std::sync::Arc<dyn FileSystem>, app_data_dir: impl Into<String>) -> Self {
        Self {
            fs,
            app_data_dir: app_data_dir.into(),
        }
    }

    fn path(&self) -> RemotePath {
        let base = self.app_data_dir.trim_end_matches(['/', '\\']);
        RemotePath::new(format!("{base}/{MODEL_SERVERS_FILE}"))
    }

    async fn read_doc(&self) -> Result<ModelServersDoc, ModelServerError> {
        match self.fs.read(&self.path()).await {
            Ok(bytes) => match serde_json::from_slice::<PersistedModelServersDoc>(&bytes)
                .map_err(|e| ModelServerError::Store(format!("serialization failed: {e}")))?
            {
                PersistedModelServersDoc::V2(doc) => Ok(ModelServersDoc {
                    version: MODEL_SERVERS_VERSION,
                    servers: doc.servers,
                }),
                PersistedModelServersDoc::V1(doc) => Ok(doc.into()),
            },
            Err(domain::ports::FsError::NotFound(_)) => Ok(ModelServersDoc::default()),
            Err(domain::ports::FsError::PermissionDenied(p)) => {
                Err(ModelServerError::PermissionDenied(p))
            }
            Err(e) => Err(ModelServerError::Store(e.to_string())),
        }
    }

    async fn write_doc(&self, doc: &ModelServersDoc) -> Result<(), ModelServerError> {
        let dir = RemotePath::new(self.app_data_dir.trim_end_matches(['/', '\\']).to_owned());
        self.fs.create_dir_all(&dir).await.map_err(|e| match e {
            domain::ports::FsError::PermissionDenied(p) => ModelServerError::PermissionDenied(p),
            other => ModelServerError::Store(other.to_string()),
        })?;
        let bytes = serde_json::to_vec_pretty(doc)
            .map_err(|e| ModelServerError::Store(format!("serialization failed: {e}")))?;
        self.fs
            .write(&self.path(), &bytes)
            .await
            .map_err(|e| match e {
                domain::ports::FsError::PermissionDenied(p) => {
                    ModelServerError::PermissionDenied(p)
                }
                other => ModelServerError::Store(other.to_string()),
            })
    }
}

#[async_trait]
impl ModelServerRegistry for FsModelServerRegistry {
    async fn get(
        &self,
        id: &LocalModelServerId,
    ) -> Result<Option<LocalModelServerConfig>, ModelServerError> {
        Ok(self
            .read_doc()
            .await?
            .servers
            .into_iter()
            .find(|server| &server.id == id))
    }

    async fn list(&self) -> Result<Vec<LocalModelServerConfig>, ModelServerError> {
        Ok(self.read_doc().await?.servers)
    }

    async fn save(&self, config: LocalModelServerConfig) -> Result<(), ModelServerError> {
        config
            .options
            .validate()
            .map_err(|e| ModelServerError::Invalid(e.to_string()))?;
        domain::model_server::validate_free_args(&config.args)
            .map_err(|e| ModelServerError::Invalid(e.to_string()))?;
        let mut doc = self.read_doc().await?;
        doc.version = MODEL_SERVERS_VERSION;
        if let Some(slot) = doc.servers.iter_mut().find(|server| server.id == config.id) {
            *slot = config;
        } else {
            doc.servers.push(config);
        }
        self.write_doc(&doc).await
    }

    async fn delete(&self, id: LocalModelServerId) -> Result<(), ModelServerError> {
        let mut doc = self.read_doc().await?;
        doc.servers.retain(|server| server.id != id);
        self.write_doc(&doc).await
    }
}

#[cfg(test)]
mod gguf_selection_tests {
    use super::{select_gguf_files, HfSibling};

    fn sibling(name: &str) -> HfSibling {
        HfSibling {
            filename: name.to_owned(),
        }
    }

    #[test]
    fn prefers_merged_file_over_colliding_shards() {
        let siblings = vec![
            sibling("qwen2.5-coder-7b-instruct-q5_k_m-00001-of-00002.gguf"),
            sibling("qwen2.5-coder-7b-instruct-q5_k_m-00002-of-00002.gguf"),
            sibling("qwen2.5-coder-7b-instruct-q5_k_m.gguf"),
        ];

        let files = select_gguf_files(&siblings, Some("q5_k_m"));

        assert_eq!(files, vec!["qwen2.5-coder-7b-instruct-q5_k_m.gguf"]);
    }

    #[test]
    fn selects_all_shards_sorted_by_index_when_no_merged_file() {
        let siblings = vec![
            sibling("model-q4_k_m-00002-of-00003.gguf"),
            sibling("model-q4_k_m-00001-of-00003.gguf"),
            sibling("model-q4_k_m-00003-of-00003.gguf"),
            sibling("model-q8_0-00001-of-00002.gguf"),
        ];

        let files = select_gguf_files(&siblings, Some("q4_k_m"));

        assert_eq!(
            files,
            vec![
                "model-q4_k_m-00001-of-00003.gguf",
                "model-q4_k_m-00002-of-00003.gguf",
                "model-q4_k_m-00003-of-00003.gguf",
            ]
        );
    }

    #[test]
    fn returns_empty_when_no_match() {
        let siblings = vec![sibling("model-q4_k_m.gguf")];

        assert!(select_gguf_files(&siblings, Some("q5_k_m")).is_empty());
    }
}