//! Load model weights from GGUF into dequantized f32 tensors.

use crate::gguf::{GgufError, GgufHeader, MappedWeights};
use crate::quant;
use crate::tensor::Tensor;
use thiserror::Error;

#[derive(Error, Debug)]
pub enum WeightError {
    #[error("Tensor not found in GGUF: {0}")]
    NotFound(String),
    #[error("Dequantization error: {0}")]
    Dequant(#[from] quant::QuantError),
    #[error("GGUF read error: {0}")]
    Gguf(#[from] GgufError),
}

/// Load a named tensor from GGUF, dequantize to f32.
pub fn load_tensor(
    name: &str,
    header: &GgufHeader,
    mapped: &MappedWeights,
) -> Result<Tensor, WeightError> {
    let info = header
        .tensors
        .iter()
        .find(|t| t.name == name)
        .ok_or_else(|| WeightError::NotFound(name.into()))?;

    let n_elements: usize = info.dimensions.iter().product::<u64>() as usize;
    let block_bytes = quant::block_size_bytes(info.dtype).unwrap_or(4);
    let block_elems = quant::block_elements(info.dtype).unwrap_or(1);
    let n_blocks = n_elements.div_ceil(block_elems);
    let byte_len = n_blocks * block_bytes;

    let raw = mapped.tensor_bytes(info.offset, byte_len)?;
    let data = quant::dequantize(raw, info.dtype, n_elements)?;

    // GGUF stores dimensions in column-major order [cols, rows].
    // Reverse to get row-major [rows, cols] for our tensor operations.
    let shape: Vec<usize> = info.dimensions.iter().rev().map(|&d| d as usize).collect();
    Ok(Tensor { data, shape })
}
