hf_fetch_model/

inspect.rs

// SPDX-License-Identifier: MIT OR Apache-2.0

//! Tensor-file header inspection (local and remote).
//!
//! Reads tensor metadata (names, shapes, dtypes, byte offsets) without
//! downloading full weight data. `.safetensors` files resolve cache-first
//! with HTTP Range request fallback; `.gguf` / `.npz` / `.pth` files are
//! inspected from the local cache via the `anamnesis` parser crate
//! ([`inspect_gguf_cached`] / [`inspect_npz_cached`] / [`inspect_pth_cached`],
//! v0.10.2–v0.10.3 — remote inspect for those three is planned for v0.11).
//!
//! The primary types are [`TensorInfo`] (per-tensor metadata),
//! [`SafetensorsHeaderInfo`] (the format-agnostic parsed-header shape all
//! four formats return), and [`ShardedIndex`] (shard-to-tensor mapping for
//! sharded models). For cheap discovery without header parsing,
//! [`list_cached_tensor_files`] enumerates a cached repo's tensor files
//! across all four formats ([`list_cached_safetensors`] is the
//! `.safetensors`-only subset).
//!
//! The module also reads small JSON sidecars from the same cache-first /
//! HTTP-fallback path: [`AdapterConfig`] (`adapter_config.json`, for `PEFT`
//! adapters) and [`ModelConfig`] (`config.json`, the architecture parameters
//! that drive `inspect --check-gpu --context` KV-cache budgeting), via
//! [`fetch_model_config`] / [`fetch_model_config_cached`].

use std::collections::HashMap;
use std::path::{Path, PathBuf};

use serde::Serialize;
use tokio::task::JoinSet;

use crate::cache;
use crate::cache_layout;
use crate::chunked;
use crate::error::FetchError;

// -----------------------------------------------------------------------
// Types
// -----------------------------------------------------------------------

/// Metadata for a single tensor from a `.safetensors` header.
///
/// This is hf-fetch-model's own type — lightweight, no quantization logic.
/// Consumers (e.g., anamnesis) map this into their own richer types.
#[derive(Debug, Clone, Serialize)]
pub struct TensorInfo {
    /// Tensor name (e.g., `"model.layers.0.self_attn.q_proj.weight"`).
    pub name: String,
    /// Element dtype string as it appears in the header (e.g., `"F8_E4M3"`, `"BF16"`).
    pub dtype: String,
    /// Tensor shape (e.g., `[7168, 7168]`).
    pub shape: Vec<usize>,
    /// Byte offset range `[start, end)` within the data section of the file.
    pub data_offsets: (u64, u64),
}

impl TensorInfo {
    /// Total number of elements (product of shape dimensions).
    ///
    /// Returns `1` for a scalar (empty shape).
    #[must_use]
    pub fn num_elements(&self) -> u64 {
        self.shape.iter().fold(1u64, |acc, &d| {
            // CAST: usize → u64, dimension values fit in u64
            #[allow(clippy::as_conversions)]
            let dim = d as u64;
            acc.saturating_mul(dim)
        })
    }

    /// Byte length of the tensor data (`end - start`).
    #[must_use]
    pub const fn byte_len(&self) -> u64 {
        self.data_offsets.1.saturating_sub(self.data_offsets.0)
    }

    /// Bytes per element for the tensor's dtype, if recognized.
    ///
    /// Returns `None` for unknown dtype strings. Recognized dtypes:
    ///
    /// | Dtype string | Bytes | Notes |
    /// |-------------|-------|-------|
    /// | `"BOOL"` | 1 | |
    /// | `"U8"`, `"I8"` | 1 | |
    /// | `"F8_E4M3"`, `"F8_E5M2"` | 1 | FP8 variants |
    /// | `"U16"`, `"I16"`, `"F16"`, `"BF16"` | 2 | |
    /// | `"U32"`, `"I32"`, `"F32"` | 4 | |
    /// | `"U64"`, `"I64"`, `"F64"` | 8 | |
    #[must_use]
    pub fn dtype_bytes(&self) -> Option<usize> {
        // BORROW: explicit .as_str() instead of Deref coercion
        match self.dtype.as_str() {
            "BOOL" | "U8" | "I8" | "F8_E4M3" | "F8_E5M2" => Some(1),
            "U16" | "I16" | "F16" | "BF16" => Some(2),
            "U32" | "I32" | "F32" => Some(4),
            "U64" | "I64" | "F64" => Some(8),
            _ => None,
        }
    }
}

/// Bytes per element for a model's activation dtype, as spelled in a
/// `config.json` `torch_dtype` field.
///
/// Distinct from [`TensorInfo::dtype_bytes`], which maps the *safetensors*
/// header spellings (`"BF16"`, `"F16"`, …); `config.json` uses the `PyTorch`
/// spellings (`"bfloat16"`, `"float16"`, `"float32"`, `"float8_e4m3fn"`).
/// Used to size the KV cache, whose element dtype tracks the model's
/// activations (typically `bf16` / `fp16`) independently of weight
/// quantization. Defaults to `2` when the dtype is absent or unrecognized —
/// the modern inference default and the safe assumption for KV sizing.
#[must_use]
pub fn torch_dtype_bytes(torch_dtype: Option<&str>) -> u8 {
    match torch_dtype {
        Some("float32" | "float") => 4,
        Some("float8_e4m3fn" | "float8_e5m2") => 1,
        // `bf16` / `fp16`, and any unknown or absent dtype: 2-byte activations.
        _ => 2,
    }
}

/// Quantization scheme + size estimates for a cached `.safetensors` file.
///
/// Populated by [`inspect_safetensors_local`] (the cache-hit path) via
/// `anamnesis::InspectInfo::from(&header)`. Absent (`None`) when:
/// - the safetensors file has no detected quantization (`QuantScheme::Unquantized`),
/// - the inspect ran over HTTP Range (the remote path's bespoke parser
///   doesn't go through anamnesis until v0.11.1), or
/// - the file format isn't safetensors (`GGUF` / `NPZ` / `PTH` carry no
///   quant-method metadata).
///
/// Decoupled from `anamnesis::QuantScheme` (a `#[non_exhaustive]` enum) so
/// downstream library consumers (`candle-mi`, `anamnesis`) aren't forced to
/// match every variant. The `scheme` field stores `QuantScheme`'s `Display`
/// output (`"FineGrainedFp8"`, `"Bnb4"`, `"Gptq"`, `"Awq"`, …); consumers
/// that need to match exact variants should call
/// `anamnesis::parse_safetensors_header` themselves.
#[derive(Debug, Clone, Serialize)]
pub struct QuantInfo {
    /// Detected quantization scheme as the `Display` form of
    /// `anamnesis::QuantScheme` (e.g. `"FineGrainedFp8"`, `"Bnb4"`).
    pub scheme: String,
    /// Bytes stored on disk for tensor data (header excluded).
    pub stored_bytes: u64,
    /// Estimated bytes after dequantising to `BF16`. For `BnB-NF4`/`FP4`
    /// (`U8`-packed nibbles), this is `stored_bytes × 4`; for `FP8` / `GPTQ` /
    /// `AWQ` / `BnB-INT8` it's `num_elements × 2` summed over weight
    /// tensors, plus passthrough tensors copied as-is. The formula lives
    /// in `anamnesis::InspectInfo::from(&SafetensorsHeader)` — hf-fm just
    /// reads the result.
    pub dequantized_bytes: u64,
}

/// Parsed safetensors header metadata.
///
/// Marked `#[non_exhaustive]` (since v0.10.3) — the struct has been
/// growing through v0.10.x (the `quant_info` field landed in Phase C)
/// and will keep growing in v0.11.x. External library consumers should
/// pattern-match with `..` or use field reads, not exhaustive struct
/// literals.
#[derive(Debug, Clone, Serialize)]
#[non_exhaustive]
pub struct SafetensorsHeaderInfo {
    /// All tensors in the header, in the order they appear in the JSON.
    pub tensors: Vec<TensorInfo>,
    /// Raw `__metadata__` entries, if present.
    ///
    /// For quantized models, this typically contains entries like
    /// `quant_method`, `bits`, `group_size` that consumers like anamnesis
    /// use to distinguish GPTQ from AWQ without downloading weights.
    pub metadata: Option<HashMap<String, String>>,
    /// Size of the JSON header in bytes.
    pub header_size: u64,
    /// Total file size in bytes (header + data), if known.
    ///
    /// **Source:** for local files, from `std::fs::metadata().len()`. For HTTP
    /// Range requests, extracted from the `Content-Range` response header of
    /// the first request (`bytes 0-7/TOTAL` → `TOTAL`). This is free — no
    /// extra request needed.
    pub file_size: Option<u64>,
    /// Quantization scheme + size estimates (cached safetensors only).
    ///
    /// Populated by [`inspect_safetensors_local`] when a non-`Unquantized`
    /// `QuantScheme` is detected. `None` for unquantized safetensors,
    /// `GGUF` / `NPZ` / `PTH` files, and the remote safetensors path
    /// (until v0.11.1 migrates that path to anamnesis).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub quant_info: Option<QuantInfo>,
}

