//! A named registry for the models and providers your application uses,
//! local and remote, in one place.
//!
//! Register any provider — a remote API client ([`Model`]) or a local
//! in-process backend ([`LocalModel`], or one that is both) — under a name,
//! then dispatch by name with [`ModelManager::complete`] /
//! [`ModelManager::stream`], or pull the instance back out with
//! [`ModelManager::get`] to pass around and compose. This is the "scale up
//! from one provider to many" tier: construct a provider directly for the
//! simple case, reach for the manager when you juggle several.
//!
//! # Memory budgeting (local weights)
//!
//! For entries that own real weights, the manager doubles as a
//! memory-budget bookkeeper with per-pool LRU eviction. Each [`LocalModel`]
//! carries an estimated footprint organised by [`Pool`] (one budget for
//! host RAM, one per GPU). Loading a model that would exceed its pool's
//! budget evicts the least-recently-used loaded model **in the same pool**
//! until it fits; models in different pools never evict each other. Remote
//! providers ([`Pool::Remote`]) hold no local weights, so they never
//! participate in eviction and never count against a budget.
//!
//! These budgets answer "will this fit?" — not "will this run fast?".
//! Whether a 70B model loaded on CPU is *useful* at 1–3 tok/s is a
//! workload-choice question that this manager intentionally does not
//! answer. It only prevents OOM.
//!
//! # Dispatch is the text [`Model`] surface
//!
//! [`ModelManager::complete`] / [`ModelManager::stream`] dispatch the chat
//! completion interface. Media providers (image / video / audio) can still
//! be registered for lifecycle and budget bookkeeping, but are invoked on
//! the concrete handle you hold (or fetch via [`ModelManager::get`]), not
//! through `complete`.

pub mod hf_loader;

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::pin::Pin;
use std::sync::Arc;
use std::time::Instant;

use blazen_llm::{
    AdapterHandle, AdapterOptions, AdapterStatus, BlazenError, LocalModel, Model, ModelRequest,
    ModelResponse, Pool, StreamChunk,
};
use futures_util::Stream;
use tokio::sync::Mutex;

#[cfg(feature = "hf-loader")]
use crate::hf_loader::{BackendHint, HfLoadOptions};

#[cfg(feature = "training")]
pub use blazen_train::{
    BlazenTrainError, DistributedConfig, DpoConfig, FullFineTuneConfig, FullFineTuneResult,
    GrpoConfig, GrpoDataset, GrpoTrainer, KtoConfig, LoraConfig, MixedPrecision, OptimConfig,
    OrpoConfig, PpoConfig, PpoDataset, PpoTrainer, PreferenceDataset, RatedDataset,
    SchedulerConfig, SchedulerKind, SimpoConfig, TrainConfig, TrainCoreConfig, TrainedAdapter,
    TrainingDataset, TrainingEvent, TrainingProgress,
};

// Why: the concrete JSONL loaders (`PreferenceJsonlDataset`,
// `RatedJsonlDataset`, `JsonlDataset`) intentionally stay namespaced under
// `blazen_train::dataset::*` to match the existing convention — only the
// trait + result types live at the root re-export surface. Callers that
// want a JSONL dataset construct it through `blazen_train::dataset::...`
// and pass it as an `Arc<dyn PreferenceDataset>` / `Arc<dyn RatedDataset>`
// / `Box<dyn TrainingDataset>` into the verbs below.
//
// The `From<BlazenTrainError> for BlazenError` impl that powers `?` inside
// every training verb lives in `blazen-train` under its
// `blazen-llm-interop` feature (orphan rule — neither type is local to
// `blazen-manager`). `training` activates that feature, and the existing
// match arms already cover Cancelled / InvalidConfig / Unsupported / all
// other variants, so the PR8 verbs reuse the same `?`-propagation path
// without any new conversion code.

/// Status of a registered model.
#[derive(Debug, Clone)]
pub struct ModelStatus {
    /// Identifier of the model.
    pub id: String,
    /// Whether the model is currently loaded.
    pub loaded: bool,
    /// Estimated memory footprint in bytes. Includes the base model plus any
    /// mounted adapters.
    pub memory_estimate_bytes: u64,
    /// Pool this model is registered against.
    pub pool: Pool,
    /// Adapters currently mounted on this model (empty if none).
    pub adapters: Vec<AdapterStatus>,
}

struct RegisteredModel {
    /// Chat-completion dispatch target. `Some` for remote providers and for
    /// local backends that implement [`Model`]; `None` for lifecycle-only
    /// entries (e.g. a local media backend registered purely for memory
    /// budgeting). At least one of `model` / `local` is always `Some`.
    model: Option<Arc<dyn Model>>,
    /// Local-weights lifecycle handle. `Some` when the entry owns in-process
    /// weights that can be loaded / unloaded / adapter-mounted; `None` for
    /// remote providers.
    local: Option<Arc<dyn LocalModel>>,
    memory_estimate_bytes: u64,
    pool: Pool,
    loaded: bool,
    last_used: Option<Instant>,
    adapters: HashMap<String, AdapterStatus>,
}

/// Memory-budget-aware model manager with per-pool LRU eviction.
///
/// See the crate-level docs for the capacity-vs-performance distinction.
pub struct ModelManager {
    pool_budgets: HashMap<Pool, u64>,
    state: Mutex<HashMap<String, RegisteredModel>>,
}

impl ModelManager {
    /// Create a new manager with the given per-pool budgets (in bytes).
    ///
    /// Pools not present in the map have an implicit budget of 0, meaning
    /// no model targeting that pool can ever load. Add explicit entries for
    /// every pool you expect to use.
    #[must_use]
    pub fn new(pool_budgets: HashMap<Pool, u64>) -> Self {
        Self {
            pool_budgets,
            state: Mutex::new(HashMap::new()),
        }
    }

    /// Convenience constructor for the common single-GPU desktop case:
    /// one CPU pool sized in GB and one GPU pool (`Pool::Gpu(0)`) sized in GB.
    ///
    /// Pass `0.0` for either argument to disable that pool. These are
    /// *capacity* budgets — they govern how many bytes of model weights can
    /// be resident concurrently, not how fast inference will run.
    #[must_use]
    pub fn with_budgets_gb(cpu_ram_gb: f64, gpu_vram_gb: f64) -> Self {
        #[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
        let cpu_bytes = (cpu_ram_gb * 1_073_741_824.0) as u64;
        #[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
        let gpu_bytes = (gpu_vram_gb * 1_073_741_824.0) as u64;
        let mut budgets = HashMap::new();
        budgets.insert(Pool::Cpu, cpu_bytes);
        budgets.insert(Pool::Gpu(0), gpu_bytes);
        Self::new(budgets)
    }

    /// Register a local model for lifecycle and memory management under `id`.
    ///
    /// The pool is derived from `model.device()` at registration time. This
    /// is the lifecycle-only path: the entry can be loaded / unloaded /
    /// adapter-mounted and counts against its pool budget, but is **not**
    /// dispatchable via [`Self::complete`] unless the backend also
    /// implements [`Model`] and is registered through [`Self::register_provider`].
    ///
    /// For the common case of a local backend that is *both* a [`LocalModel`]
    /// and a chat [`Model`] (mistral.rs, llama.cpp, candle), prefer
    /// [`Self::register_provider`] so it can be dispatched by name.
    pub async fn register(&self, id: &str, model: Arc<dyn LocalModel>, memory_estimate_bytes: u64) {
        let pool: Pool = model.device().into();
        let mut state = self.state.lock().await;
        state.insert(
            id.to_owned(),
            RegisteredModel {
                model: None,
                local: Some(model),
                memory_estimate_bytes,
                pool,
                loaded: false,
                last_used: None,
                adapters: HashMap::new(),
            },
        );
    }

    /// Register a provider under `id` so it can be dispatched by name with
    /// [`Self::complete`] / [`Self::stream`] and fetched with [`Self::get`].
    ///
    /// This is the unified entry point that puts remote API providers and
    /// local in-process backends in one place:
    ///
    /// - **Remote provider** (`OpenAI`, Anthropic, fal.ai, …): pass the
    ///   [`Model`] and `local = None`. The entry is filed under
    ///   [`Pool::Remote`], holds no local weights, never participates in
    ///   eviction, and never counts against a memory budget — pass
    ///   `memory_estimate_bytes = 0`.
    /// - **Local backend that is also a chat model** (mistral.rs, llama.cpp,
    ///   candle): pass the same instance as both `model` and
    ///   `Some(local)` (concrete providers implement both traits, so the two
    ///   `Arc`s coerce from one value). The pool is derived from
    ///   `local.device()` and the entry participates in memory budgeting and
    ///   LRU eviction exactly like [`Self::register`].
    ///
    /// `complete`/`stream` auto-load a local entry on first use via
    /// [`Self::ensure_loaded`]; for remote entries that load is a no-op.
    pub async fn register_provider(
        &self,
        id: &str,
        model: Arc<dyn Model>,
        local: Option<Arc<dyn LocalModel>>,
        memory_estimate_bytes: u64,
    ) {
        let pool: Pool = match &local {
            Some(l) => l.device().into(),
            None => Pool::Remote,
        };
        let mut state = self.state.lock().await;
        state.insert(
            id.to_owned(),
            RegisteredModel {
                model: Some(model),
                local,
                memory_estimate_bytes,
                pool,
                loaded: false,
                last_used: None,
                adapters: HashMap::new(),
            },
        );
    }

    /// Fetch the chat interface registered under `id` to use or compose
    /// directly (e.g. wrap several in a
    /// [`FallbackModel`](blazen_llm::fallback::FallbackModel)).
    ///
    /// Returns `None` if `id` is not registered or was registered for
    /// lifecycle only (no chat [`Model`]).
    pub async fn get(&self, id: &str) -> Option<Arc<dyn Model>> {
        let state = self.state.lock().await;
        state.get(id).and_then(|e| e.model.clone())
    }

    /// Run a chat completion against the provider registered under `id`.
    ///
    /// Local entries are auto-loaded on first use (subject to the pool
    /// budget); remote entries dispatch straight through. The entry's
    /// last-used timestamp is bumped so LRU eviction sees the access.
    ///
    /// # Errors
    /// - [`BlazenError::validation`] if `id` is not registered.
    /// - [`BlazenError::unsupported`] if `id` was registered for lifecycle
    ///   only (no chat [`Model`] — see [`Self::register_provider`]).
    /// - Any error from [`Self::ensure_loaded`] (local budget / load) or the
    ///   provider's own [`Model::complete`].
    pub async fn complete(
        &self,
        id: &str,
        request: ModelRequest,
    ) -> Result<ModelResponse, BlazenError> {
        let model = self.dispatch_target(id).await?;
        self.ensure_loaded(id).await?;
        self.touch(id).await;
        model.complete(request).await
    }

    /// Streaming counterpart to [`Self::complete`].
    ///
    /// # Errors
    /// Same surface as [`Self::complete`], plus any error from the
    /// provider's own [`Model::stream`].
    pub async fn stream(
        &self,
        id: &str,
        request: ModelRequest,
    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, BlazenError>> + Send>>, BlazenError>
    {
        let model = self.dispatch_target(id).await?;
        self.ensure_loaded(id).await?;
        self.touch(id).await;
        model.stream(request).await
    }

    /// Resolve `id` to its chat dispatch target, erroring if it is unknown
    /// or lifecycle-only.
    async fn dispatch_target(&self, id: &str) -> Result<Arc<dyn Model>, BlazenError> {
        let state = self.state.lock().await;
        let entry = state
            .get(id)
            .ok_or_else(|| BlazenError::validation(format!("model '{id}' is not registered")))?;
        entry.model.clone().ok_or_else(|| {
            BlazenError::unsupported(format!(
                "model '{id}' is registered for lifecycle/memory management only and does not \
                 implement the chat Model interface; register it via register_provider with a \
                 Model to dispatch completions"
            ))
        })
    }

    /// Bump the last-used timestamp for LRU bookkeeping. Silent no-op if the
    /// entry vanished between dispatch and here.
    async fn touch(&self, id: &str) {
        let mut state = self.state.lock().await;
        if let Some(e) = state.get_mut(id) {
            e.last_used = Some(Instant::now());
        }
    }

    /// Load a model, evicting LRU models in the same pool if necessary.
    ///
    /// # Errors
    /// Returns [`BlazenError::Validation`] when the model is not registered,
    /// when its memory estimate exceeds its pool's total budget, or when
    /// no eviction can free enough space. Propagates errors from the
    /// underlying [`LocalModel::load`] / [`LocalModel::unload`] calls.
    ///
    /// # Panics
    /// Panics only if internal state becomes inconsistent (a model ID that
    /// was just verified to exist is no longer present). This cannot happen
    /// under normal operation.
    pub async fn load(&self, id: &str) -> Result<(), BlazenError> {
        let mut state = self.state.lock().await;

        let entry = state
            .get(id)
            .ok_or_else(|| BlazenError::validation(format!("model '{id}' is not registered")))?;

        if entry.loaded {
            let entry = state.get_mut(id).expect("checked above");
            entry.last_used = Some(Instant::now());
            return Ok(());
        }

        // Remote / dispatch-only entries own no local weights: there is
        // nothing to load, no budget to charge, and no eviction to run.
        // Mark them ready so `complete`/`stream` can treat both tiers
        // uniformly.
        if entry.local.is_none() {
            let entry = state.get_mut(id).expect("checked above");
            entry.loaded = true;
            entry.last_used = Some(Instant::now());
            return Ok(());
        }

        let needed = entry.memory_estimate_bytes;
        let pool = entry.pool;
        let budget = self.pool_budgets.get(&pool).copied().unwrap_or(0);

        if needed > budget {
            return Err(BlazenError::validation(format!(
                "model '{id}' requires {needed} bytes but pool {pool} budget is only {budget} bytes",
            )));
        }

        let mut used = Self::used_bytes_in_pool(&state, pool);
        while used + needed > budget {
            let lru_id = state
                .iter()
                .filter(|(k, v)| {
                    v.loaded && v.local.is_some() && v.pool == pool && k.as_str() != id
                })
                .min_by_key(|(_, v)| v.last_used)
                .map(|(k, _)| k.clone());

            let Some(lru_id) = lru_id else {
                return Err(BlazenError::validation(format!(
                    "cannot free enough memory to load model '{id}' in pool {pool} \
                     (need {needed}, used {used}, budget {budget})",
                )));
            };

            let lru_model = state
                .get(&lru_id)
                .expect("lru_id came from iteration")
                .local
                .clone()
                .expect("LRU candidates are filtered to local entries");
            drop(state);
            lru_model.unload().await?;
            state = self.state.lock().await;
            if let Some(e) = state.get_mut(&lru_id) {
                e.loaded = false;
                e.last_used = None;
            }
            used = Self::used_bytes_in_pool(&state, pool);
        }

        let model = state
            .get(id)
            .expect("checked at top")
            .local
            .clone()
            .expect("local presence checked above");
        drop(state);
        model.load().await?;
        let mut state = self.state.lock().await;
        if let Some(e) = state.get_mut(id) {
            e.loaded = true;
            e.last_used = Some(Instant::now());
        }
        Ok(())
    }

    /// Unload a model, freeing its budget within its pool.
    ///
    /// # Errors
    /// Returns [`BlazenError::Validation`] when the model is not registered.
    /// Propagates errors from the underlying [`LocalModel::unload`] call.
    pub async fn unload(&self, id: &str) -> Result<(), BlazenError> {
        let state = self.state.lock().await;
        let entry = state
            .get(id)
            .ok_or_else(|| BlazenError::validation(format!("model '{id}' is not registered")))?;
        if !entry.loaded {
            return Ok(());
        }
        // Remote / dispatch-only entries have no weights to free; just clear
        // the ready flag.
        let Some(model) = entry.local.clone() else {
            let mut state = state;
            if let Some(e) = state.get_mut(id) {
                e.loaded = false;
                e.last_used = None;
            }
            return Ok(());
        };
        drop(state);
        model.unload().await?;
        let mut state = self.state.lock().await;
        if let Some(e) = state.get_mut(id) {
            e.loaded = false;
            e.last_used = None;
        }
        Ok(())
    }

    /// Check if a model is currently loaded.
    ///
    /// Remote / dispatch-only entries report `true` — they own no local
    /// weights and are always ready to dispatch.
    pub async fn is_loaded(&self, id: &str) -> bool {
        let state = self.state.lock().await;
        state.get(id).is_some_and(|e| e.loaded || e.local.is_none())
    }

    /// Ensure a model is loaded. Equivalent to [`Self::load`].
    ///
    /// # Errors
    /// See [`Self::load`].
    pub async fn ensure_loaded(&self, id: &str) -> Result<(), BlazenError> {
        self.load(id).await
    }

    /// Total memory currently used by loaded models in the given pool.
    pub async fn used_bytes(&self, pool: Pool) -> u64 {
        let state = self.state.lock().await;
        Self::used_bytes_in_pool(&state, pool)
    }

    /// Available memory within the given pool's budget.
    pub async fn available_bytes(&self, pool: Pool) -> u64 {
        let used = self.used_bytes(pool).await;
        let budget = self.pool_budgets.get(&pool).copied().unwrap_or(0);
        budget.saturating_sub(used)
    }

    /// All configured pools and their budgets (in bytes).
    #[must_use]
    pub fn pools(&self) -> Vec<(Pool, u64)> {
        self.pool_budgets.iter().map(|(p, b)| (*p, *b)).collect()
    }

    /// Status of all registered models.
    pub async fn status(&self) -> Vec<ModelStatus> {
        let state = self.state.lock().await;
        state
            .iter()
            .map(|(id, entry)| ModelStatus {
                id: id.clone(),
                loaded: entry.loaded,
                memory_estimate_bytes: entry.memory_estimate_bytes,
                pool: entry.pool,
                adapters: entry.adapters.values().cloned().collect(),
            })
            .collect()
    }

    fn used_bytes_in_pool(state: &HashMap<String, RegisteredModel>, pool: Pool) -> u64 {
        state
            .values()
            .filter(|e| e.loaded && e.pool == pool)
            .map(|e| e.memory_estimate_bytes)
            .sum()
    }

    /// Mount a `PEFT`-format `LoRA` adapter on a previously-registered model.
    ///
    /// The base model is loaded automatically via [`Self::ensure_loaded`]
    /// before the adapter is mounted. Adapter file sizes are summed from
    /// `adapter_dir` and checked against the pool budget before forwarding
    /// to the backend; if the result would exceed the budget the call
    /// fails with [`BlazenError::validation`] and the pool is NOT auto-
    /// evicted (evicting another tenant's model to make room for this
    /// tenant's adapter is action-at-a-distance and breaks isolation).
    ///
    /// After the backend reports success the manager bumps the model's
    /// `memory_estimate_bytes` by the adapter's `memory_bytes` so future
    /// LRU decisions see the updated footprint.
    ///
    /// # Errors
    ///
    /// - [`BlazenError::validation`] if `model_id` is not registered, the
    ///   adapter directory is missing required files, the adapter id is
    ///   already mounted, or the pool budget would be exceeded.
    /// - [`BlazenError::unsupported`] if the backend cannot mount adapters.
    /// - Any error returned by [`LocalModel::load`] (during the
    ///   `ensure_loaded` pre-flight) or [`LocalModel::load_adapter`].
    pub async fn load_adapter(
        &self,
        model_id: &str,
        adapter_dir: &Path,
        options: AdapterOptions,
    ) -> Result<AdapterHandle, BlazenError> {
        self.ensure_loaded(model_id).await?;

        let adapter_id = options.adapter_id.clone();
        let scale = options.scale;
        let source_dir = adapter_dir.to_path_buf();

        let (model, pool, used, budget, on_disk_estimate) = {
            let state = self.state.lock().await;
            let entry = state.get(model_id).ok_or_else(|| {
                BlazenError::validation(format!("model '{model_id}' is not registered"))
            })?;
            if entry.adapters.contains_key(&adapter_id) {
                return Err(BlazenError::validation(format!(
                    "adapter '{adapter_id}' is already mounted on model '{model_id}'; \
                     unload it first to replace",
                )));
            }
            let pool = entry.pool;
            let budget = self.pool_budgets.get(&pool).copied().unwrap_or(0);
            let used = Self::used_bytes_in_pool(&state, pool);
            let model = entry.local.clone().ok_or_else(|| {
                BlazenError::unsupported(format!(
                    "model '{model_id}' is a remote provider and does not support LoRA adapters"
                ))
            })?;
            let on_disk = sum_adapter_bytes(adapter_dir)?;
            (model, pool, used, budget, on_disk)
        };

        if used + on_disk_estimate > budget {
            return Err(BlazenError::validation(format!(
                "mounting adapter '{adapter_id}' on '{model_id}' would exceed pool {pool} budget \
                 (need {on_disk_estimate} more, used {used}, budget {budget})",
            )));
        }

        let handle = model.load_adapter(adapter_dir, options).await?;

        let status = AdapterStatus {
            adapter_id: handle.adapter_id.clone(),
            scale,
            source_dir,
            memory_bytes: handle.memory_bytes,
        };
        let mut state = self.state.lock().await;
        if let Some(entry) = state.get_mut(model_id) {
            entry.memory_estimate_bytes = entry
                .memory_estimate_bytes
                .saturating_add(handle.memory_bytes);
            entry.adapters.insert(handle.adapter_id.clone(), status);
            entry.last_used = Some(Instant::now());
        }
        Ok(handle)
    }

    /// Remove an adapter previously mounted via [`Self::load_adapter`].
    ///
    /// # Errors
    ///
    /// - [`BlazenError::validation`] if `model_id` is not registered or
    ///   `adapter_id` is not currently mounted on it.
    /// - Any error returned by [`LocalModel::unload_adapter`].
    pub async fn unload_adapter(
        &self,
        model_id: &str,
        adapter_id: &str,
    ) -> Result<(), BlazenError> {
        let (model, status) = {
            let state = self.state.lock().await;
            let entry = state.get(model_id).ok_or_else(|| {
                BlazenError::validation(format!("model '{model_id}' is not registered"))
            })?;
            let status = entry.adapters.get(adapter_id).cloned().ok_or_else(|| {
                BlazenError::validation(format!(
                    "adapter '{adapter_id}' is not mounted on model '{model_id}'",
                ))
            })?;
            let model = entry.local.clone().ok_or_else(|| {
                BlazenError::unsupported(format!(
                    "model '{model_id}' is a remote provider and does not support LoRA adapters"
                ))
            })?;
            (model, status)
        };

        let handle = AdapterHandle {
            adapter_id: status.adapter_id.clone(),
            memory_bytes: status.memory_bytes,
            mount_strategy: blazen_llm::AdapterMountStrategy::Attached,
        };
        model.unload_adapter(&handle).await?;

        let mut state = self.state.lock().await;
        if let Some(entry) = state.get_mut(model_id) {
            entry.memory_estimate_bytes = entry
                .memory_estimate_bytes
                .saturating_sub(status.memory_bytes);
            entry.adapters.remove(adapter_id);
        }
        Ok(())
    }

    /// Auto-detect the right local-inference backend for the Hugging Face
    /// repo at `repo` and register a fresh provider under `id`.
    ///
    /// Hides backend choice from casual users; advanced users still
    /// construct providers directly and call [`Self::register`]. The chosen
    /// backend is returned so the caller can inspect or log it.
    ///
    /// See [`hf_loader::choose_backend`] for the selection rules. The repo
    /// metadata is probed once via [`hf_loader::detect_layout`] (one HTTP
    /// `GET /api/models/{repo}` — no weight downloads); selection plus
    /// memory estimation are then pure.
    ///
    /// # Errors
    ///
    /// - [`BlazenError::Unsupported`] when the backend the loader picked
    ///   is not compiled in (feature gate `mistralrs-provider`,
    ///   `candle-provider`, or `llamacpp-provider` — `live-models`
    ///   activates all three).
    /// - [`BlazenError::Validation`] when the repo has no recognisable
    ///   model weights, is PEFT-adapter-only (use [`Self::load_adapter`]
    ///   instead), or is otherwise unloadable as a base model.
    /// - Any error from [`hf_loader::detect_layout`] (network / parse) or
    ///   the chosen provider's own construction.
    #[cfg(feature = "hf-loader")]
    pub async fn load_from_hf(
        &self,
        id: String,
        repo: &str,
        options: HfLoadOptions,
    ) -> Result<BackendHint, BlazenError> {
        let layout = hf_loader::detect_layout(repo, &options).await?;
        let backend =
            hf_loader::choose_backend(&layout, options.backend_hint, options.gguf_file.as_deref())?;

        // Why: layout-derived sizes are only available when the Hub returned
        // them. The user override always wins; otherwise fall back to summing
        // siblings, then to layout.total_weight_bytes (rounded), then bail.
        let memory_bytes = options
            .memory_estimate_bytes
            .or_else(|| {
                hf_loader::estimate_backend_bytes(
                    backend,
                    &layout,
                    &layout.file_sizes,
                    options.gguf_file.as_deref(),
                )
            })
            .or_else(|| layout.total_weight_bytes.map(hf_loader::round_up_to_mb))
            .ok_or_else(|| {
                BlazenError::validation(format!(
                    "could not determine memory footprint for repo '{repo}'; \
                     pass HfLoadOptions::memory_estimate_bytes explicitly"
                ))
            })?;

        let model = build_provider(backend, repo, &options, &layout).await?;
        self.register(&id, model, memory_bytes).await;
        Ok(backend)
    }

    /// Train a `LoRA` adapter on the base model named by `config.base_model_repo`.
    ///
    /// Builds a fresh `VarMap` for the adapter params, resolves the device
    /// from `config.device` (defaults to `"cpu"`), drives the
    /// [`blazen_train::Trainer`] state machine through the full dataset,
    /// and returns the on-disk [`TrainedAdapter`] descriptor. The
    /// `adapter_dir` field on the result is immediately mountable via
    /// [`Self::load_adapter`] on any backend that supports `PEFT`-format
    /// `LoRA` adapters (mistralrs today; candle and llama.cpp pending the
    /// PR3 wave).
    ///
    /// `progress`, when `Some`, receives [`TrainingEvent`]s for every
    /// `Started` / `StepCompleted` / `CheckpointSaved` / `Finished`
    /// transition. Returning `Err(...)` from the callback aborts the run
    /// and surfaces as [`BlazenError::cancelled`].
    ///
    /// # Errors
    ///
    /// - [`BlazenError::Validation`] if the device string is unrecognised,
    ///   the chosen CUDA / Metal device cannot be opened, or the
    ///   [`TrainConfig`] fails [`blazen_train::Trainer::new`]'s validation
    ///   (zero `lora.rank` / zero `max_steps`).
    /// - [`BlazenError::Request`] (via [`BlazenError::internal`]) for any
    ///   other training-time failure: HF-Hub download, safetensors mmap,
    ///   forward/backward/optimizer error, checkpoint or adapter export
    ///   failure, dataset I/O error, or candle tensor error.
    /// - [`BlazenError::cancelled`] when the progress callback aborts.
    #[cfg(feature = "training")]
    pub async fn train_lora(
        &self,
        config: TrainConfig,
        dataset: Box<dyn TrainingDataset>,
        progress: Option<Arc<dyn TrainingProgress>>,
    ) -> Result<TrainedAdapter, BlazenError> {
        let device = parse_train_device(config.device.as_deref().unwrap_or("cpu"))?;
        let varmap = candle_nn::VarMap::new();

        let mut trainer = blazen_train::Trainer::new(config, varmap, device)?;
        if let Some(sink) = progress {
            trainer = trainer.with_progress(sink);
        }
        let adapter = trainer.run(dataset).await?;
        Ok(adapter)
    }

    /// Train a `LoRA` adapter via Direct Preference Optimization (DPO).
    ///
    /// Same shape as [`Self::train_lora`] but expects a
    /// [`PreferenceDataset`] of `(prompt, chosen, rejected)` triples and a
    /// frozen reference model (defaults to `config.core.base_model_repo`
    /// when `config.reference_model_repo` is `None`). The returned
    /// [`TrainedAdapter`] is a `PEFT`-format `LoRA`, mountable via
    /// [`Self::load_adapter`] on any backend that supports `PEFT` `LoRA`.
    ///
    /// # Errors
    ///
    /// - [`BlazenError::Validation`] for unknown device strings or any
    ///   [`blazen_train::Trainer::new_dpo`] config validation failure
    ///   (zero `lora.rank` / zero `core.max_steps`).
    /// - [`BlazenError::Request`] (via [`BlazenError::internal`]) for any
    ///   training-time failure: HF-Hub download, safetensors mmap, forward
    ///   / backward / optimizer error, checkpoint or adapter export
    ///   failure, dataset I/O error, or candle tensor error.
    /// - [`BlazenError::cancelled`] when the progress callback aborts.
    #[cfg(feature = "training")]
    pub async fn train_dpo(
        &self,
        config: DpoConfig,
        dataset: Arc<dyn PreferenceDataset>,
        progress: Option<Arc<dyn TrainingProgress>>,
    ) -> Result<TrainedAdapter, BlazenError> {
        let device = parse_train_device(config.core.device.as_deref().unwrap_or("cpu"))?;
        let varmap = candle_nn::VarMap::new();

        let mut trainer = blazen_train::Trainer::new_dpo(config, varmap, device, progress)?;
        trainer.load_models_dpo().await?;
        let adapter = trainer.run_dpo(dataset).await?;
        Ok(adapter)
    }

    /// Train a `LoRA` adapter via Odds Ratio Preference Optimization (ORPO).
    ///
    /// Reference-free — combines a standard SFT loss on chosen completions
    /// with an odds-ratio preference term weighted by `config.lambda`.
    /// Consumes the same [`PreferenceDataset`] schema as
    /// [`Self::train_dpo`].
    ///
    /// # Errors
    ///
    /// Same surface as [`Self::train_dpo`].
    #[cfg(feature = "training")]
    pub async fn train_orpo(
        &self,
        config: OrpoConfig,
        dataset: Arc<dyn PreferenceDataset>,
        progress: Option<Arc<dyn TrainingProgress>>,
    ) -> Result<TrainedAdapter, BlazenError> {
        let device = parse_train_device(config.core.device.as_deref().unwrap_or("cpu"))?;
        let varmap = candle_nn::VarMap::new();

        let mut trainer = blazen_train::Trainer::new_orpo(config, varmap, device, progress)?;
        trainer.load_models_orpo().await?;
        let adapter = trainer.run_orpo(dataset).await?;
        Ok(adapter)
    }

    /// Train a `LoRA` adapter via Simple Preference Optimization (`SimPO`).
    ///
    /// Reference-free and length-normalized. `config.beta` scales the
    /// preference logits and `config.gamma` sets the target reward margin.
    /// Consumes the same [`PreferenceDataset`] schema as
    /// [`Self::train_dpo`].
    ///
    /// # Errors
    ///
    /// Same surface as [`Self::train_dpo`].
    #[cfg(feature = "training")]
    pub async fn train_simpo(
        &self,
        config: SimpoConfig,
        dataset: Arc<dyn PreferenceDataset>,
        progress: Option<Arc<dyn TrainingProgress>>,
    ) -> Result<TrainedAdapter, BlazenError> {
        let device = parse_train_device(config.core.device.as_deref().unwrap_or("cpu"))?;
        let varmap = candle_nn::VarMap::new();

        let mut trainer = blazen_train::Trainer::new_simpo(config, varmap, device, progress)?;
        trainer.load_models_simpo().await?;
        let adapter = trainer.run_simpo(dataset).await?;
        Ok(adapter)
    }

    /// Train a `LoRA` adapter via Kahneman-Tversky Optimization (KTO).
    ///
    /// Like DPO, KTO requires a frozen reference model (defaults to
    /// `config.core.base_model_repo`) — but the dataset schema differs:
    /// each row is a `(prompt, completion, desirable)` triple
    /// ([`RatedDataset`]), not a chosen/rejected pair.
    ///
    /// # Errors
    ///
    /// Same surface as [`Self::train_dpo`].
    #[cfg(feature = "training")]
    pub async fn train_kto(
        &self,
        config: KtoConfig,
        dataset: Arc<dyn RatedDataset>,
        progress: Option<Arc<dyn TrainingProgress>>,
    ) -> Result<TrainedAdapter, BlazenError> {
        let device = parse_train_device(config.core.device.as_deref().unwrap_or("cpu"))?;
        let varmap = candle_nn::VarMap::new();

        let mut trainer = blazen_train::Trainer::new_kto(config, varmap, device, progress)?;
        trainer.load_models_kto().await?;
        let adapter = trainer.run_kto(dataset).await?;
        Ok(adapter)
    }

    /// Run a full fine-tune (every parameter trainable; no `LoRA` adapter).
    ///
    /// Returns [`FullFineTuneResult`] — not [`TrainedAdapter`] — because the
    /// output is a complete set of model weights in `config.core.output_dir`
    /// rather than a mountable PEFT delta. Setting
    /// `config.gradient_checkpointing = true` is rejected up-front with
    /// [`BlazenError::Validation`] (mapped from
    /// [`BlazenTrainError::Unsupported`]) because candle 0.10.2 has no
    /// activation-checkpointing primitive.
    ///
    /// # Errors
    ///
    /// - [`BlazenError::Validation`] for unknown device strings,
    ///   `gradient_checkpointing = true`, or any
    ///   [`blazen_train::Trainer::new_full_finetune`] validation failure.
    /// - [`BlazenError::Request`] (via [`BlazenError::internal`]) for any
    ///   training-time failure: HF-Hub download, safetensors mmap, OOM-
    ///   bounded `>1B`-param rejection, forward / backward / optimizer
    ///   error, checkpoint or full-weights export failure, dataset I/O
    ///   error, or candle tensor error.
    /// - [`BlazenError::cancelled`] when the progress callback aborts.
    #[cfg(feature = "training")]
    pub async fn fine_tune(
        &self,
        config: FullFineTuneConfig,
        dataset: Arc<dyn TrainingDataset>,
        progress: Option<Arc<dyn TrainingProgress>>,
    ) -> Result<FullFineTuneResult, BlazenError> {
        let device = parse_train_device(config.core.device.as_deref().unwrap_or("cpu"))?;
        let varmap = candle_nn::VarMap::new();

        let mut trainer =
            blazen_train::Trainer::new_full_finetune(config, varmap, device, progress)?;
        trainer.load_models_full_finetune().await?;
        let result = trainer.run_full_finetune(dataset).await?;
        Ok(result)
    }

    /// Run Group Relative Policy Optimization (GRPO) against a pre-built
    /// [`GrpoTrainer`] and a stream of pre-sampled, pre-scored batches.
    ///
    /// Unlike [`Self::train_lora`] / [`Self::train_dpo`] / etc., this verb
    /// does *not* construct the policy / reference models — GRPO batches
    /// require K sampled completions per prompt and per-row rewards from
    /// a frozen reward model, and the policy + reference + reward-model
    /// loading path for GRPO is application-specific (different reward
    /// models, different samplers). Callers build the [`GrpoTrainer`] via
    /// `GrpoTrainer::new(cfg, varmap, device, policy, reference, ..)` and
    /// hand it in. The verb then owns the driving loop, the optional
    /// `AllReduce` wrapping, progress emission, and PEFT adapter export
    /// at the end.
    ///
    /// When `distributed` is `Some`, the inner trainer is wrapped in
    /// [`blazen_train::distributed::AllReduceGrpoTrainer`] and every
    /// `step` participates in a ring-AllReduce gradient average across
    /// `distributed.peers` before the optimizer kick. The local gRPC
    /// `AllReduce` server is bound on `0.0.0.0:<distributed.master_port>`
    /// (this rank's slot in the peer list); the `next` peer's endpoint
    /// is dialed via tonic. When `None`, the inner sync `step` is driven
    /// directly with no networking.
    ///
    /// # Errors
    ///
    /// - [`BlazenError::Validation`] if `distributed` is `Some` and
    ///   `RingTopology::new` rejects the rank / peers combination.
    /// - [`BlazenError::Request`] (via [`BlazenError::internal`]) for any
    ///   transport bootstrap failure (server bind, peer connect),
    ///   training-time failure (forward / backward / optimizer /
    ///   `AllReduce`), dataset I/O error, or adapter-export failure.
    /// - [`BlazenError::cancelled`] when the progress callback aborts.
    #[cfg(feature = "training")]
    pub async fn train_grpo(
        &self,
        trainer: blazen_train::GrpoTrainer,
        mut dataset: Box<dyn blazen_train::GrpoDataset>,
        progress: Option<Arc<dyn TrainingProgress>>,
        #[cfg(feature = "distributed")] distributed: Option<
            blazen_train::config::DistributedConfig,
        >,
        #[cfg(not(feature = "distributed"))] distributed: Option<()>,
    ) -> Result<TrainedAdapter, BlazenError> {
        let _ = progress; // GrpoTrainer captures its progress sink at construction time.
        let output_dir = trainer.output_dir();
        let base_repo = trainer.config().base_model_repo.clone();
        let lora_cfg = trainer.config().lora.clone();

        #[cfg(feature = "distributed")]
        let (final_loss, total_steps, finished) = if let Some(dist_cfg) = distributed {
            run_grpo_distributed(trainer, &mut *dataset, dist_cfg).await?
        } else {
            run_grpo_local(trainer, &mut *dataset).await?
        };
        #[cfg(not(feature = "distributed"))]
        let (final_loss, total_steps, finished) = {
            let _ = distributed; // unused when the `distributed` feature is off
            run_grpo_local(trainer, &mut *dataset).await?
        };

        blazen_train::export::save_peft_adapter(
            finished.varmap(),
            &output_dir,
            &lora_cfg,
            &base_repo,
        )
        .map_err(BlazenError::from)?;

        Ok(TrainedAdapter {
            adapter_dir: output_dir,
            final_loss,
            total_steps,
        })
    }

    /// Run Proximal Policy Optimization (PPO) against a pre-built
    /// [`PpoTrainer`] and a stream of pre-rolled batches.
    ///
    /// Mirrors [`Self::train_grpo`]: the caller builds the trainer
    /// (policy + critic + reference) because PPO rollouts are
    /// application-specific (different sampler, different reward model,
    /// different KL configuration). The verb owns the loop, the optional
    /// distributed wrapping via
    /// [`blazen_train::distributed::AllReducePpoTrainer`], progress
    /// emission, and PEFT adapter export.
    ///
    /// # Errors
    ///
    /// Same surface as [`Self::train_grpo`].
    #[cfg(feature = "training")]
    pub async fn train_ppo(
        &self,
        trainer: blazen_train::PpoTrainer,
        mut dataset: Box<dyn blazen_train::PpoDataset>,
        progress: Option<Arc<dyn TrainingProgress>>,
        #[cfg(feature = "distributed")] distributed: Option<
            blazen_train::config::DistributedConfig,
        >,
        #[cfg(not(feature = "distributed"))] distributed: Option<()>,
    ) -> Result<TrainedAdapter, BlazenError> {
        let _ = progress; // PpoTrainer captures its progress sink at construction time.
        let output_dir = trainer.config().output_dir.clone();
        let base_repo = trainer.config().base_model_repo.clone();
        let lora_cfg = trainer.config().lora.clone();

        #[cfg(feature = "distributed")]
        let (final_loss, total_steps, finished) = if let Some(dist_cfg) = distributed {
            run_ppo_distributed(trainer, &mut *dataset, dist_cfg).await?
        } else {
            run_ppo_local(trainer, &mut *dataset).await?
        };
        #[cfg(not(feature = "distributed"))]
        let (final_loss, total_steps, finished) = {
            let _ = distributed;
            run_ppo_local(trainer, &mut *dataset).await?
        };

        blazen_train::export::save_peft_adapter(
            finished.varmap(),
            &output_dir,
            &lora_cfg,
            &base_repo,
        )
        .map_err(BlazenError::from)?;

        Ok(TrainedAdapter {
            adapter_dir: output_dir,
            final_loss,
            total_steps,
        })
    }

    /// List adapters mounted on a registered model.
    ///
    /// # Errors
    /// [`BlazenError::validation`] if `model_id` is not registered.
    pub async fn list_adapters(&self, model_id: &str) -> Result<Vec<AdapterStatus>, BlazenError> {
        let state = self.state.lock().await;
        let entry = state.get(model_id).ok_or_else(|| {
            BlazenError::validation(format!("model '{model_id}' is not registered"))
        })?;
        Ok(entry.adapters.values().cloned().collect())
    }
}

/// Construct an `Arc<dyn LocalModel>` for the chosen backend, mapping
/// [`HfLoadOptions`] onto the backend's own options struct.
///
/// Each backend is independently feature-gated; calling [`ModelManager::load_from_hf`]
/// for a backend whose feature is off returns [`BlazenError::unsupported`]
/// naming the missing feature.
#[cfg(feature = "hf-loader")]
#[allow(unused_variables)] // Why: `layout` is consumed only by the candle/llamacpp arms.
#[allow(clippy::unused_async)] // Why: only the llamacpp arm awaits; the other arms are sync or disabled-feature error branches.
async fn build_provider(
    backend: BackendHint,
    repo: &str,
    options: &HfLoadOptions,
    layout: &hf_loader::DetectedLayout,
) -> Result<Arc<dyn LocalModel>, BlazenError> {
    match backend {
        BackendHint::Mistralrs => {
            #[cfg(feature = "mistralrs-provider")]
            {
                let mr_opts = blazen_llm_mistralrs::MistralRsOptions {
                    base: blazen_local_llm::LocalLlmOptions {
                        model_id: Some(repo.to_string()),
                        device: options.device.clone(),
                        cache_dir: options.cache_dir.clone(),
                        ..blazen_local_llm::LocalLlmOptions::default()
                    },
                    max_batch_size: None,
                    chat_template: None,
                    vision: false,
                };
                let p = blazen_llm_mistralrs::MistralRsProvider::from_options(mr_opts)
                    .map_err(|e| BlazenError::internal(format!("mistralrs build: {e}")))?;
                Ok(Arc::new(p))
            }
            #[cfg(not(feature = "mistralrs-provider"))]
            {
                Err(BlazenError::unsupported(
                    "load_from_hf chose the mistralrs backend but the \
                     `mistralrs-provider` feature is not enabled on blazen-manager",
                ))
            }
        }
        BackendHint::Candle => {
            #[cfg(feature = "candle-provider")]
            {
                let force_st =
                    !layout.gguf_files.is_empty() && !layout.safetensors_files.is_empty();
                let c_opts = blazen_llm_candle::CandleLlmOptions {
                    base: blazen_local_llm::LocalLlmOptions {
                        model_id: Some(repo.to_string()),
                        device: options.device.clone(),
                        revision: options.revision.clone(),
                        cache_dir: options.cache_dir.clone(),
                        ..blazen_local_llm::LocalLlmOptions::default()
                    },
                    force_safetensors: force_st,
                };
                let p = blazen_llm_candle::CandleLlmProvider::from_options(c_opts)
                    .map_err(|e| BlazenError::internal(format!("candle build: {e}")))?;
                Ok(Arc::new(p))
            }
            #[cfg(not(feature = "candle-provider"))]
            {
                Err(BlazenError::unsupported(
                    "load_from_hf chose the candle backend but the \
                     `candle-provider` feature is not enabled on blazen-manager",
                ))
            }
        }
        BackendHint::Llamacpp => {
            #[cfg(feature = "llamacpp-provider")]
            {
                // Why: llama.cpp's LlamaCppOptions::model_path accepts either a
                // local path or the HF triple "org/repo/file.gguf". We must
                // pick a specific GGUF — either the caller-supplied gguf_file
                // or the first one in the repo — because the provider rejects
                // bare "org/repo" without a filename.
                let gguf = options
                    .gguf_file
                    .clone()
                    .or_else(|| layout.gguf_files.first().cloned())
                    .ok_or_else(|| {
                        BlazenError::validation(format!(
                            "llamacpp backend requires a *.gguf file but repo '{repo}' \
                             has none; pass HfLoadOptions::gguf_file or use a different backend"
                        ))
                    })?;
                let model_path = format!("{repo}/{gguf}");
                let l_opts = blazen_llm_llamacpp::LlamaCppOptions {
                    base: blazen_local_llm::LocalLlmOptions {
                        model_id: Some(model_path),
                        device: options.device.clone(),
                        cache_dir: options.cache_dir.clone(),
                        ..blazen_local_llm::LocalLlmOptions::default()
                    },
                    n_gpu_layers: None,
                };
                let p = blazen_llm_llamacpp::LlamaCppProvider::from_options(l_opts)
                    .await
                    .map_err(|e| BlazenError::internal(format!("llamacpp build: {e}")))?;
                Ok(Arc::new(p))
            }
            #[cfg(not(feature = "llamacpp-provider"))]
            {
                Err(BlazenError::unsupported(
                    "load_from_hf chose the llamacpp backend but the \
                     `llamacpp-provider` feature is not enabled on blazen-manager",
                ))
            }
        }
    }
}

/// Parse a `TrainConfig::device` string into a `candle_core::Device`.
///
/// Recognised forms: `"cpu"`, `"cuda"` (= `cuda:0`), `"cuda:N"`, `"metal"`
/// (= `metal:0`), `"metal:N"`. Anything else returns
/// [`BlazenError::validation`].
#[cfg(feature = "training")]
fn parse_train_device(spec: &str) -> Result<candle_core::Device, BlazenError> {
    let normalized = spec.trim().to_ascii_lowercase();
    if normalized == "cpu" {
        return Ok(candle_core::Device::Cpu);
    }
    let (kind, idx) = match normalized.split_once(':') {
        Some((k, rest)) => {
            let parsed = rest.parse::<usize>().map_err(|e| {
                BlazenError::validation(format!(
                    "training device '{spec}' has non-numeric index '{rest}': {e}"
                ))
            })?;
            (k, parsed)
        }
        None => (normalized.as_str(), 0),
    };
    match kind {
        "cuda" => candle_core::Device::new_cuda(idx)
            .map_err(|e| BlazenError::validation(format!("cuda:{idx} unavailable: {e}"))),
        "metal" => candle_core::Device::new_metal(idx)
            .map_err(|e| BlazenError::validation(format!("metal:{idx} unavailable: {e}"))),
        other => Err(BlazenError::validation(format!(
            "unknown training device '{other}' (want one of: cpu, cuda[:N], metal[:N])"
        ))),
    }
}

/// Sum the on-disk sizes of the canonical PEFT-layout files in
/// `adapter_dir`: `adapter_model.safetensors` (required) +
/// `adapter_config.json` (required) + `tokenizer.json` (optional). Used as
/// a conservative pre-flight estimate before calling
/// [`LocalModel::load_adapter`]. The backend reports the true runtime
/// footprint after the mount succeeds.
fn sum_adapter_bytes(adapter_dir: &Path) -> Result<u64, BlazenError> {
    let required = ["adapter_model.safetensors", "adapter_config.json"];
    let optional = ["tokenizer.json"];

    let mut total: u64 = 0;
    for name in required {
        let path: PathBuf = adapter_dir.join(name);
        let meta = std::fs::metadata(&path).map_err(|e| {
            BlazenError::validation(format!(
                "adapter directory missing required file '{name}' at {}: {e}",
                path.display(),
            ))
        })?;
        total = total.saturating_add(meta.len());
    }
    for name in optional {
        let path = adapter_dir.join(name);
        if let Ok(meta) = std::fs::metadata(&path) {
            total = total.saturating_add(meta.len());
        }
    }
    Ok(total)
}

/// Drive a [`blazen_train::GrpoTrainer`] to dataset exhaustion on a single
/// device, returning `(final_loss, total_steps, trainer)`.
#[cfg(feature = "training")]
async fn run_grpo_local(
    mut trainer: blazen_train::GrpoTrainer,
    dataset: &mut dyn blazen_train::GrpoDataset,
) -> Result<(f32, usize, blazen_train::GrpoTrainer), BlazenError> {
    let mut final_loss: f32 = 0.0;
    while let Some(batch) = dataset.next_batch().await? {
        final_loss = trainer.step(&batch)?;
    }
    let total_steps = trainer.global_step();
    Ok((final_loss, total_steps, trainer))
}

/// Drive a [`blazen_train::GrpoTrainer`] under ring-AllReduce gradient
/// averaging across `distributed.peers`. Bootstraps the gRPC `AllReduce`
/// server on `0.0.0.0:<dist_cfg.master_port>` (this rank's slot in the
/// peer list) and dials the `next` peer's endpoint before the first
/// step. Returns `(final_loss, total_steps, unwrapped_trainer)`.
#[cfg(all(feature = "training", feature = "distributed"))]
async fn run_grpo_distributed(
    trainer: blazen_train::GrpoTrainer,
    dataset: &mut dyn blazen_train::GrpoDataset,
    dist_cfg: blazen_train::config::DistributedConfig,
) -> Result<(f32, usize, blazen_train::GrpoTrainer), BlazenError> {
    use blazen_train::distributed::{
        AllReduceGrpoTrainer, RingTopology, transport::grpc::GrpcRingTransport,
    };

    let topology = RingTopology::new(dist_cfg.rank, dist_cfg.peers.clone())
        .map_err(|e| BlazenError::validation(format!("invalid distributed config: {e}")))?;
    let bind_addr: std::net::SocketAddr = format!("0.0.0.0:{}", dist_cfg.master_port)
        .parse()
        .map_err(|e| {
            BlazenError::validation(format!(
                "invalid master_port {} for AllReduce bind: {e}",
                dist_cfg.master_port
            ))
        })?;
    let (transport, _server_handle) = GrpcRingTransport::bootstrap(&topology, bind_addr)
        .await
        .map_err(BlazenError::from)?;
    let transport: Arc<dyn blazen_train::distributed::RingTransport> = Arc::new(transport);

    let mut wrapped = AllReduceGrpoTrainer::new(trainer, topology, transport);
    let mut final_loss: f32 = 0.0;
    while let Some(batch) = dataset.next_batch().await? {
        final_loss = wrapped.step(&batch).await?;
    }
    let total_steps = wrapped.inner().global_step();
    let inner = wrapped.into_inner();
    Ok((final_loss, total_steps, inner))
}

/// Drive a [`blazen_train::PpoTrainer`] to dataset exhaustion on a single
/// device, returning `(final_loss, total_steps, trainer)`.
#[cfg(feature = "training")]
async fn run_ppo_local(
    mut trainer: blazen_train::PpoTrainer,
    dataset: &mut dyn blazen_train::PpoDataset,
) -> Result<(f32, usize, blazen_train::PpoTrainer), BlazenError> {
    let mut final_loss: f32 = 0.0;
    while let Some(batch) = dataset.next_batch().await? {
        final_loss = trainer.step(&batch)?;
    }
    let total_steps = trainer.global_step();
    Ok((final_loss, total_steps, trainer))
}

/// Drive a [`blazen_train::PpoTrainer`] under ring-AllReduce gradient
/// averaging across `distributed.peers`. See [`run_grpo_distributed`]
/// for bootstrap details.
#[cfg(all(feature = "training", feature = "distributed"))]
async fn run_ppo_distributed(
    trainer: blazen_train::PpoTrainer,
    dataset: &mut dyn blazen_train::PpoDataset,
    dist_cfg: blazen_train::config::DistributedConfig,
) -> Result<(f32, usize, blazen_train::PpoTrainer), BlazenError> {
    use blazen_train::distributed::{
        AllReducePpoTrainer, RingTopology, transport::grpc::GrpcRingTransport,
    };

    let topology = RingTopology::new(dist_cfg.rank, dist_cfg.peers.clone())
        .map_err(|e| BlazenError::validation(format!("invalid distributed config: {e}")))?;
    let bind_addr: std::net::SocketAddr = format!("0.0.0.0:{}", dist_cfg.master_port)
        .parse()
        .map_err(|e| {
            BlazenError::validation(format!(
                "invalid master_port {} for AllReduce bind: {e}",
                dist_cfg.master_port
            ))
        })?;
    let (transport, _server_handle) = GrpcRingTransport::bootstrap(&topology, bind_addr)
        .await
        .map_err(BlazenError::from)?;
    let transport: Arc<dyn blazen_train::distributed::RingTransport> = Arc::new(transport);

    let mut wrapped = AllReducePpoTrainer::new(trainer, topology, transport);
    let mut final_loss: f32 = 0.0;
    while let Some(batch) = dataset.next_batch().await? {
        final_loss = wrapped.step(&batch).await?;
    }
    let total_steps = wrapped.inner().global_step();
    let inner = wrapped.into_inner();
    Ok((final_loss, total_steps, inner))
}

#[cfg(test)]
mod tests {
    use super::*;
    use blazen_llm::{AdapterMountStrategy, Device};
    use std::sync::Mutex as StdMutex;