impl SafetensorsHeaderInfo {
    /// Total parameter count across all tensors.
    #[must_use]
    pub fn total_params(&self) -> u64 {
        self.tensors
            .iter()
            .map(TensorInfo::num_elements)
            .fold(0u64, u64::saturating_add)
    }

    /// Returns tensors matching a dtype string (e.g., `"F8_E4M3"`).
    #[must_use]
    pub fn tensors_with_dtype(&self, dtype: &str) -> Vec<&TensorInfo> {
        self.tensors
            .iter()
            // BORROW: explicit .as_str() instead of Deref coercion
            .filter(|t| t.dtype.as_str() == dtype)
            .collect()
    }

    /// Constructs a new [`SafetensorsHeaderInfo`] from its core fields.
    ///
    /// Since v0.10.3 the struct is `#[non_exhaustive]` — this constructor is
    /// the canonical way to build one from outside the `hf-fetch-model` lib
    /// crate (e.g. the `hf-fm` binary crate, downstream consumers like
    /// `candle-mi`). Inside the lib crate, struct-literal syntax stays
    /// available for the inspect entry points.
    ///
    /// `quant_info` is typically `None`; populated only by
    /// [`inspect_safetensors_local`] for cached, quantized safetensors files.
    #[must_use]
    pub fn new(
        tensors: Vec<TensorInfo>,
        metadata: Option<HashMap<String, String>>,
        header_size: u64,
        file_size: Option<u64>,
        quant_info: Option<QuantInfo>,
    ) -> Self {
        Self {
            tensors,
            metadata,
            header_size,
            file_size,
            quant_info,
        }
    }
}

/// The source from which a header was read.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum InspectSource {
    /// Read from local cache (no network).
    Cached,
    /// Fetched via HTTP Range requests.
    Remote,
}

/// Parsed `model.safetensors.index.json` for a sharded model.
#[derive(Debug, Clone, Serialize)]
pub struct ShardedIndex {
    /// Mapping from tensor name to shard filename.
    pub weight_map: HashMap<String, String>,
    /// Ordered list of unique shard filenames.
    pub shards: Vec<String>,
    /// Raw metadata from the index, if present.
    pub metadata: Option<HashMap<String, serde_json::Value>>,
}

/// `PEFT` adapter configuration parsed from `adapter_config.json`.
///
/// Contains the key fields that identify an adapter: the `PEFT` type,
/// base model, `LoRA` rank and scaling parameters, and target modules.
/// All fields are optional because adapter configs vary across `PEFT` methods.
#[derive(Debug, Clone, Serialize)]
pub struct AdapterConfig {
    /// `PEFT` method type (e.g., `"LORA"`, `"ADALORA"`, `"IA3"`).
    pub peft_type: Option<String>,
    /// The base model this adapter was trained on.
    pub base_model_name_or_path: Option<String>,
    /// `LoRA` rank (the `r` parameter). Only meaningful for `LoRA`-family methods.
    pub r: Option<u32>,
    /// `LoRA` alpha scaling factor. Only meaningful for `LoRA`-family methods.
    pub lora_alpha: Option<f64>,
    /// List of model modules targeted by the adapter.
    pub target_modules: Vec<String>,
    /// Task type the adapter was trained for (e.g., `"CAUSAL_LM"`).
    pub task_type: Option<String>,
}

/// Attention- and cache-relevant fields parsed from a model's `config.json`.
///
/// Every field is [`Option`] because configs vary across architecture
/// families; the KV-cache estimator decides which combinations are
/// computable and which fall back to an "unavailable" verdict. Legacy
/// `n_layer` / `n_head` / `n_head_kv` spellings are absorbed by serde aliases
/// on the private deserialization struct. Drives KV-cache budgeting for
/// `inspect --check-gpu --context`.
#[derive(Debug, Clone, Default, Serialize)]
#[non_exhaustive]
pub struct ModelConfig {
    /// Architecture tag (e.g. `"llama"`, `"qwen3"`, `"gemma2"`, `"deepseek_v2"`).
    pub model_type: Option<String>,
    /// Number of transformer layers (`num_hidden_layers` / `n_layer`).
    pub num_hidden_layers: Option<u32>,
    /// Number of query attention heads (`num_attention_heads` / `n_head`).
    pub num_attention_heads: Option<u32>,
    /// Number of key/value heads for `GQA` (`num_key_value_heads` /
    /// `num_kv_heads` / `n_head_kv`). Absent ⇒ `MHA` (equals
    /// `num_attention_heads`).
    pub num_key_value_heads: Option<u32>,
    /// Explicit per-head dimension when stated (Gemma = 256, Qwen3 = 128).
    /// Absent ⇒ derived as `hidden_size / num_attention_heads`.
    pub head_dim: Option<u32>,
    /// Model hidden size, used to derive `head_dim` when it is not explicit.
    pub hidden_size: Option<u32>,
    /// Activation dtype spelling (`"bfloat16"`, `"float16"`, …); sizes the
    /// KV-cache element via [`torch_dtype_bytes`].
    pub torch_dtype: Option<String>,
    /// Sliding-window span in tokens when the model uses windowed attention.
    /// `null` / absent ⇒ full attention.
    pub sliding_window: Option<u32>,
    /// Global-attention period for mixed local/global layouts (Gemma-3:
    /// every `N`-th layer is a full-attention layer).
    pub sliding_window_pattern: Option<u32>,
    /// Explicit on/off switch for sliding-window attention — Qwen2/3 ship a
    /// `sliding_window` value but disable it with `false`.
    pub use_sliding_window: Option<bool>,
    /// `MLA` latent-KV rank (`DeepSeek`). Presence marks multi-head latent
    /// attention, where the naive KV formula does not apply.
    pub kv_lora_rank: Option<u32>,
    /// `MLA` decoupled-`RoPE` key dimension (`DeepSeek`); part of the latent-KV
    /// size used by the documented `MLA` estimate.
    pub qk_rope_head_dim: Option<u32>,
    /// Per-layer kind tags for hybrid models (`"attention"` / `"mamba"` /
    /// `"linear_attention"` / …). Primary hybrid-layout signal (Granite-4).
    pub layer_types: Option<Vec<String>>,
    /// Nemotron-H layer-layout string (`"M-M-M-M*-…"`: `*` = attention,
    /// `M` = Mamba, `-` = FFN-only). Alternative hybrid-layout signal.
    pub hybrid_override_pattern: Option<String>,
    /// Explicit indices of the attention layers (Bamba). Alternative
    /// hybrid-layout signal; the remaining layers are recurrent.
    pub attn_layer_indices: Option<Vec<u32>>,
    /// Period of full-attention layers — every `N`-th layer is attention, the
    /// rest recurrent (Qwen3-Next). Alternative hybrid-layout signal.
    pub full_attention_interval: Option<u32>,
    /// Mamba2 SSM head count (`mamba_n_heads` / `mamba_num_heads`).
    pub mamba_n_heads: Option<u32>,
    /// Mamba2 SSM per-head dimension (`mamba_d_head` / `mamba_head_dim`).
    pub mamba_d_head: Option<u32>,
    /// Mamba2 SSM state size (`mamba_d_state` / `ssm_state_size`).
    pub mamba_d_state: Option<u32>,
    /// Mamba2 causal-convolution width (`mamba_d_conv` / `conv_kernel`).
    pub mamba_d_conv: Option<u32>,
    /// Mamba2 group count for the convolution (`mamba_n_groups` / `n_groups`).
    pub mamba_n_groups: Option<u32>,
}

// -----------------------------------------------------------------------
// JSON parsing
// -----------------------------------------------------------------------

/// Raw tensor entry as it appears in the safetensors JSON header.
#[derive(serde::Deserialize)]
struct RawTensorEntry {
    dtype: String,
    shape: Vec<usize>,
    data_offsets: (u64, u64),
}

/// Parsed tensor list and optional metadata from a safetensors header.
type ParsedHeader = (Vec<TensorInfo>, Option<HashMap<String, String>>);

/// Parses the safetensors JSON header bytes into tensor metadata.
///
/// Extracts the `__metadata__` key separately (if present).
fn parse_header_json(json_bytes: &[u8], filename: &str) -> Result<ParsedHeader, FetchError> {
    let raw: HashMap<String, serde_json::Value> =
        serde_json::from_slice(json_bytes).map_err(|e| FetchError::SafetensorsHeader {
            filename: filename.to_owned(),
            reason: format!("failed to parse header JSON: {e}"),
        })?;

    let mut metadata: Option<HashMap<String, String>> = None;
    let mut tensors = Vec::new();

    for (key, value) in raw {
        if key == "__metadata__" {
            if let serde_json::Value::Object(obj) = value {
                let mut meta_map = HashMap::new();
                for (mk, mv) in obj {
                    // BORROW: explicit .to_owned()/.to_string() for Value → String conversion
                    let v_str = if let Some(s) = mv.as_str() {
                        s.to_owned()
                    } else {
                        mv.to_string()
                    };
                    meta_map.insert(mk, v_str);
                }
                metadata = Some(meta_map);
            }
            continue;
        }

        let entry: RawTensorEntry =
            serde_json::from_value(value).map_err(|e| FetchError::SafetensorsHeader {
                filename: filename.to_owned(),
                reason: format!("failed to parse tensor \"{key}\": {e}"),
            })?;

        tensors.push(TensorInfo {
            name: key,
            dtype: entry.dtype,
            shape: entry.shape,
            data_offsets: entry.data_offsets,
        });
    }

    // Sort by data offset start to preserve file order.
    tensors.sort_by_key(|t| t.data_offsets.0);

    Ok((tensors, metadata))
}