    const GB: u64 = 1_073_741_824;

    struct MockLocalModel {
        loaded: StdMutex<bool>,
        device: Device,
        adapters: StdMutex<HashMap<String, AdapterStatus>>,
        next_adapter_memory_bytes: StdMutex<u64>,
    }

    impl MockLocalModel {
        fn new(device: Device) -> Self {
            Self {
                loaded: StdMutex::new(false),
                device,
                adapters: StdMutex::new(HashMap::new()),
                next_adapter_memory_bytes: StdMutex::new(1024),
            }
        }

        fn with_adapter_memory(self, bytes: u64) -> Self {
            *self.next_adapter_memory_bytes.lock().unwrap() = bytes;
            self
        }
    }

    #[async_trait::async_trait]
    impl LocalModel for MockLocalModel {
        async fn load(&self) -> Result<(), BlazenError> {
            *self.loaded.lock().unwrap() = true;
            Ok(())
        }
        async fn unload(&self) -> Result<(), BlazenError> {
            *self.loaded.lock().unwrap() = false;
            Ok(())
        }
        async fn is_loaded(&self) -> bool {
            *self.loaded.lock().unwrap()
        }
        fn device(&self) -> Device {
            self.device.clone()
        }
        async fn load_adapter(
            &self,
            adapter_dir: &std::path::Path,
            options: AdapterOptions,
        ) -> Result<AdapterHandle, BlazenError> {
            let memory_bytes = *self.next_adapter_memory_bytes.lock().unwrap();
            let status = AdapterStatus {
                adapter_id: options.adapter_id.clone(),
                scale: options.scale,
                source_dir: adapter_dir.to_path_buf(),
                memory_bytes,
            };
            self.adapters
                .lock()
                .unwrap()
                .insert(options.adapter_id.clone(), status);
            Ok(AdapterHandle {
                adapter_id: options.adapter_id,
                memory_bytes,
                mount_strategy: AdapterMountStrategy::Attached,
            })
        }
        async fn unload_adapter(&self, handle: &AdapterHandle) -> Result<(), BlazenError> {
            let mut adapters = self.adapters.lock().unwrap();
            if adapters.remove(&handle.adapter_id).is_none() {
                return Err(BlazenError::validation(format!(
                    "mock backend has no adapter '{}'",
                    handle.adapter_id
                )));
            }
            Ok(())
        }
        async fn list_adapters(&self) -> Vec<AdapterStatus> {
            self.adapters.lock().unwrap().values().cloned().collect()
        }
    }