// -----------------------------------------------------------------------
// Cache resolution
// -----------------------------------------------------------------------

/// Resolves a cached file path for a given repo, revision, and filename.
///
/// Returns `None` if the file is not in the local cache.
fn resolve_cached_path(repo_id: &str, revision: &str, filename: &str) -> Option<PathBuf> {
    let cache_dir = cache::hf_cache_dir().ok()?;
    let repo_dir = cache_layout::repo_dir(&cache_dir, repo_id);
    let commit_hash = cache::read_ref(&repo_dir, revision)?;
    let cached_path = cache_layout::pointer_path(&repo_dir, &commit_hash, filename);
    if cached_path.exists() {
        Some(cached_path)
    } else {
        None
    }
}

// -----------------------------------------------------------------------
// Local file reading
// -----------------------------------------------------------------------

/// Inspects a single `.safetensors` file's header from a local file path.
///
/// Reads the first `8 + header_size` bytes from disk. Does not read tensor data.
///
/// # Blocking I/O
///
/// This function performs synchronous filesystem I/O. In async contexts, wrap
/// it in [`tokio::task::spawn_blocking`] so the calling task does not stall
/// the runtime — particularly important on network-mounted caches (NFS/CIFS)
/// where `read`/`stat` calls can take tens of milliseconds each.
///
/// # Errors
///
/// Returns [`FetchError::Io`] if the file cannot be read.
/// Returns [`FetchError::SafetensorsHeader`] if the header is malformed.
pub fn inspect_safetensors_local(path: &Path) -> Result<SafetensorsHeaderInfo, FetchError> {
    let file_size = std::fs::metadata(path)
        .map_err(|e| FetchError::Io {
            path: path.to_path_buf(),
            source: e,
        })?
        .len();

    // BORROW: explicit .to_string_lossy() for Path → str conversion
    let filename = path.file_name().map_or_else(
        || path.display().to_string(),
        |n| n.to_string_lossy().to_string(),
    );

    let file = std::fs::File::open(path).map_err(|e| FetchError::Io {
        path: path.to_path_buf(),
        source: e,
    })?;

    // Cache-hit path delegates to anamnesis (v0.10.3 Phase B commit 4):
    // single source of truth for safetensors layout. The reader variant
    // reads the 8-byte u64 prefix + the JSON header bytes from the `Read`
    // impl itself, *without* requiring the data section to be present
    // (it bypasses `safetensors::SafeTensors::read_metadata` for exactly
    // this reason). Anamnesis caps the declared header length at 100 MiB
    // internally so the worst-case allocation is bounded.
    //
    // The remote path's `fetch_header_bytes` + `parse_header_json` chain
    // stays bespoke until v0.11.1, when anamnesis's reader variant adapts
    // to an HTTP Range source.
    let header = anamnesis::parse_safetensors_header_from_reader(file).map_err(|e| {
        FetchError::SafetensorsHeader {
            // BORROW: explicit .clone() for the error variant's owned String field
            filename: filename.clone(),
            reason: format!("failed to parse safetensors header: {e}"),
        }
    })?;
    // CAST: usize → u64, anamnesis caps header_size at 100 MiB so it always fits in u64
    #[allow(clippy::as_conversions)]
    let header_size = header.header_size as u64;

    // Phase C: derive `quant_info` before the `header.tensors` move below.
    // `anamnesis::InspectInfo::from(&header)` iterates the already-parsed
    // tensors and aggregates per-role byte sums — pure computation, no I/O.
    // Unquantized models produce no `quant_info` so the renderer suppresses
    // the new `Format:` / `Size:` lines (absence communicates full precision).
    let quant_info = if header.scheme == anamnesis::QuantScheme::Unquantized {
        None
    } else {
        let info = anamnesis::InspectInfo::from(&header);
        Some(QuantInfo {
            // BORROW: explicit .to_string() — anamnesis `QuantScheme` → owned `String`
            scheme: info.format.to_string(),
            stored_bytes: info.current_size,
            dequantized_bytes: info.dequantized_size,
        })
    };

    let mut tensors: Vec<TensorInfo> = header
        .tensors
        .into_iter()
        .map(|t| {
            // CAST: usize → u64, header data_offsets fit in u64 by definition (file size is u64)
            #[allow(clippy::as_conversions)]
            let start = t.data_offsets.0 as u64;
            // CAST: usize → u64, same rationale as above
            #[allow(clippy::as_conversions)]
            let end = t.data_offsets.1 as u64;
            TensorInfo {
                name: t.name,
                // BORROW: explicit .to_string() — anamnesis `Dtype` enum → owned `String`
                dtype: t.dtype.to_string(),
                shape: t.shape,
                data_offsets: (start, end),
            }
        })
        .collect();

    // Preserve hf-fm's v0.10.2 sort order. Anamnesis returns tensors sorted
    // alphabetically by name; hf-fm's inspect table has always been
    // file-ordered (sorted by start offset) so users can spot first/last
    // tensors per shard at a glance.
    tensors.sort_by_key(|t| t.data_offsets.0);

    Ok(SafetensorsHeaderInfo {
        tensors,
        metadata: header.metadata,
        header_size,
        file_size: Some(file_size),
        quant_info,
    })
}

// -----------------------------------------------------------------------
// Remote fetching (HTTP Range requests)
// -----------------------------------------------------------------------

/// Fetches safetensors header bytes via two HTTP Range requests.
///
/// 1. `Range: bytes=0-7` → 8-byte header length (little-endian `u64`)
/// 2. `Range: bytes=8-{8+length-1}` → JSON header
///
/// Returns `(json_bytes, total_file_size)`. The file size is extracted from
/// the `Content-Range` header of the first request.
async fn fetch_header_bytes(
    client: &reqwest::Client,
    url: &str,
    filename: &str,
) -> Result<(Vec<u8>, Option<u64>), FetchError> {
    // Request 1: 8-byte length prefix.
    let resp1 = client
        .get(url)
        .header(reqwest::header::RANGE, "bytes=0-7")
        .send()
        .await
        .map_err(|e| {
            FetchError::Http(format!("failed to fetch header length for {filename}: {e}"))
        })?;

    if !resp1.status().is_success() && resp1.status() != reqwest::StatusCode::PARTIAL_CONTENT {
        return Err(FetchError::Http(format!(
            "Range request for {filename} returned status {}",
            resp1.status()
        )));
    }

    // Extract total file size from Content-Range: bytes 0-7/{total}
    let file_size = resp1
        .headers()
        .get(reqwest::header::CONTENT_RANGE)
        .and_then(|v| v.to_str().ok())
        .and_then(|s| s.split('/').next_back())
        .and_then(|s| s.parse::<u64>().ok());

    let len_bytes = resp1.bytes().await.map_err(|e| {
        FetchError::Http(format!("failed to read header length for {filename}: {e}"))
    })?;

    if len_bytes.len() < 8 {
        return Err(FetchError::SafetensorsHeader {
            filename: filename.to_owned(),
            reason: format!(
                "expected 8 bytes for length prefix, got {}",
                len_bytes.len()
            ),
        });
    }

    // INDEX: first 8 bytes guaranteed by length check above
    #[allow(clippy::indexing_slicing)]
    let header_size = u64::from_le_bytes([
        len_bytes[0],
        len_bytes[1],
        len_bytes[2],
        len_bytes[3],
        len_bytes[4],
        len_bytes[5],
        len_bytes[6],
        len_bytes[7],
    ]);

    // Request 2: JSON header.
    let range_end = 8u64.saturating_add(header_size).saturating_sub(1);
    let range_header = format!("bytes=8-{range_end}");
    let resp2 = client
        .get(url)
        // BORROW: explicit .as_str() instead of Deref coercion
        .header(reqwest::header::RANGE, range_header.as_str())
        .send()
        .await
        .map_err(|e| {
            FetchError::Http(format!("failed to fetch header JSON for {filename}: {e}"))
        })?;

    if !resp2.status().is_success() && resp2.status() != reqwest::StatusCode::PARTIAL_CONTENT {
        return Err(FetchError::Http(format!(
            "Range request for {filename} header JSON returned status {}",
            resp2.status()
        )));
    }

    let json_bytes = resp2
        .bytes()
        .await
        .map_err(|e| FetchError::Http(format!("failed to read header JSON for {filename}: {e}")))?;

    Ok((json_bytes.to_vec(), file_size))
}

// -----------------------------------------------------------------------
// Public API: single-file inspection
// -----------------------------------------------------------------------