    /// Build a minimal canned [`ModelResponse`] for mock providers.
    fn mock_response(model: &str, content: &str) -> ModelResponse {
        ModelResponse {
            content: Some(content.to_owned()),
            tool_calls: Vec::new(),
            reasoning: None,
            citations: Vec::new(),
            artifacts: Vec::new(),
            usage: None,
            model: model.to_owned(),
            finish_reason: Some("stop".to_owned()),
            cost: None,
            timing: None,
            images: Vec::new(),
            audio: Vec::new(),
            videos: Vec::new(),
            metadata: serde_json::Value::Null,
        }
    }

    /// A remote-style provider: implements [`Model`] only (no local weights).
    struct MockRemoteModel {
        id: String,
    }

    impl MockRemoteModel {
        fn new(id: &str) -> Self {
            Self { id: id.to_owned() }
        }
    }

    #[async_trait::async_trait]
    impl Model for MockRemoteModel {
        fn model_id(&self) -> &str {
            &self.id
        }
        async fn complete(&self, _request: ModelRequest) -> Result<ModelResponse, BlazenError> {
            Ok(mock_response(&self.id, &format!("echo from {}", self.id)))
        }
        async fn stream(
            &self,
            _request: ModelRequest,
        ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, BlazenError>> + Send>>, BlazenError>
        {
            let chunk = StreamChunk {
                delta: Some(format!("echo from {}", self.id)),
                tool_calls: Vec::new(),
                finish_reason: Some("stop".to_owned()),
                reasoning_delta: None,
                citations: Vec::new(),
                artifacts: Vec::new(),
            };
            Ok(Box::pin(futures_util::stream::iter(vec![Ok(chunk)])))
        }
    }

    /// A local backend that is *both* a [`Model`] and a [`LocalModel`] — the
    /// mistral.rs / llama.cpp shape. Used to prove a local entry registered
    /// via [`ModelManager::register_provider`] is both dispatchable and
    /// budget-managed.
    struct MockLocalLlm {
        id: String,
        loaded: StdMutex<bool>,
        device: Device,
    }

    impl MockLocalLlm {
        fn new(id: &str, device: Device) -> Self {
            Self {
                id: id.to_owned(),
                loaded: StdMutex::new(false),
                device,
            }
        }
    }

    #[async_trait::async_trait]
    impl Model for MockLocalLlm {
        fn model_id(&self) -> &str {
            &self.id
        }
        async fn complete(&self, _request: ModelRequest) -> Result<ModelResponse, BlazenError> {
            if !*self.loaded.lock().unwrap() {
                return Err(BlazenError::validation(format!(
                    "mock local llm '{}' completed before being loaded",
                    self.id
                )));
            }
            Ok(mock_response(&self.id, &format!("local {}", self.id)))
        }
        async fn stream(
            &self,
            _request: ModelRequest,
        ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, BlazenError>> + Send>>, BlazenError>
        {
            Ok(Box::pin(futures_util::stream::iter(Vec::new())))
        }
    }