/// Inspects a single `.safetensors` file's header (cache-first).
///
/// Checks the local HF cache first. If the file is cached, reads the header
/// from disk with zero network requests. Otherwise, falls back to two HTTP
/// Range requests (8-byte length prefix + JSON header). Does not download
/// tensor data in either case.
///
/// # Errors
///
/// Returns [`FetchError::Http`] if the Range requests fail.
/// Returns [`FetchError::SafetensorsHeader`] if the header is malformed.
pub async fn inspect_safetensors(
    repo_id: &str,
    filename: &str,
    token: Option<&str>,
    revision: Option<&str>,
) -> Result<(SafetensorsHeaderInfo, InspectSource), FetchError> {
    let rev = revision.unwrap_or("main");

    // Try local cache first.
    if let Some(cached_path) = resolve_cached_path(repo_id, rev, filename) {
        let info = inspect_safetensors_local(&cached_path)?;
        return Ok((info, InspectSource::Cached));
    }

    // Fall back to HTTP Range requests.
    let client = chunked::build_client(token)?;
    let url = chunked::build_download_url(repo_id, rev, filename);

    // BORROW: explicit .as_str() instead of Deref coercion
    let (json_bytes, file_size) = fetch_header_bytes(&client, url.as_str(), filename).await?;

    // CAST: usize → u64, JSON buffer length is always small
    #[allow(clippy::as_conversions)]
    let header_size = json_bytes.len() as u64;

    let (tensors, metadata) = parse_header_json(&json_bytes, filename)?;

    Ok((
        SafetensorsHeaderInfo {
            tensors,
            metadata,
            header_size,
            file_size,
            // Remote path doesn't go through anamnesis yet; quant detection
            // arrives in v0.11.1 when this path migrates to
            // `anamnesis::parse_safetensors_header_from_reader` over an
            // HTTP Range adapter.
            quant_info: None,
        },
        InspectSource::Remote,
    ))
}

/// Inspects a single `.safetensors` file from cache only.
///
/// Resolves the file in the local HF cache using the given `repo_id`,
/// `revision`, and `filename`. Returns an error if the file is not cached.
///
/// # Blocking I/O
///
/// Performs synchronous filesystem I/O; wrap in [`tokio::task::spawn_blocking`]
/// from async contexts. See [`inspect_safetensors_local`] for rationale.
///
/// # Errors
///
/// Returns [`FetchError::SafetensorsHeader`] if the file is not in the cache.
/// Returns [`FetchError::Io`] if the cached file cannot be read.
/// Returns [`FetchError::SafetensorsHeader`] if the header is malformed.
pub fn inspect_safetensors_cached(
    repo_id: &str,
    filename: &str,
    revision: Option<&str>,
) -> Result<SafetensorsHeaderInfo, FetchError> {
    let rev = revision.unwrap_or("main");

    let cached_path = resolve_cached_path(repo_id, rev, filename).ok_or_else(|| {
        FetchError::SafetensorsHeader {
            filename: filename.to_owned(),
            reason: format!("file not found in local cache for {repo_id} ({rev})"),
        }
    })?;

    inspect_safetensors_local(&cached_path)
}

/// Inspects a `.gguf` file's metadata from the local `HuggingFace` cache.
///
/// Delegates to [`anamnesis::parse_gguf`] for the on-disk parse, then maps the
/// result into the format-agnostic [`SafetensorsHeaderInfo`] shape used by
/// hf-fm's existing render path. Tensor names, GGUF-native shape order, and
/// dtype name strings carry over directly; per-tensor `data_offsets` are
/// `(data_offset, data_offset + byte_len)` (with `byte_len = 0` for tensors
/// whose dtype has no known byte size in anamnesis yet).
///
/// **Naming note:** the returned type is still called [`SafetensorsHeaderInfo`]
/// in v0.10.x because renaming a public type is a breaking change; the
/// uniform-dispatch rename to a format-agnostic name is scheduled for v0.10.3
/// when the dispatcher extends across `.npz` / `.pth` (see the cache-management
/// roadmap). For now, treat the type name as "header / file-level inspect
/// info" regardless of format.
///
/// **Metadata surfacing:** the GGUF metadata table can contain very large
/// arrays (e.g. tokenizer.ggml.tokens with 50K+ entries). To keep `Metadata:`
/// rendering useful, this function surfaces *scalar* metadata values only —
/// strings, booleans, integers, floats — and skips arrays. The GGUF format
/// version is surfaced under the synthetic key `gguf.version`, the effective
/// alignment under `gguf.alignment`. The original `general.architecture`,
/// `general.name`, and friends pass through unchanged.
///
/// **Blocking I/O:** anamnesis's GGUF parser mmaps the file; this function is
/// synchronous and should be wrapped in [`tokio::task::spawn_blocking`] from
/// async contexts.
///
/// # Errors
///
/// Returns [`FetchError::SafetensorsHeader`] if the file is not in the cache.
/// Returns [`FetchError::SafetensorsHeader`] if anamnesis rejects the GGUF
/// file (malformed header, truncated tensor table, etc.).
pub fn inspect_gguf_cached(
    repo_id: &str,
    filename: &str,
    revision: Option<&str>,
) -> Result<SafetensorsHeaderInfo, FetchError> {
    let rev = revision.unwrap_or("main");

    let cached_path = resolve_cached_path(repo_id, rev, filename).ok_or_else(|| {
        FetchError::SafetensorsHeader {
            // BORROW: explicit .to_owned() for owned String in the error variant
            filename: filename.to_owned(),
            reason: format!("file not found in local cache for {repo_id} ({rev})"),
        }
    })?;

    let file_size = std::fs::metadata(&cached_path).ok().map(|m| m.len());

    let parsed =
        anamnesis::parse_gguf(&cached_path).map_err(|e| FetchError::SafetensorsHeader {
            // BORROW: explicit .to_owned() for owned String in the error variant
            filename: filename.to_owned(),
            reason: format!("failed to parse GGUF: {e}"),
        })?;

    let tensors: Vec<TensorInfo> = parsed
        .tensor_info()
        .iter()
        .map(|info| {
            let start = info.data_offset;
            let end = info.byte_len.map_or(start, |b| start.saturating_add(b));
            TensorInfo {
                // BORROW: explicit .clone() / .to_string() to materialise owned
                // String + Vec<usize> from anamnesis's borrowed metadata
                name: info.name.clone(),
                dtype: info.dtype.to_string(),
                shape: info.shape.clone(),
                data_offsets: (start, end),
            }
        })
        .collect();

    // Stringify scalar metadata only; skip arrays (potentially huge — e.g.
    // tokenizer.ggml.tokens). Add synthetic keys for the format version and
    // alignment so they appear in the `Metadata:` block.
    let mut metadata: HashMap<String, String> = parsed
        .metadata()
        .iter()
        // BORROW: explicit .clone() to materialise an owned String key from
        // the borrowed HashMap iteration
        .filter_map(|(k, v)| stringify_gguf_scalar(v).map(|s| (k.clone(), s)))
        .collect();
    // BORROW: explicit .to_owned() for owned String keys
    metadata.insert("gguf.version".to_owned(), parsed.version().to_string());
    metadata.insert("gguf.alignment".to_owned(), parsed.alignment().to_string());

    Ok(SafetensorsHeaderInfo {
        tensors,
        metadata: Some(metadata),
        // GGUF has no discrete "header size" like safetensors's
        // u64-length-prefix + JSON. The value is left at 0 here; consumers
        // that care can derive an approximation from `file_size` minus the
        // tensor byte sum. The `Metadata:` block's `gguf.version` /
        // `gguf.alignment` keys surface the equivalent format-level info.
        header_size: 0,
        file_size,
        // GGUF quant info (Q4_K_M etc.) is implicit in per-tensor dtypes;
        // the v0.10.3 Phase C `Format:` / `Size:` lines are safetensors-only.
        quant_info: None,
    })
}

/// Inspects a `.npz` file's metadata from the local `HuggingFace` cache.
///
/// Delegates to [`anamnesis::inspect_npz`] for the on-disk parse (which
/// reads only the ZIP central directory + per-entry NPY headers — no
/// tensor data), then maps the result into the format-agnostic
/// [`SafetensorsHeaderInfo`] shape used by hf-fm's existing render path.
///
/// **Synthesised offsets.** Anamnesis exposes per-tensor `byte_len` but
/// not on-disk byte offsets (NPZ tensors live inside a ZIP archive;
/// offsets are not part of the inspect surface). hf-fm's
/// `TensorInfo::data_offsets` is synthesised as cumulative `(start, end)`
/// pairs — `start = sum of previous byte_lens` — so `byte_len()`
/// (= `end - start`) renders the actual storage size. The synthetic
/// offsets are NOT on-disk truth; consumers that need real offsets
/// should defer to v0.11.x once remote-inspect lands.
///
/// **Metadata.** `metadata: None` — NPZ has no metadata block analogous
/// to safetensors's `__metadata__` or GGUF's KV table.
///
/// **Header size.** Always `0` — NPZ has no discrete header analogous
/// to safetensors's `u64`-length-prefix + JSON. Mirrors the GGUF convention.
///
/// **Blocking I/O:** anamnesis's NPZ parser opens the file with
/// `std::fs::File`; this function is synchronous and should be wrapped
/// in [`tokio::task::spawn_blocking`] from async contexts.
///
/// # Errors
///
/// Returns [`FetchError::SafetensorsHeader`] if the file is not in the cache.
/// Returns [`FetchError::SafetensorsHeader`] if anamnesis rejects the NPZ
/// file (malformed ZIP central directory, unsupported NPY dtype, etc.).
pub fn inspect_npz_cached(
    repo_id: &str,
    filename: &str,
    revision: Option<&str>,
) -> Result<SafetensorsHeaderInfo, FetchError> {
    let rev = revision.unwrap_or("main");

    let cached_path = resolve_cached_path(repo_id, rev, filename).ok_or_else(|| {
        FetchError::SafetensorsHeader {
            // BORROW: explicit .to_owned() for owned String in the error variant
            filename: filename.to_owned(),
            reason: format!("file not found in local cache for {repo_id} ({rev})"),
        }
    })?;

    let file_size = std::fs::metadata(&cached_path).ok().map(|m| m.len());

    let parsed =
        anamnesis::inspect_npz(&cached_path).map_err(|e| FetchError::SafetensorsHeader {
            // BORROW: explicit .to_owned() for owned String in the error variant
            filename: filename.to_owned(),
            reason: format!("failed to parse NPZ: {e}"),
        })?;

    let mut tensors: Vec<TensorInfo> = Vec::with_capacity(parsed.tensors.len());
    let mut cursor: u64 = 0;
    for t in parsed.tensors {
        // CAST: usize → u64, byte_len fits in u64 by definition (in-memory size).
        #[allow(clippy::as_conversions)]
        let len = t.byte_len as u64;
        let start = cursor;
        let end = cursor.saturating_add(len);
        cursor = end;
        tensors.push(TensorInfo {
            name: t.name,
            // BORROW: explicit .to_string() — anamnesis `NpzDtype` enum → owned `String`
            dtype: t.dtype.to_string(),
            shape: t.shape,
            data_offsets: (start, end),
        });
    }

    Ok(SafetensorsHeaderInfo {
        tensors,
        metadata: None,
        header_size: 0,
        file_size,
        // NPZ has no quant-method metadata; quant_info stays None.
        quant_info: None,
    })
}