    #[async_trait::async_trait]
    impl LocalModel for MockLocalLlm {
        async fn load(&self) -> Result<(), BlazenError> {
            *self.loaded.lock().unwrap() = true;
            Ok(())
        }
        async fn unload(&self) -> Result<(), BlazenError> {
            *self.loaded.lock().unwrap() = false;
            Ok(())
        }
        async fn is_loaded(&self) -> bool {
            *self.loaded.lock().unwrap()
        }
        fn device(&self) -> Device {
            self.device.clone()
        }
    }

    fn make_peft_adapter_dir(bytes_a: usize, bytes_b: usize) -> tempfile::TempDir {
        let dir = tempfile::tempdir().expect("tempdir");
        std::fs::write(
            dir.path().join("adapter_model.safetensors"),
            vec![0u8; bytes_a],
        )
        .expect("write safetensors");
        std::fs::write(dir.path().join("adapter_config.json"), vec![0u8; bytes_b])
            .expect("write config");
        dir
    }

    fn cpu_gpu_mgr(cpu_gb: u64, gpu_gb: u64) -> ModelManager {
        let mut budgets = HashMap::new();
        budgets.insert(Pool::Cpu, cpu_gb * GB);
        budgets.insert(Pool::Gpu(0), gpu_gb * GB);
        ModelManager::new(budgets)
    }

    #[tokio::test]
    async fn test_register_and_load_cpu() {
        let mgr = cpu_gpu_mgr(24, 24);
        let model = Arc::new(MockLocalModel::new(Device::Cpu));

        mgr.register("m1", model.clone(), 4 * GB).await;
        mgr.load("m1").await.unwrap();

        assert!(mgr.is_loaded("m1").await);
        assert_eq!(mgr.used_bytes(Pool::Cpu).await, 4 * GB);
        assert_eq!(mgr.used_bytes(Pool::Gpu(0)).await, 0);
    }

    #[tokio::test]
    async fn test_register_and_load_gpu() {
        let mgr = cpu_gpu_mgr(24, 24);
        let model = Arc::new(MockLocalModel::new(Device::Cuda(0)));

        mgr.register("g1", model.clone(), 8 * GB).await;
        mgr.load("g1").await.unwrap();

        assert!(mgr.is_loaded("g1").await);
        assert_eq!(mgr.used_bytes(Pool::Gpu(0)).await, 8 * GB);
        assert_eq!(mgr.used_bytes(Pool::Cpu).await, 0);
    }

    #[tokio::test]
    async fn test_lru_eviction_within_pool() {
        let mgr = cpu_gpu_mgr(20, 0);
        let m1 = Arc::new(MockLocalModel::new(Device::Cpu));
        let m2 = Arc::new(MockLocalModel::new(Device::Cpu));
        let m3 = Arc::new(MockLocalModel::new(Device::Cpu));

        mgr.register("m1", m1.clone(), 4 * GB).await;
        mgr.register("m2", m2.clone(), 8 * GB).await;
        mgr.register("m3", m3.clone(), 12 * GB).await;

        mgr.load("m1").await.unwrap();
        mgr.load("m2").await.unwrap();
        assert_eq!(mgr.used_bytes(Pool::Cpu).await, 12 * GB);

        // Loading m3 (12 GB) on top of 12 GB used in a 20 GB pool needs to
        // evict m1 (oldest LRU, 4 GB) — leaves 8 GB used, 12 GB free, fits.
        mgr.load("m3").await.unwrap();

        assert!(
            !mgr.is_loaded("m1").await,
            "m1 should have been evicted (oldest LRU)"
        );
        assert!(mgr.is_loaded("m2").await, "m2 should still be loaded");
        assert!(mgr.is_loaded("m3").await, "m3 should now be loaded");
    }

    #[tokio::test]
    async fn test_cross_pool_no_eviction() {
        // CPU pool: 16 GB, GPU pool: 16 GB.
        let mgr = cpu_gpu_mgr(16, 16);
        let cpu = Arc::new(MockLocalModel::new(Device::Cpu));
        let gpu_a = Arc::new(MockLocalModel::new(Device::Cuda(0)));
        let gpu_b = Arc::new(MockLocalModel::new(Device::Cuda(0)));

        // Fill each pool to its max.
        mgr.register("cpu", cpu.clone(), 16 * GB).await;
        mgr.register("gpu_a", gpu_a.clone(), 16 * GB).await;
        mgr.register("gpu_b", gpu_b.clone(), 16 * GB).await;

        mgr.load("cpu").await.unwrap();
        mgr.load("gpu_a").await.unwrap();

        // Loading another GPU-sized model must evict gpu_a, never cpu —
        // even though cpu is older.
        mgr.load("gpu_b").await.unwrap();

        assert!(
            mgr.is_loaded("cpu").await,
            "CPU model must NOT be evicted when a GPU model is loaded"
        );
        assert!(
            !mgr.is_loaded("gpu_a").await,
            "GPU LRU (gpu_a) should have been evicted"
        );
        assert!(mgr.is_loaded("gpu_b").await, "gpu_b should be loaded");
        assert_eq!(mgr.used_bytes(Pool::Cpu).await, 16 * GB);
        assert_eq!(mgr.used_bytes(Pool::Gpu(0)).await, 16 * GB);
    }

    #[tokio::test]
    async fn test_model_exceeds_pool_budget() {
        let mgr = cpu_gpu_mgr(24, 0);
        let model = Arc::new(MockLocalModel::new(Device::Cpu));

        mgr.register("big", model.clone(), 32 * GB).await;

        let err = mgr
            .load("big")
            .await
            .expect_err("loading a model larger than the pool budget must fail");
        let msg = err.to_string();
        let needed_bytes = (32 * GB).to_string();
        let budget_bytes = (24 * GB).to_string();
        assert!(
            msg.contains(&needed_bytes) && msg.contains(&budget_bytes),
            "error should mention both needed and budget bytes, got: {msg}"
        );
    }

    #[tokio::test]
    async fn test_unload_frees_pool_budget() {
        let mgr = cpu_gpu_mgr(0, 24);
        let model = Arc::new(MockLocalModel::new(Device::Cuda(0)));

        mgr.register("g1", model.clone(), 8 * GB).await;
        mgr.load("g1").await.unwrap();
        assert_eq!(mgr.used_bytes(Pool::Gpu(0)).await, 8 * GB);
        assert_eq!(mgr.available_bytes(Pool::Gpu(0)).await, 16 * GB);

        mgr.unload("g1").await.unwrap();
        assert_eq!(mgr.used_bytes(Pool::Gpu(0)).await, 0);
        assert_eq!(mgr.available_bytes(Pool::Gpu(0)).await, 24 * GB);
    }