/// Inspects a `.pth` file's metadata from the local `HuggingFace` cache.
///
/// Delegates to [`anamnesis::parse_pth`] for the on-disk parse, then uses
/// the metadata-only `ParsedPth::tensor_info()` view (new in anamnesis
/// `0.5.0`) to enumerate `(name, shape, dtype, byte_len)` per tensor — no
/// further I/O beyond the initial mmap. The earlier `.tensors()` method
/// would materialise each tensor's data via `Cow<'a, [u8]>`, which is
/// unnecessary for inspect-only use.
///
/// **Synthesised offsets.** As with NPZ, anamnesis exposes per-tensor
/// `byte_len` but not on-disk byte offsets (PTH tensors live inside a
/// ZIP archive; offsets are not part of the inspect surface). hf-fm's
/// `TensorInfo::data_offsets` is synthesised as cumulative `(start, end)`
/// pairs so `byte_len()` (= `end - start`) renders the actual storage size.
///
/// **Metadata.** `metadata: None` — PTH has no metadata block analogous
/// to safetensors's `__metadata__` or GGUF's KV table. The format-level
/// `big_endian` flag (rare, near-always `false`) is not surfaced here;
/// can be added as a synthetic `pth.big_endian` key in a future patch if
/// real users request it.
///
/// **Header size.** Always `0` — PTH has no discrete header analogous
/// to safetensors's `u64`-length-prefix + JSON. Mirrors the GGUF / NPZ
/// convention.
///
/// **Blocking I/O:** anamnesis's PTH parser mmaps the file; this function
/// is synchronous and should be wrapped in [`tokio::task::spawn_blocking`]
/// from async contexts.
///
/// # Errors
///
/// Returns [`FetchError::SafetensorsHeader`] if the file is not in the cache.
/// Returns [`FetchError::SafetensorsHeader`] if anamnesis rejects the PTH
/// file (malformed pickle stream, legacy pre-1.6 raw-pickle format,
/// unsupported tensor dtype, etc.).
pub fn inspect_pth_cached(
    repo_id: &str,
    filename: &str,
    revision: Option<&str>,
) -> Result<SafetensorsHeaderInfo, FetchError> {
    let rev = revision.unwrap_or("main");

    let cached_path = resolve_cached_path(repo_id, rev, filename).ok_or_else(|| {
        FetchError::SafetensorsHeader {
            // BORROW: explicit .to_owned() for owned String in the error variant
            filename: filename.to_owned(),
            reason: format!("file not found in local cache for {repo_id} ({rev})"),
        }
    })?;

    let file_size = std::fs::metadata(&cached_path).ok().map(|m| m.len());

    let parsed = anamnesis::parse_pth(&cached_path).map_err(|e| FetchError::SafetensorsHeader {
        // BORROW: explicit .to_owned() for owned String in the error variant
        filename: filename.to_owned(),
        reason: format!("failed to parse PTH: {e}"),
    })?;

    let pth_tensors = parsed.tensor_info();
    let mut tensors: Vec<TensorInfo> = Vec::with_capacity(pth_tensors.len());
    let mut cursor: u64 = 0;
    for t in pth_tensors {
        // CAST: usize → u64, byte_len fits in u64 by definition (in-memory size).
        #[allow(clippy::as_conversions)]
        let len = t.byte_len as u64;
        let start = cursor;
        let end = cursor.saturating_add(len);
        cursor = end;
        tensors.push(TensorInfo {
            name: t.name,
            // BORROW: explicit .to_string() — anamnesis `PthDtype` enum → owned `String`
            dtype: t.dtype.to_string(),
            shape: t.shape,
            data_offsets: (start, end),
        });
    }

    Ok(SafetensorsHeaderInfo {
        tensors,
        metadata: None,
        header_size: 0,
        file_size,
        // PTH has no quant-method metadata; quant_info stays None.
        quant_info: None,
    })
}

/// Stringifies a scalar `GgufMetadataValue` from anamnesis.
///
/// Returns `None` for array variants (potentially huge — vocab tables, merges
/// lists) and for any future `#[non_exhaustive]` variants we don't yet
/// recognise. Surfaced through the `Metadata:` block in `inspect` output by
/// [`inspect_gguf_cached`].
//
// `GgufMetadataValue` is `#[non_exhaustive]`. The explicit `V::Array(_)` arm
// and the `_ =>` catch-all both return `None`, but they document different
// intents — "array variants are deliberately skipped" vs "future unknown
// variants fall through". Clippy's `match_same_arms` flags the bodies as
// identical; the duplication is intentional.
#[allow(clippy::match_same_arms)]
fn stringify_gguf_scalar(value: &anamnesis::parse::gguf::GgufMetadataValue) -> Option<String> {
    use anamnesis::parse::gguf::GgufMetadataValue as V;
    match value {
        V::String(s) => Some(s.clone()),
        V::Bool(b) => Some(b.to_string()),
        V::U8(n) => Some(n.to_string()),
        V::I8(n) => Some(n.to_string()),
        V::U16(n) => Some(n.to_string()),
        V::I16(n) => Some(n.to_string()),
        V::U32(n) => Some(n.to_string()),
        V::I32(n) => Some(n.to_string()),
        V::U64(n) => Some(n.to_string()),
        V::I64(n) => Some(n.to_string()),
        V::F32(n) => Some(format!("{n}")),
        V::F64(n) => Some(format!("{n}")),
        V::Array(_) => None,
        _ => None,
    }
}

// -----------------------------------------------------------------------
// Public API: multi-file inspection
// -----------------------------------------------------------------------

/// Inspects all `.safetensors` files in a repository (cache-first per file).
///
/// Fetches the file listing via `list_repo_files_with_metadata()`, then
/// inspects each `.safetensors` file's header via [`inspect_safetensors()`].
/// For each file, checks the local cache first and only makes HTTP Range
/// requests on cache miss. Returns full per-shard headers in filename order.
///
/// For a lightweight summary of sharded models (tensor counts per shard
/// without fetching individual headers), use [`fetch_shard_index()`] instead.
///
/// # Errors
///
/// Returns [`FetchError::Http`] if the metadata or Range requests fail.
pub async fn inspect_repo_safetensors(
    repo_id: &str,
    token: Option<&str>,
    revision: Option<&str>,
) -> Result<Vec<(String, SafetensorsHeaderInfo, InspectSource)>, FetchError> {
    let client = crate::chunked::build_client(token)?;
    let files =
        crate::repo::list_repo_files_with_metadata(repo_id, token, revision, &client).await?;

    let safetensors_files: Vec<String> = files
        .into_iter()
        .filter(|f| f.filename.ends_with(".safetensors"))
        .map(|f| f.filename)
        .collect();

    if safetensors_files.is_empty() {
        return Ok(Vec::new());
    }

    let semaphore = std::sync::Arc::new(tokio::sync::Semaphore::new(4));
    let mut join_set = JoinSet::new();

    for filename in safetensors_files {
        // BORROW: explicit .clone()/.to_owned() to move into async task
        let sem = semaphore.clone();
        let repo = repo_id.to_owned();
        let tok = token.map(str::to_owned);
        let rev = revision.map(str::to_owned);

        join_set.spawn(async move {
            let _permit = sem
                .acquire()
                .await
                .map_err(|e| FetchError::Http(format!("semaphore error: {e}")))?;
            // BORROW: explicit .as_deref() for Option<String> → Option<&str>
            let (info, source) =
                inspect_safetensors(&repo, &filename, tok.as_deref(), rev.as_deref()).await?;
            Ok::<_, FetchError>((filename, info, source))
        });
    }

    let mut results = Vec::new();
    while let Some(join_result) = join_set.join_next().await {
        match join_result {
            Ok(Ok(item)) => results.push(item),
            Ok(Err(e)) => {
                join_set.abort_all();
                return Err(e);
            }
            Err(e) => {
                join_set.abort_all();
                return Err(FetchError::Http(format!("task join error: {e}")));
            }
        }
    }

    results.sort_by(|a, b| a.0.cmp(&b.0));

    Ok(results)
}

/// Tensor-file extensions the `inspect` dispatcher understands.
///
/// Single source of truth shared by the cached listing
/// ([`list_cached_tensor_files`]) and the CLI's remote listing / numeric
/// index / `--pick` candidate set. Matches the v0.10.3 per-file dispatch in
/// `hf-fm inspect` (`.safetensors` remote or cached; `.gguf` / `.npz` /
/// `.pth` cached-only until the v0.11 remote-inspect line).
pub const SUPPORTED_TENSOR_EXTENSIONS: [&str; 4] = ["safetensors", "gguf", "npz", "pth"];

/// Returns `true` when `filename`'s extension matches one of
/// [`SUPPORTED_TENSOR_EXTENSIONS`] (case-insensitive).
#[must_use]
pub fn is_supported_tensor_file(filename: &str) -> bool {
    Path::new(filename)
        .extension()
        .and_then(|e| e.to_str())
        .is_some_and(|ext| {
            SUPPORTED_TENSOR_EXTENSIONS
                .iter()
                .any(|supported| ext.eq_ignore_ascii_case(supported))
        })
}

/// A `(filename, size_bytes)` enumeration of tensor files in a repo,
/// paired with the commit SHA of the resolved revision (when known).
///
/// The same tuple shape serves both local and remote listings:
/// [`list_cached_tensor_files`] produces it from a cached snapshot;
/// `repo::list_repo_files_with_commit` filtered through
/// [`is_supported_tensor_file`] produces it from the `HuggingFace` API.
/// Callers that need a uniform view over "what tensor files can I inspect?"
/// regardless of source use this alias.
pub type TensorFileListing = (Vec<(String, u64)>, Option<String>);

/// Alias kept for pre-v0.10.5 callers; [`list_cached_safetensors`] returns it.
///
/// Same tuple shape as [`TensorFileListing`], restricted by convention to
/// `.safetensors` entries.
pub type SafetensorsListing = TensorFileListing;

/// Lists `.safetensors` files in the cached snapshot for `repo_id`@`revision`.
///
/// Returns `(entries, commit_sha)` where `entries` is a sorted list of
/// `(filename, size_bytes)` tuples, and `commit_sha` is the snapshot's commit
/// hash (same value stored in `refs/<revision>`). Returns empty lists when the
/// repo or revision is not cached. Unlike [`inspect_repo_safetensors_cached`],
/// this does **not** parse any headers — it is a cheap name-and-size enumeration
/// intended for discovery UI (e.g. `inspect --list --cached`).
///
/// # Blocking I/O
///
/// Performs a synchronous recursive directory walk with a `stat` call per
/// `.safetensors` entry. On local SSDs the cost is sub-millisecond; on
/// networked caches (NFS/CIFS) a large sharded repo can take seconds. Wrap
/// in [`tokio::task::spawn_blocking`] from async contexts.
///
/// # Errors
///
/// Returns [`FetchError::Io`] if the snapshot directory cannot be read.
pub fn list_cached_safetensors(
    repo_id: &str,
    revision: Option<&str>,
) -> Result<SafetensorsListing, FetchError> {
    list_cached_matching_files(repo_id, revision, |name| name.ends_with(".safetensors"))
}

/// Lists all supported tensor files in the cached snapshot for `repo_id`@`revision`.
///
/// Multi-format sibling of [`list_cached_safetensors`]: matches every
/// extension in [`SUPPORTED_TENSOR_EXTENSIONS`] (case-insensitive) instead
/// of `.safetensors` only. Returns `(entries, commit_sha)` where `entries`
/// is a sorted list of `(filename, size_bytes)` tuples, and `commit_sha` is
/// the snapshot's commit hash (same value stored in `refs/<revision>`).
/// Returns empty lists when the repo or revision is not cached. Does **not**
/// parse any headers — it is a cheap name-and-size enumeration intended for
/// discovery UI (e.g. `inspect --list --cached`, `inspect --pick --cached`).
///
/// # Blocking I/O
///
/// Performs a synchronous recursive directory walk with a `stat` call per
/// matching entry. On local SSDs the cost is sub-millisecond; on networked
/// caches (NFS/CIFS) a large sharded repo can take seconds. Wrap in
/// [`tokio::task::spawn_blocking`] from async contexts.
///
/// # Errors
///
/// Returns [`FetchError::Io`] if the snapshot directory cannot be read.
pub fn list_cached_tensor_files(
    repo_id: &str,
    revision: Option<&str>,
) -> Result<TensorFileListing, FetchError> {
    list_cached_matching_files(repo_id, revision, is_supported_tensor_file)
}

/// Shared body of [`list_cached_safetensors`] / [`list_cached_tensor_files`]:
/// resolves the snapshot directory and walks it with the given filename
/// predicate.
fn list_cached_matching_files(
    repo_id: &str,
    revision: Option<&str>,
    matches: fn(&str) -> bool,
) -> Result<TensorFileListing, FetchError> {
    let rev = revision.unwrap_or("main");
    let cache_dir = cache::hf_cache_dir()?;
    let repo_dir = cache_layout::repo_dir(&cache_dir, repo_id);

    let Some(commit_hash) = cache::read_ref(&repo_dir, rev) else {
        return Ok((Vec::new(), None));
    };

    let snapshot_dir = cache_layout::snapshot_dir(&repo_dir, &commit_hash);
    if !snapshot_dir.exists() {
        return Ok((Vec::new(), Some(commit_hash)));
    }

    let mut results = Vec::new();
    collect_matching_names_sizes(&snapshot_dir, "", matches, &mut results)?;
    results.sort_by(|a, b| a.0.cmp(&b.0));
    Ok((results, Some(commit_hash)))
}

/// Recursively collects `(filename, size)` pairs for files whose bare
/// entry name satisfies `matches` (extension predicates need no prefix).
fn collect_matching_names_sizes(
    dir: &Path,
    prefix: &str,
    matches: fn(&str) -> bool,
    results: &mut Vec<(String, u64)>,
) -> Result<(), FetchError> {
    let entries = std::fs::read_dir(dir).map_err(|e| FetchError::Io {
        path: dir.to_path_buf(),
        source: e,
    })?;

    for entry in entries {
        let Ok(entry) = entry else { continue };
        let path = entry.path();
        // BORROW: explicit .to_string_lossy() for OsString → str conversion
        let name = entry.file_name().to_string_lossy().to_string();

        if path.is_dir() {
            let child_prefix = if prefix.is_empty() {
                name
            } else {
                format!("{prefix}/{name}")
            };
            collect_matching_names_sizes(&path, &child_prefix, matches, results)?;
        } else if matches(&name) {
            let filename = if prefix.is_empty() {
                name
            } else {
                format!("{prefix}/{name}")
            };
            let size = entry.metadata().map_or(0, |m| m.len());
            results.push((filename, size));
        }
    }

    Ok(())
}

/// Inspects all `.safetensors` files in a cached repository (no network).
///
/// Walks the snapshot directory and inspects each `.safetensors` file's
/// header from local disk. Returns results in filename order.
///
/// # Blocking I/O
///
/// Walks the snapshot directory and reads each header synchronously. In async
/// contexts, wrap in [`tokio::task::spawn_blocking`] to avoid stalling the
/// runtime — multi-shard repos on network-mounted caches can take seconds.
///
/// # Errors
///
/// Returns [`FetchError::Io`] if the cache directory cannot be read.
/// Returns [`FetchError::SafetensorsHeader`] if any header is malformed.
pub fn inspect_repo_safetensors_cached(
    repo_id: &str,
    revision: Option<&str>,
) -> Result<Vec<(String, SafetensorsHeaderInfo)>, FetchError> {
    let rev = revision.unwrap_or("main");
    let cache_dir = cache::hf_cache_dir()?;
    let repo_dir = cache_layout::repo_dir(&cache_dir, repo_id);

    let Some(commit_hash) = cache::read_ref(&repo_dir, rev) else {
        return Ok(Vec::new());
    };

    let snapshot_dir = cache_layout::snapshot_dir(&repo_dir, &commit_hash);
    if !snapshot_dir.exists() {
        return Ok(Vec::new());
    }

    let mut results = Vec::new();
    collect_safetensors_recursive(&snapshot_dir, "", &mut results)?;
    results.sort_by(|a, b| a.0.cmp(&b.0));

    Ok(results)
}