    #[tokio::test]
    async fn test_status_includes_pool() {
        let mgr = cpu_gpu_mgr(24, 24);
        let cpu = Arc::new(MockLocalModel::new(Device::Cpu));
        let gpu = Arc::new(MockLocalModel::new(Device::Cuda(0)));

        mgr.register("cpu", cpu.clone(), 4 * GB).await;
        mgr.register("gpu", gpu.clone(), 8 * GB).await;
        mgr.load("cpu").await.unwrap();

        let statuses = mgr.status().await;
        assert_eq!(statuses.len(), 2);

        let cpu_status = statuses
            .iter()
            .find(|s| s.id == "cpu")
            .expect("cpu missing");
        let gpu_status = statuses
            .iter()
            .find(|s| s.id == "gpu")
            .expect("gpu missing");

        assert_eq!(cpu_status.pool, Pool::Cpu);
        assert!(cpu_status.loaded);
        assert_eq!(cpu_status.memory_estimate_bytes, 4 * GB);

        assert_eq!(gpu_status.pool, Pool::Gpu(0));
        assert!(!gpu_status.loaded);
        assert_eq!(gpu_status.memory_estimate_bytes, 8 * GB);
    }

    #[tokio::test]
    async fn test_with_budgets_gb() {
        let mgr = ModelManager::with_budgets_gb(100.0, 24.0);
        assert_eq!(mgr.available_bytes(Pool::Cpu).await, 100 * GB);
        assert_eq!(mgr.available_bytes(Pool::Gpu(0)).await, 24 * GB);
    }

    #[tokio::test]
    async fn test_unknown_pool_implicit_zero_budget() {
        // Manager only knows about Pool::Cpu and Pool::Gpu(0).
        let mgr = cpu_gpu_mgr(24, 24);
        // But this model targets Pool::Gpu(7) which is implicitly 0.
        let model = Arc::new(MockLocalModel::new(Device::Cuda(7)));

        mgr.register("orphan", model.clone(), GB).await;

        let err = mgr
            .load("orphan")
            .await
            .expect_err("model on unbudgeted pool must fail to load");
        let msg = err.to_string();
        assert!(
            msg.contains("budget is only 0 bytes"),
            "error should mention zero budget, got: {msg}"
        );
        assert!(!mgr.is_loaded("orphan").await);
        assert_eq!(mgr.available_bytes(Pool::Gpu(7)).await, 0);
    }

    #[tokio::test]
    async fn test_load_adapter_records_memory_and_status() {
        let mgr = cpu_gpu_mgr(16, 0);
        let model = Arc::new(MockLocalModel::new(Device::Cpu).with_adapter_memory(2 * GB));
        mgr.register("base", model.clone(), 4 * GB).await;
        mgr.load("base").await.unwrap();

        let dir = make_peft_adapter_dir(1024, 256);
        let handle = mgr
            .load_adapter("base", dir.path(), AdapterOptions::new("a1"))
            .await
            .expect("load_adapter should succeed");

        assert_eq!(handle.adapter_id, "a1");
        assert_eq!(handle.memory_bytes, 2 * GB);
        assert_eq!(handle.mount_strategy, AdapterMountStrategy::Attached);

        let listed = mgr.list_adapters("base").await.unwrap();
        assert_eq!(listed.len(), 1);
        let entry = &listed[0];
        assert_eq!(entry.adapter_id, "a1");
        assert!((entry.scale - 1.0).abs() < f32::EPSILON);
        assert_eq!(entry.source_dir, dir.path().to_path_buf());
        assert_eq!(entry.memory_bytes, 2 * GB);

        let statuses = mgr.status().await;
        let base_status = statuses.iter().find(|s| s.id == "base").expect("base");
        assert_eq!(base_status.memory_estimate_bytes, 4 * GB + 2 * GB);
        assert_eq!(base_status.adapters.len(), 1);
    }

    #[tokio::test]
    async fn test_load_adapter_refuses_duplicate_id() {
        let mgr = cpu_gpu_mgr(16, 0);
        let model = Arc::new(MockLocalModel::new(Device::Cpu).with_adapter_memory(GB));
        mgr.register("base", model.clone(), 4 * GB).await;
        mgr.load("base").await.unwrap();

        let dir = make_peft_adapter_dir(1024, 256);
        mgr.load_adapter("base", dir.path(), AdapterOptions::new("dup"))
            .await
            .expect("first mount should succeed");

        let dir2 = make_peft_adapter_dir(1024, 256);
        let err = mgr
            .load_adapter("base", dir2.path(), AdapterOptions::new("dup"))
            .await
            .expect_err("second mount with same id must fail");
        let msg = err.to_string();
        assert!(
            msg.contains("already mounted"),
            "expected 'already mounted' in error, got: {msg}"
        );
    }

    #[tokio::test]
    async fn test_load_adapter_fails_if_pool_full() {
        let mut budgets = HashMap::new();
        budgets.insert(Pool::Cpu, 8 * 1024_u64);
        budgets.insert(Pool::Gpu(0), 0);
        let mgr = ModelManager::new(budgets);

        let model = Arc::new(MockLocalModel::new(Device::Cpu));
        mgr.register("base", model.clone(), 4 * 1024).await;
        mgr.load("base").await.unwrap();

        let dir = make_peft_adapter_dir(4 * 1024, 4 * 1024);
        let err = mgr
            .load_adapter("base", dir.path(), AdapterOptions::new("oversized"))
            .await
            .expect_err("mounting an adapter that exceeds budget must fail");
        let msg = err.to_string();
        assert!(
            msg.contains("would exceed pool"),
            "expected 'would exceed pool' in error, got: {msg}"
        );
        assert!(
            msg.contains("CPU") || msg.contains("cpu") || msg.contains("Cpu"),
            "expected pool name in error, got: {msg}"
        );

        assert!(mgr.list_adapters("base").await.unwrap().is_empty());
    }

    #[tokio::test]
    async fn test_unload_adapter_decrements_memory() {
        let mgr = cpu_gpu_mgr(16, 0);
        let model = Arc::new(MockLocalModel::new(Device::Cpu).with_adapter_memory(2 * GB));
        mgr.register("base", model.clone(), 4 * GB).await;
        mgr.load("base").await.unwrap();

        let dir = make_peft_adapter_dir(1024, 256);
        mgr.load_adapter("base", dir.path(), AdapterOptions::new("rm"))
            .await
            .unwrap();

        let statuses = mgr.status().await;
        let base = statuses.iter().find(|s| s.id == "base").expect("base");
        assert_eq!(base.memory_estimate_bytes, 4 * GB + 2 * GB);

        mgr.unload_adapter("base", "rm").await.unwrap();

        let statuses = mgr.status().await;
        let base = statuses.iter().find(|s| s.id == "base").expect("base");
        assert_eq!(base.memory_estimate_bytes, 4 * GB);
        assert!(base.adapters.is_empty());

        let err = mgr
            .unload_adapter("base", "rm")
            .await
            .expect_err("second unload must fail");
        let msg = err.to_string();
        assert!(
            msg.contains("is not mounted"),
            "expected 'is not mounted' in error, got: {msg}"
        );
    }

    #[tokio::test]
    async fn test_list_adapters_reflects_state() {
        use std::collections::HashSet;

        let mgr = cpu_gpu_mgr(16, 0);
        let model = Arc::new(MockLocalModel::new(Device::Cpu).with_adapter_memory(1024));
        mgr.register("base", model.clone(), 4 * GB).await;
        mgr.load("base").await.unwrap();

        for id in ["a", "b", "c"] {
            let dir = make_peft_adapter_dir(512, 128);
            mgr.load_adapter("base", dir.path(), AdapterOptions::new(id))
                .await
                .unwrap();
        }

        let list = mgr.list_adapters("base").await.unwrap();
        assert_eq!(list.len(), 3);
        let ids: HashSet<String> = list.iter().map(|a| a.adapter_id.clone()).collect();
        assert!(ids.contains("a") && ids.contains("b") && ids.contains("c"));

        mgr.unload_adapter("base", "b").await.unwrap();

        let list = mgr.list_adapters("base").await.unwrap();
        assert_eq!(list.len(), 2);
        let ids: HashSet<String> = list.iter().map(|a| a.adapter_id.clone()).collect();
        assert!(ids.contains("a") && ids.contains("c"));
        assert!(!ids.contains("b"));

        let err = mgr
            .list_adapters("nonexistent")
            .await
            .expect_err("list_adapters on unknown model must fail");
        let msg = err.to_string();
        assert!(
            msg.contains("not registered"),
            "expected 'not registered' in error, got: {msg}"
        );
    }

    // -- unified provider registry (remote + local dispatch) ----------------

    #[tokio::test]
    async fn test_register_remote_provider_and_complete_by_name() {
        use futures_util::StreamExt;

        let mgr = cpu_gpu_mgr(24, 24);
        mgr.register_provider("fast", Arc::new(MockRemoteModel::new("fast")), None, 0)
            .await;

        // Dispatch by name.
        let resp = mgr
            .complete("fast", ModelRequest::new(Vec::new()))
            .await
            .expect("remote complete by name should succeed");
        assert_eq!(resp.content.as_deref(), Some("echo from fast"));
        assert_eq!(resp.model, "fast");

        // Streaming counterpart.
        let mut stream = mgr
            .stream("fast", ModelRequest::new(Vec::new()))
            .await
            .expect("remote stream by name should succeed");
        let first = stream.next().await.expect("one chunk").expect("ok chunk");
        assert_eq!(first.delta.as_deref(), Some("echo from fast"));

        // The instance can be fetched back out to use / compose directly.
        let handle = mgr.get("fast").await.expect("get returns the model");
        assert_eq!(handle.model_id(), "fast");

        // Remote providers are always "ready" and never count against a budget.
        assert!(mgr.is_loaded("fast").await);
        assert_eq!(mgr.used_bytes(Pool::Remote).await, 0);
        assert_eq!(mgr.used_bytes(Pool::Cpu).await, 0);

        let status = mgr.status().await;
        let fast = status.iter().find(|s| s.id == "fast").expect("fast status");
        assert_eq!(fast.pool, Pool::Remote);
        assert_eq!(fast.memory_estimate_bytes, 0);
    }

    #[tokio::test]
    async fn test_remote_provider_does_not_evict_or_charge_local_pool() {
        // Pool fits exactly one 16 GB local model.
        let mgr = cpu_gpu_mgr(16, 0);
        let local = Arc::new(MockLocalModel::new(Device::Cpu));
        mgr.register("local", local.clone(), 16 * GB).await;
        mgr.load("local").await.unwrap();

        // Registering and using a remote provider must not disturb the CPU
        // pool or evict the resident local model.
        mgr.register_provider("remote", Arc::new(MockRemoteModel::new("remote")), None, 0)
            .await;
        let resp = mgr
            .complete("remote", ModelRequest::new(Vec::new()))
            .await
            .expect("remote complete should succeed even with a full local pool");
        assert_eq!(resp.content.as_deref(), Some("echo from remote"));

        assert!(
            mgr.is_loaded("local").await,
            "remote dispatch must not evict the resident local model"
        );
        assert_eq!(mgr.used_bytes(Pool::Cpu).await, 16 * GB);
    }