/// Recursively finds and inspects `.safetensors` files in a snapshot directory.
fn collect_safetensors_recursive(
    dir: &Path,
    prefix: &str,
    results: &mut Vec<(String, SafetensorsHeaderInfo)>,
) -> Result<(), FetchError> {
    let entries = std::fs::read_dir(dir).map_err(|e| FetchError::Io {
        path: dir.to_path_buf(),
        source: e,
    })?;

    for entry in entries {
        let Ok(entry) = entry else { continue };
        let path = entry.path();
        // BORROW: explicit .to_string_lossy() for OsString → str conversion
        let name = entry.file_name().to_string_lossy().to_string();

        if path.is_dir() {
            let child_prefix = if prefix.is_empty() {
                name
            } else {
                format!("{prefix}/{name}")
            };
            collect_safetensors_recursive(&path, &child_prefix, results)?;
        } else if name.ends_with(".safetensors") {
            let filename = if prefix.is_empty() {
                name
            } else {
                format!("{prefix}/{name}")
            };
            let info = inspect_safetensors_local(&path)?;
            results.push((filename, info));
        }
    }

    Ok(())
}

// -----------------------------------------------------------------------
// Shard index
// -----------------------------------------------------------------------

/// Raw JSON structure of `model.safetensors.index.json`.
#[derive(serde::Deserialize)]
struct RawShardIndex {
    weight_map: HashMap<String, String>,
    #[serde(default)]
    metadata: Option<HashMap<String, serde_json::Value>>,
}

/// Fetches and parses the shard index for a sharded `.safetensors` model (cache-first).
///
/// Returns `Ok(None)` if the repo has no `model.safetensors.index.json` (i.e.,
/// the model is not sharded or uses a single `.safetensors` file).
///
/// # Errors
///
/// Returns [`FetchError::Http`] if the index fetch fails.
/// Returns [`FetchError::SafetensorsHeader`] if the index JSON is malformed.
pub async fn fetch_shard_index(
    repo_id: &str,
    token: Option<&str>,
    revision: Option<&str>,
) -> Result<Option<ShardedIndex>, FetchError> {
    let rev = revision.unwrap_or("main");
    let index_filename = "model.safetensors.index.json";

    // Try local cache first.
    if let Some(cached_path) = resolve_cached_path(repo_id, rev, index_filename) {
        let content = std::fs::read_to_string(&cached_path).map_err(|e| FetchError::Io {
            path: cached_path,
            source: e,
        })?;
        let index = parse_shard_index_json(&content, repo_id)?;
        return Ok(Some(index));
    }

    // Fall back to HTTP.
    let client = chunked::build_client(token)?;
    let url = chunked::build_download_url(repo_id, rev, index_filename);

    // BORROW: explicit .as_str() instead of Deref coercion
    let response =
        client.get(url.as_str()).send().await.map_err(|e| {
            FetchError::Http(format!("failed to fetch shard index for {repo_id}: {e}"))
        })?;

    if response.status() == reqwest::StatusCode::NOT_FOUND {
        return Ok(None);
    }

    if !response.status().is_success() {
        return Err(FetchError::Http(format!(
            "shard index request for {repo_id} returned status {}",
            response.status()
        )));
    }

    let content = response
        .text()
        .await
        .map_err(|e| FetchError::Http(format!("failed to read shard index for {repo_id}: {e}")))?;

    let index = parse_shard_index_json(&content, repo_id)?;
    Ok(Some(index))
}

/// Fetches the shard index from cache only (no network).
///
/// Returns `Ok(None)` if the index file is not cached.
///
/// # Errors
///
/// Returns [`FetchError::Io`] if the cached file cannot be read.
/// Returns [`FetchError::SafetensorsHeader`] if the index JSON is malformed.
pub fn fetch_shard_index_cached(
    repo_id: &str,
    revision: Option<&str>,
) -> Result<Option<ShardedIndex>, FetchError> {
    let rev = revision.unwrap_or("main");
    let index_filename = "model.safetensors.index.json";

    let Some(cached_path) = resolve_cached_path(repo_id, rev, index_filename) else {
        return Ok(None);
    };

    let content = std::fs::read_to_string(&cached_path).map_err(|e| FetchError::Io {
        path: cached_path,
        source: e,
    })?;

    let index = parse_shard_index_json(&content, repo_id)?;
    Ok(Some(index))
}

/// Parses shard index JSON into a `ShardedIndex`.
fn parse_shard_index_json(content: &str, repo_id: &str) -> Result<ShardedIndex, FetchError> {
    let raw: RawShardIndex =
        serde_json::from_str(content).map_err(|e| FetchError::SafetensorsHeader {
            filename: "model.safetensors.index.json".to_owned(),
            reason: format!("failed to parse shard index for {repo_id}: {e}"),
        })?;

    // Collect unique shard filenames in sorted order.
    let mut shard_set: Vec<String> = raw.weight_map.values().cloned().collect();
    shard_set.sort();
    shard_set.dedup();

    Ok(ShardedIndex {
        weight_map: raw.weight_map,
        shards: shard_set,
        metadata: raw.metadata,
    })
}

// -----------------------------------------------------------------------
// Param formatting helper
// -----------------------------------------------------------------------

/// Formats a parameter count with a compact suffix (e.g., `927.0M`, `1.02B`).
#[must_use]
pub fn format_params(count: u64) -> String {
    // CAST: u64 → f64, precision loss acceptable; value is a display-only scalar
    #[allow(clippy::cast_precision_loss, clippy::as_conversions)]
    let val = count as f64;

    if count >= 1_000_000_000 {
        format!("{:.2}B", val / 1_000_000_000.0)
    } else if count >= 1_000_000 {
        format!("{:.1}M", val / 1_000_000.0)
    } else if count >= 1_000 {
        format!("{:.1}K", val / 1_000.0)
    } else {
        count.to_string()
    }
}

// -----------------------------------------------------------------------
// Adapter config
// -----------------------------------------------------------------------

/// Raw JSON structure of `adapter_config.json`.
#[derive(serde::Deserialize)]
struct RawAdapterConfig {
    #[serde(default)]
    peft_type: Option<String>,
    #[serde(default)]
    base_model_name_or_path: Option<String>,
    #[serde(default)]
    r: Option<u32>,
    #[serde(default)]
    lora_alpha: Option<f64>,
    #[serde(default)]
    target_modules: Option<AdapterTargetModules>,
    #[serde(default)]
    task_type: Option<String>,
}

/// `target_modules` in adapter configs can be a list of strings or a single string.
#[derive(serde::Deserialize)]
#[serde(untagged)]
enum AdapterTargetModules {
    /// A list of module name strings.
    List(Vec<String>),
    /// A single module name string.
    Single(String),
}

/// Fetches and parses `adapter_config.json` for a `PEFT` adapter repository (cache-first).
///
/// Returns `Ok(None)` if the file does not exist (HTTP 404), meaning the
/// repository is not a `PEFT` adapter.
///
/// # Errors
///
/// Returns [`FetchError::Http`] if the request fails (other than 404).
/// Returns [`FetchError::SafetensorsHeader`] if the JSON is malformed.
pub async fn fetch_adapter_config(
    repo_id: &str,
    token: Option<&str>,
    revision: Option<&str>,
) -> Result<Option<AdapterConfig>, FetchError> {
    let rev = revision.unwrap_or("main");
    let config_filename = "adapter_config.json";

    // Try local cache first.
    if let Some(cached_path) = resolve_cached_path(repo_id, rev, config_filename) {
        let content = std::fs::read_to_string(&cached_path).map_err(|e| FetchError::Io {
            path: cached_path,
            source: e,
        })?;
        let config = parse_adapter_config_json(&content, repo_id)?;
        return Ok(Some(config));
    }

    // Fall back to HTTP.
    let client = chunked::build_client(token)?;
    let url = chunked::build_download_url(repo_id, rev, config_filename);

    // BORROW: explicit .as_str() instead of Deref coercion
    let response = client.get(url.as_str()).send().await.map_err(|e| {
        FetchError::Http(format!("failed to fetch adapter config for {repo_id}: {e}"))
    })?;

    if response.status() == reqwest::StatusCode::NOT_FOUND {
        return Ok(None);
    }

    if !response.status().is_success() {
        return Err(FetchError::Http(format!(
            "adapter config request for {repo_id} returned status {}",
            response.status()
        )));
    }

    let content = response.text().await.map_err(|e| {
        FetchError::Http(format!("failed to read adapter config for {repo_id}: {e}"))
    })?;

    let config = parse_adapter_config_json(&content, repo_id)?;
    Ok(Some(config))
}

/// Fetches the adapter config from cache only (no network).
///
/// Returns `Ok(None)` if the file is not cached.
///
/// # Errors
///
/// Returns [`FetchError::Io`] if the cached file cannot be read.
/// Returns [`FetchError::SafetensorsHeader`] if the JSON is malformed.
pub fn fetch_adapter_config_cached(
    repo_id: &str,
    revision: Option<&str>,
) -> Result<Option<AdapterConfig>, FetchError> {
    let rev = revision.unwrap_or("main");
    let config_filename = "adapter_config.json";

    let Some(cached_path) = resolve_cached_path(repo_id, rev, config_filename) else {
        return Ok(None);
    };

    let content = std::fs::read_to_string(&cached_path).map_err(|e| FetchError::Io {
        path: cached_path,
        source: e,
    })?;

    let config = parse_adapter_config_json(&content, repo_id)?;
    Ok(Some(config))
}

/// Parses adapter config JSON into an [`AdapterConfig`].
fn parse_adapter_config_json(content: &str, repo_id: &str) -> Result<AdapterConfig, FetchError> {
    let raw: RawAdapterConfig =
        serde_json::from_str(content).map_err(|e| FetchError::SafetensorsHeader {
            filename: "adapter_config.json".to_owned(),
            reason: format!("failed to parse adapter config for {repo_id}: {e}"),
        })?;

    let target_modules = match raw.target_modules {
        Some(AdapterTargetModules::List(v)) => v,
        Some(AdapterTargetModules::Single(s)) => vec![s],
        None => Vec::new(),
    };

    Ok(AdapterConfig {
        peft_type: raw.peft_type,
        base_model_name_or_path: raw.base_model_name_or_path,
        r: raw.r,
        lora_alpha: raw.lora_alpha,
        target_modules,
        task_type: raw.task_type,
    })
}

/// Raw JSON structure of `config.json` (only the fields hf-fm reads for
/// KV-cache budgeting).
///
/// Serde aliases absorb the legacy GPT-NeoX / Falcon spellings. `text_config`
/// holds the nested language-model config of multimodal repos (Gemma-3) and
/// is used as a fallback when the attention dims are absent at top level.
#[derive(serde::Deserialize)]
struct RawModelConfig {
    #[serde(default)]
    model_type: Option<String>,
    #[serde(default, alias = "n_layer")]
    num_hidden_layers: Option<u32>,
    #[serde(default, alias = "n_head")]
    num_attention_heads: Option<u32>,
    #[serde(default, alias = "num_kv_heads", alias = "n_head_kv")]
    num_key_value_heads: Option<u32>,
    #[serde(default, alias = "attention_head_dim")]
    head_dim: Option<u32>,
    #[serde(default)]
    hidden_size: Option<u32>,
    #[serde(default)]
    torch_dtype: Option<String>,
    #[serde(default)]
    sliding_window: Option<u32>,
    #[serde(default)]
    sliding_window_pattern: Option<u32>,
    #[serde(default)]
    use_sliding_window: Option<bool>,
    #[serde(default)]
    kv_lora_rank: Option<u32>,
    #[serde(default)]
    qk_rope_head_dim: Option<u32>,
    #[serde(default)]
    layer_types: Option<Vec<String>>,
    #[serde(default)]
    hybrid_override_pattern: Option<String>,
    #[serde(default)]
    attn_layer_indices: Option<Vec<u32>>,
    #[serde(default)]
    full_attention_interval: Option<u32>,
    #[serde(default, alias = "mamba_num_heads")]
    mamba_n_heads: Option<u32>,
    #[serde(default, alias = "mamba_head_dim")]
    mamba_d_head: Option<u32>,
    #[serde(default, alias = "ssm_state_size")]
    mamba_d_state: Option<u32>,
    #[serde(default, alias = "conv_kernel")]
    mamba_d_conv: Option<u32>,
    #[serde(default, alias = "n_groups")]
    mamba_n_groups: Option<u32>,
    #[serde(default)]
    text_config: Option<Box<RawModelConfig>>,
}

/// Lowers a [`RawModelConfig`] into the public [`ModelConfig`].
///
/// Multimodal configs (Gemma-3) nest the language-model dims under
/// `text_config`; when the top level carries no attention dims, this recurses
/// into that nested config so the KV estimator sees the real numbers.
fn model_config_from_raw(raw: RawModelConfig) -> ModelConfig {
    if raw.num_hidden_layers.is_none()
        && raw.num_attention_heads.is_none()
        && raw.hidden_size.is_none()
    {
        if let Some(text) = raw.text_config {
            return model_config_from_raw(*text);
        }
    }

    ModelConfig {
        model_type: raw.model_type,
        num_hidden_layers: raw.num_hidden_layers,
        num_attention_heads: raw.num_attention_heads,
        num_key_value_heads: raw.num_key_value_heads,
        head_dim: raw.head_dim,
        hidden_size: raw.hidden_size,
        torch_dtype: raw.torch_dtype,
        sliding_window: raw.sliding_window,
        sliding_window_pattern: raw.sliding_window_pattern,
        use_sliding_window: raw.use_sliding_window,
        kv_lora_rank: raw.kv_lora_rank,
        qk_rope_head_dim: raw.qk_rope_head_dim,
        layer_types: raw.layer_types,
        hybrid_override_pattern: raw.hybrid_override_pattern,
        attn_layer_indices: raw.attn_layer_indices,
        full_attention_interval: raw.full_attention_interval,
        mamba_n_heads: raw.mamba_n_heads,
        mamba_d_head: raw.mamba_d_head,
        mamba_d_state: raw.mamba_d_state,
        mamba_d_conv: raw.mamba_d_conv,
        mamba_n_groups: raw.mamba_n_groups,
    }
}

/// Parses `config.json` content into a [`ModelConfig`].
fn parse_model_config_json(content: &str, repo_id: &str) -> Result<ModelConfig, FetchError> {
    let raw: RawModelConfig =
        serde_json::from_str(content).map_err(|e| FetchError::SafetensorsHeader {
            filename: "config.json".to_owned(),
            reason: format!("failed to parse model config for {repo_id}: {e}"),
        })?;

    Ok(model_config_from_raw(raw))
}

/// Fetches and parses a model's `config.json` (cache-first, then HTTP).
///
/// Returns `Ok(None)` when the repository has no `config.json` (HTTP 404) —
/// e.g. a non-model repo or a raw-weights upload.
///
/// # Errors
///
/// Returns [`FetchError::Http`] if the request fails (other than 404).
/// Returns [`FetchError::Io`] if a cached `config.json` cannot be read.
/// Returns [`FetchError::SafetensorsHeader`] if the JSON is malformed.
pub async fn fetch_model_config(
    repo_id: &str,
    token: Option<&str>,
    revision: Option<&str>,
) -> Result<Option<ModelConfig>, FetchError> {
    let rev = revision.unwrap_or("main");
    let config_filename = "config.json";

    // Try local cache first.
    if let Some(cached_path) = resolve_cached_path(repo_id, rev, config_filename) {
        let content = std::fs::read_to_string(&cached_path).map_err(|e| FetchError::Io {
            path: cached_path,
            source: e,
        })?;
        let config = parse_model_config_json(&content, repo_id)?;
        return Ok(Some(config));
    }

    // Fall back to HTTP.
    let client = chunked::build_client(token)?;
    let url = chunked::build_download_url(repo_id, rev, config_filename);

    // BORROW: explicit .as_str() instead of Deref coercion
    let response = client.get(url.as_str()).send().await.map_err(|e| {
        FetchError::Http(format!("failed to fetch model config for {repo_id}: {e}"))
    })?;

    if response.status() == reqwest::StatusCode::NOT_FOUND {
        return Ok(None);
    }

    if !response.status().is_success() {
        return Err(FetchError::Http(format!(
            "model config request for {repo_id} returned status {}",
            response.status()
        )));
    }

    let content = response
        .text()
        .await
        .map_err(|e| FetchError::Http(format!("failed to read model config for {repo_id}: {e}")))?;

    let config = parse_model_config_json(&content, repo_id)?;
    Ok(Some(config))
}

/// Fetches a model's `config.json` from the local cache only (no network).
///
/// Returns `Ok(None)` if the file is not cached.
///
/// # Errors
///
/// Returns [`FetchError::Io`] if the cached file cannot be read.
/// Returns [`FetchError::SafetensorsHeader`] if the JSON is malformed.
pub fn fetch_model_config_cached(
    repo_id: &str,
    revision: Option<&str>,
) -> Result<Option<ModelConfig>, FetchError> {
    let rev = revision.unwrap_or("main");
    let config_filename = "config.json";

    let Some(cached_path) = resolve_cached_path(repo_id, rev, config_filename) else {
        return Ok(None);
    };

    let content = std::fs::read_to_string(&cached_path).map_err(|e| FetchError::Io {
        path: cached_path,
        source: e,
    })?;

    let config = parse_model_config_json(&content, repo_id)?;
    Ok(Some(config))
}

#[cfg(test)]
mod tests {
    #![allow(clippy::panic)]

    use super::is_supported_tensor_file;

    #[test]
    fn is_supported_tensor_file_accepts_all_four_formats() {
        assert!(is_supported_tensor_file("model.safetensors"));
        assert!(is_supported_tensor_file("model.gguf"));
        assert!(is_supported_tensor_file("params.npz"));
        assert!(is_supported_tensor_file("weights.pth"));
    }

    #[test]
    fn is_supported_tensor_file_is_case_insensitive_on_extension() {
        assert!(is_supported_tensor_file("MODEL.SAFETENSORS"));
        assert!(is_supported_tensor_file("model.GGUF"));
    }

    #[test]
    fn is_supported_tensor_file_handles_nested_paths() {
        assert!(is_supported_tensor_file(
            "transformer/demonCORESFWNSFW_fluxV13.safetensors"
        ));
    }

    #[test]
    fn is_supported_tensor_file_rejects_other_extensions() {
        assert!(!is_supported_tensor_file("config.json"));
        assert!(!is_supported_tensor_file("model.bin"));
        assert!(!is_supported_tensor_file("archive.npy"));
        assert!(!is_supported_tensor_file("README.md"));
        assert!(!is_supported_tensor_file("no_extension"));
        // The extension must be the FINAL path segment suffix, not a
        // substring elsewhere in the name.
        assert!(!is_supported_tensor_file("model.safetensors.bak"));
    }
}