    #[tokio::test]
    async fn test_complete_on_lifecycle_only_entry_is_unsupported() {
        let mgr = cpu_gpu_mgr(24, 0);
        // Registered via the lifecycle-only path: no chat Model attached.
        let media = Arc::new(MockLocalModel::new(Device::Cpu));
        mgr.register("media", media, GB).await;

        let err = mgr
            .complete("media", ModelRequest::new(Vec::new()))
            .await
            .expect_err("completing a lifecycle-only entry must fail");
        let msg = err.to_string();
        assert!(
            msg.contains("does not implement the chat Model interface"),
            "expected an unsupported-dispatch error, got: {msg}"
        );

        // get() returns None for lifecycle-only entries.
        assert!(mgr.get("media").await.is_none());
    }

    #[tokio::test]
    async fn test_complete_on_unknown_id_is_validation_error() {
        let mgr = cpu_gpu_mgr(24, 0);
        let err = mgr
            .complete("ghost", ModelRequest::new(Vec::new()))
            .await
            .expect_err("completing an unregistered id must fail");
        assert!(
            matches!(err, BlazenError::Validation { .. }),
            "expected BlazenError::Validation, got {err:?}"
        );
    }

    #[tokio::test]
    async fn test_local_llm_via_register_provider_auto_loads_and_charges_pool() {
        let mgr = cpu_gpu_mgr(0, 24);
        let llm = Arc::new(MockLocalLlm::new("llama", Device::Cuda(0)));
        // Same instance is both the chat Model and the LocalModel.
        mgr.register_provider("llama", llm.clone(), Some(llm.clone()), 8 * GB)
            .await;

        // Not loaded yet — pool is empty.
        assert!(!mgr.is_loaded("llama").await);
        assert_eq!(mgr.used_bytes(Pool::Gpu(0)).await, 0);

        // complete() auto-loads the local weights, then dispatches.
        let resp = mgr
            .complete("llama", ModelRequest::new(Vec::new()))
            .await
            .expect("local llm complete should auto-load and succeed");
        assert_eq!(resp.content.as_deref(), Some("local llama"));

        assert!(mgr.is_loaded("llama").await);
        assert_eq!(mgr.used_bytes(Pool::Gpu(0)).await, 8 * GB);

        // Unload frees the pool and the dispatch handle remains fetchable.
        mgr.unload("llama").await.unwrap();
        assert!(!mgr.is_loaded("llama").await);
        assert_eq!(mgr.used_bytes(Pool::Gpu(0)).await, 0);
        assert!(mgr.get("llama").await.is_some());
    }

    #[tokio::test]
    async fn test_remote_provider_rejects_adapter_mount() {
        let mgr = cpu_gpu_mgr(24, 0);
        mgr.register_provider("remote", Arc::new(MockRemoteModel::new("remote")), None, 0)
            .await;
        let dir = make_peft_adapter_dir(1024, 256);
        let err = mgr
            .load_adapter("remote", dir.path(), AdapterOptions::new("a1"))
            .await
            .expect_err("mounting an adapter on a remote provider must fail");
        let msg = err.to_string();
        assert!(
            msg.contains("remote provider") && msg.contains("does not support LoRA adapters"),
            "expected a remote-adapter-unsupported error, got: {msg}"
        );
    }
}

#[cfg(all(test, feature = "training"))]
mod training_tests {
    use super::*;
    use async_trait::async_trait;

    struct EmptyDataset;

    #[async_trait]
    impl TrainingDataset for EmptyDataset {
        fn len(&self) -> usize {
            0
        }
        async fn batch(
            &self,
            _batch_size: usize,
            _idx: usize,
        ) -> Result<blazen_train::TrainingBatch, BlazenTrainError> {
            Err(BlazenTrainError::Dataset(
                "training_tests dataset is intentionally empty".to_string(),
            ))
        }
    }

    #[tokio::test]
    async fn train_lora_validation_error_propagates() {
        let mgr = ModelManager::with_budgets_gb(8.0, 0.0);
        let cfg = TrainConfig {
            max_steps: 0,
            device: Some("cpu".to_string()),
            ..TrainConfig::default()
        };

        let err = mgr
            .train_lora(cfg, Box::new(EmptyDataset), None)
            .await
            .expect_err("max_steps=0 must surface as a validation error");

        match err {
            BlazenError::Validation { message, .. } => {
                assert!(
                    message.contains("max_steps"),
                    "expected 'max_steps' in validation message, got: {message}"
                );
            }
            other => panic!("expected BlazenError::Validation, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn train_lora_rejects_unknown_device() {
        let mgr = ModelManager::with_budgets_gb(8.0, 0.0);
        let cfg = TrainConfig {
            device: Some("opencl:0".to_string()),
            ..TrainConfig::default()
        };

        let err = mgr
            .train_lora(cfg, Box::new(EmptyDataset), None)
            .await
            .expect_err("unknown device must surface as a validation error");
        match err {
            BlazenError::Validation { message, .. } => {
                assert!(
                    message.contains("opencl"),
                    "expected device name in message, got: {message}"
                );
            }
            other => panic!("expected BlazenError::Validation, got {other:?}"),
        }
    }

    // PR8 Wave 10: dummy preference / rated datasets used by the
    // validation-error tests. Both implementations return Dataset errors
    // from `batch()`, but neither verb under test reaches `batch()` —
    // every assertion below trips on config validation in
    // `Trainer::new_dpo/orpo/simpo/kto/full_finetune` before any dataset
    // method is called.
    struct EmptyPreferenceDataset;

    #[async_trait]
    impl PreferenceDataset for EmptyPreferenceDataset {
        fn len(&self) -> usize {
            0
        }
        async fn batch(
            &self,
            _batch_size: usize,
            _idx: usize,
        ) -> Result<blazen_train::PreferenceBatch, BlazenTrainError> {
            Err(BlazenTrainError::Dataset(
                "EmptyPreferenceDataset is intentionally empty".to_string(),
            ))
        }
    }

    struct EmptyRatedDataset;

    #[async_trait]
    impl RatedDataset for EmptyRatedDataset {
        fn len(&self) -> usize {
            0
        }
        async fn batch(
            &self,
            _batch_size: usize,
            _idx: usize,
        ) -> Result<blazen_train::KtoBatch, BlazenTrainError> {
            Err(BlazenTrainError::Dataset(
                "EmptyRatedDataset is intentionally empty".to_string(),
            ))
        }
    }

    fn cpu_core(max_steps: usize) -> TrainCoreConfig {
        TrainCoreConfig {
            max_steps,
            device: Some("cpu".to_string()),
            ..TrainCoreConfig::default()
        }
    }

    #[tokio::test]
    async fn train_dpo_rejects_invalid_config() {
        let mgr = ModelManager::with_budgets_gb(8.0, 0.0);
        let cfg = DpoConfig {
            core: cpu_core(0),
            ..DpoConfig::default()
        };
        let err = mgr
            .train_dpo(cfg, Arc::new(EmptyPreferenceDataset), None)
            .await
            .expect_err("max_steps=0 must surface as a validation error");
        match err {
            BlazenError::Validation { message, .. } => assert!(
                message.contains("max_steps"),
                "expected 'max_steps' in validation message, got: {message}"
            ),
            other => panic!("expected BlazenError::Validation, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn train_orpo_rejects_invalid_config() {
        let mgr = ModelManager::with_budgets_gb(8.0, 0.0);
        let cfg = OrpoConfig {
            core: cpu_core(0),
            ..OrpoConfig::default()
        };
        let err = mgr
            .train_orpo(cfg, Arc::new(EmptyPreferenceDataset), None)
            .await
            .expect_err("max_steps=0 must surface as a validation error");
        assert!(
            matches!(err, BlazenError::Validation { .. }),
            "expected BlazenError::Validation, got {err:?}"
        );
    }

    #[tokio::test]
    async fn train_orpo_rejects_unknown_device() {
        let mgr = ModelManager::with_budgets_gb(8.0, 0.0);
        let cfg = OrpoConfig {
            core: TrainCoreConfig {
                device: Some("opencl:0".to_string()),
                ..TrainCoreConfig::default()
            },
            ..OrpoConfig::default()
        };
        let err = mgr
            .train_orpo(cfg, Arc::new(EmptyPreferenceDataset), None)
            .await
            .expect_err("unknown device must surface as a validation error");
        match err {
            BlazenError::Validation { message, .. } => assert!(
                message.contains("opencl"),
                "expected device name in message, got: {message}"
            ),
            other => panic!("expected BlazenError::Validation, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn train_simpo_rejects_invalid_config() {
        let mgr = ModelManager::with_budgets_gb(8.0, 0.0);
        let cfg = SimpoConfig {
            core: cpu_core(0),
            ..SimpoConfig::default()
        };
        let err = mgr
            .train_simpo(cfg, Arc::new(EmptyPreferenceDataset), None)
            .await
            .expect_err("max_steps=0 must surface as a validation error");
        assert!(
            matches!(err, BlazenError::Validation { .. }),
            "expected BlazenError::Validation, got {err:?}"
        );
    }

    #[tokio::test]
    async fn train_kto_rejects_invalid_config() {
        let mgr = ModelManager::with_budgets_gb(8.0, 0.0);
        let cfg = KtoConfig {
            core: cpu_core(0),
            ..KtoConfig::default()
        };
        let err = mgr
            .train_kto(cfg, Arc::new(EmptyRatedDataset), None)
            .await
            .expect_err("max_steps=0 must surface as a validation error");
        assert!(
            matches!(err, BlazenError::Validation { .. }),
            "expected BlazenError::Validation, got {err:?}"
        );
    }

    #[tokio::test]
    async fn fine_tune_rejects_gradient_checkpointing() {
        // gradient_checkpointing=true is rejected up-front by
        // Trainer::new_full_finetune because candle 0.10.2 has no
        // activation-checkpointing primitive. The BlazenTrainError::Unsupported
        // arm of the From impl folds into BlazenError::Validation with a
        // "training: " prefix, so we assert on the canonical substring.
        let mgr = ModelManager::with_budgets_gb(8.0, 0.0);
        let cfg = FullFineTuneConfig {
            core: cpu_core(1),
            gradient_checkpointing: true,
        };
        let err = mgr
            .fine_tune(cfg, Arc::new(EmptyDataset), None)
            .await
            .expect_err("gradient_checkpointing=true must be rejected");
        match err {
            BlazenError::Validation { message, .. } => assert!(
                message.contains("gradient checkpointing"),
                "expected 'gradient checkpointing' in message, got: {message}"
            ),
            other => panic!("expected BlazenError::Validation, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn fine_tune_rejects_invalid_config() {
        let mgr = ModelManager::with_budgets_gb(8.0, 0.0);
        let cfg = FullFineTuneConfig {
            core: cpu_core(0),
            gradient_checkpointing: false,
        };
        let err = mgr
            .fine_tune(cfg, Arc::new(EmptyDataset), None)
            .await
            .expect_err("max_steps=0 must surface as a validation error");
        assert!(
            matches!(err, BlazenError::Validation { .. }),
            "expected BlazenError::Validation, got {err:?}"
        );
    }
}
