use serde::de::{MapAccess, Visitor};
use serde::{Deserialize, Deserializer};
use serde_json::Value;
use sha2::{Digest, Sha256};
use std::collections::{BTreeMap, BTreeSet, HashSet};
use std::env;
use std::fmt;
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::time::Duration;

const MAX_SAFETENSORS_HEADER: u64 = 16 * 1024 * 1024;
const CONVERTER: &str = "qwen38-artifacts-v1-identity";
const LICENSE: &str = "Qwen Community License 1.0";
const SOURCE_REPOSITORY: &str = "Youssofal/Qwen3.8-Flash-Next-MTPLX-Bare-Speed";
const SOURCE_REVISION: &str = "74559cdf34fbfc0b593de72d17e93f37fd4f9ea7";
const BASE_REPOSITORY: &str = "Qwen/Qwen3.8-Flash-Next";
const BASE_REVISION: &str = "de4b8e4d43b917e7706784d8bb445c9af86a3540";
const DOWNLOAD_TIMEOUT: Duration = Duration::from_secs(30);

#[derive(Deserialize)]
struct Manifest {
    format: u32,
    source: Source,
    tensor_inventory: String,
    tensor_inventory_sha256: String,
    files: Vec<Artifact>,
    config: BTreeMap<String, Value>,
    runtime: BTreeMap<String, Value>,
    representative_rows: Vec<Sample>,
    excluded: Vec<Excluded>,
}

#[derive(Deserialize)]
struct Source {
    repository: String,
    url: String,
    revision: String,
    base_repository: String,
    base_url: String,
    base_revision: String,
    license: String,
    converter: String,
}

#[derive(Deserialize)]
struct Artifact {
    path: String,
    role: String,
    size: u64,
    sha256: String,
}

#[derive(Deserialize)]
struct Excluded {
    path: String,
    reason: String,
    size: u64,
    sha256: String,
}

#[derive(Deserialize)]
struct Sample {
    class: String,
    file: String,
    tensor: String,
    row: u64,
    values: u64,
    bits: Option<u32>,
    group_size: Option<u64>,
    sha256: String,
}

fn main() {
    if let Err(error) = run() {
        eprintln!("qwen38-artifacts: {error}");
        std::process::exit(1);
    }
}

fn run() -> Result<(), String> {
    let args = env::args_os().skip(1).collect::<Vec<_>>();
    match args.as_slice() {
        [command, manifest, source] if command == "verify" => {
            let manifest = Path::new(manifest);
            verify(manifest, Path::new(source))?;
            println!("verified {}", manifest.display());
            Ok(())
        }
        [command, manifest, source, output] if command == "inventory" => {
            let (manifest, _) = load_manifest(Path::new(manifest))?;
            let inventory = build_inventory(&manifest, Path::new(source), false)?;
            fs::write(output, inventory).map_err(|error| error.to_string())
        }
        [command, manifest, source, output] if command == "materialize" => {
            materialize(Path::new(manifest), Path::new(source), Path::new(output))
        }
        [command, manifest, output] if command == "fetch" => {
            fetch(Path::new(manifest), Path::new(output))
        }
        [command, manifest, source] if command == "sample-digests" => {
            let (manifest, _) = load_manifest(Path::new(manifest))?;
            for sample in &manifest.representative_rows {
                println!("{}\t{}", sample.class, sample_digest(Path::new(source), sample)?);
            }
            Ok(())
        }
        _ => Err("usage: qwen38-artifacts fetch MANIFEST OUTPUT | verify MANIFEST SOURCE | inventory MANIFEST SOURCE OUTPUT | sample-digests MANIFEST SOURCE | materialize MANIFEST SOURCE OUTPUT".into()),
    }
}

fn load_manifest(path: &Path) -> Result<(Manifest, PathBuf), String> {
    let bytes = fs::read(path).map_err(|error| format!("{}: {error}", path.display()))?;
    let manifest: Manifest =
        serde_json::from_slice(&bytes).map_err(|error| format!("{}: {error}", path.display()))?;
    if manifest.format != 1 {
        return Err(format!("unsupported manifest format {}", manifest.format));
    }
    validate_source(&manifest.source)?;
    let directory = path
        .parent()
        .ok_or_else(|| format!("{} has no parent directory", path.display()))?;
    Ok((manifest, directory.to_owned()))
}

fn validate_source(source: &Source) -> Result<(), String> {
    for (name, value) in [
        ("source repository", source.repository.as_str()),
        ("source URL", source.url.as_str()),
        ("source revision", source.revision.as_str()),
        ("base repository", source.base_repository.as_str()),
        ("base URL", source.base_url.as_str()),
        ("base revision", source.base_revision.as_str()),
        ("license", source.license.as_str()),
        ("converter", source.converter.as_str()),
    ] {
        if value.is_empty() {
            return Err(format!("{name} is empty"));
        }
    }
    if source.revision.len() != 40 || source.base_revision.len() != 40 {
        return Err("source revisions must be full 40-character commits".into());
    }
    if source.repository != SOURCE_REPOSITORY
        || source.revision != SOURCE_REVISION
        || source.base_repository != BASE_REPOSITORY
        || source.base_revision != BASE_REVISION
    {
        return Err("manifest does not identify the pinned Bare Speed source and base".into());
    }
    if source.url
        != format!(
            "https://huggingface.co/{}/tree/{}",
            source.repository, source.revision
        )
        || source.base_url
            != format!(
                "https://huggingface.co/{}/tree/{}",
                source.base_repository, source.base_revision
            )
    {
        return Err("source URLs do not match their repositories and revisions".into());
    }
    if source.converter != CONVERTER || source.license != LICENSE {
        return Err("manifest converter or license does not match this tool".into());
    }
    Ok(())
}

fn verify(manifest_path: &Path, source: &Path) -> Result<(), String> {
    let (manifest, manifest_directory) = load_manifest(manifest_path)?;
    validate_file_contract(&manifest)?;
    for artifact in &manifest.files {
        verify_artifact(&source.join(&artifact.path), artifact)?;
    }
    validate_json_contract(source, "config.json", &manifest.config)?;
    validate_json_contract(source, "mtplx_runtime.json", &manifest.runtime)?;
    validate_index(&manifest, source)?;
    validate_samples(&manifest, source)?;
    let actual_inventory = build_inventory(&manifest, source, true)?;
    let inventory_path = manifest_directory.join(&manifest.tensor_inventory);
    let expected_inventory = fs::read_to_string(&inventory_path)
        .map_err(|error| format!("{}: {error}", inventory_path.display()))?;
    let inventory_sha256 = sha256_file(&inventory_path)?;
    if inventory_sha256 != manifest.tensor_inventory_sha256 {
        return Err(format!(
            "{} has SHA-256 {inventory_sha256}, expected {}",
            inventory_path.display(),
            manifest.tensor_inventory_sha256
        ));
    }
    if actual_inventory != expected_inventory {
        return Err(format!(
            "tensor inventory does not match {}",
            inventory_path.display()
        ));
    }
    Ok(())
}

fn validate_samples(manifest: &Manifest, source: &Path) -> Result<(), String> {
    let classes = manifest
        .representative_rows
        .iter()
        .map(|sample| sample.class.as_str())
        .collect::<BTreeSet<_>>();
    let required = ["dense", "expert", "gdn", "mtp", "ple", "qsa"]
        .into_iter()
        .collect::<BTreeSet<_>>();
    if classes != required || manifest.representative_rows.len() != required.len() {
        return Err("representative rows must cover dense, expert, GDN, MTP, PLE, and QSA".into());
    }
    let precision = precision_assignments(manifest, source, true)?;
    for sample in &manifest.representative_rows {
        if sample.values == 0 || sample.values > 4096 {
            return Err(format!("{} sample has an invalid width", sample.class));
        }
        let expected_role = match sample.class.as_str() {
            "ple" => "ple",
            "mtp" => "mtp",
            _ => "core",
        };
        if !manifest
            .files
            .iter()
            .any(|artifact| artifact.path == sample.file && artifact.role == expected_role)
        {
            return Err(format!(
                "{} sample is not in its artifact role",
                sample.class
            ));
        }
        match precision.get(&sample.tensor) {
            Some(contract)
                if sample.bits == Some(contract.bits)
                    && sample.group_size == Some(contract.group_size) => {}
            None if sample.bits.is_none() && sample.group_size.is_none() => {}
            _ => {
                return Err(format!(
                    "{} sample does not match its tensor precision",
                    sample.class
                ));
            }
        }
        let actual = sample_digest(source, sample)?;
        if actual != sample.sha256 {
            return Err(format!(
                "{} sample has SHA-256 {actual}, expected {}",
                sample.class, sample.sha256
            ));
        }
    }
    Ok(())
}

fn sample_digest(source: &Path, sample: &Sample) -> Result<String, String> {
    let path = source.join(&sample.file);
    let (data_start, tensors) = read_safetensors_header(&path, true)?;
    let tensor = tensors
        .get(&sample.tensor)
        .ok_or_else(|| format!("{} is missing {}", path.display(), sample.tensor))?;
    let mut values = Vec::with_capacity(sample.values as usize);
    match tensor.dtype.as_str() {
        "BF16" => {
            if sample.bits.is_some() || sample.group_size.is_some() {
                return Err(format!(
                    "{} BF16 sample declares quantization",
                    sample.class
                ));
            }
            let columns = *tensor
                .shape
                .last()
                .ok_or_else(|| format!("{} has no dimensions", sample.tensor))?;
            if sample.values > columns {
                return Err(format!(
                    "{} sample is wider than its tensor row",
                    sample.class
                ));
            }
            let start = tensor_byte_offset(data_start, tensor, sample.row, columns * 2)?;
            let bytes = read_bytes(&path, start, sample.values * 2)?;
            for bytes in bytes.chunks_exact(2) {
                values.push(f32::from_bits(
                    u32::from(u16::from_le_bytes([bytes[0], bytes[1]])) << 16,
                ));
            }
        }
        "U32" => {
            let bits = sample
                .bits
                .ok_or_else(|| format!("{} sample has no bit width", sample.class))?;
            let group_size = sample
                .group_size
                .ok_or_else(|| format!("{} sample has no group size", sample.class))?;
            if !matches!(bits, 2 | 4 | 8) || group_size == 0 {
                return Err(format!("{} sample has invalid quantization", sample.class));
            }
            let packed_columns = *tensor
                .shape
                .last()
                .ok_or_else(|| format!("{} has no dimensions", sample.tensor))?;
            let pack = u64::from(32 / bits);
            let columns = packed_columns
                .checked_mul(pack)
                .ok_or_else(|| format!("{} row width overflows", sample.tensor))?;
            if sample.values > columns || !columns.is_multiple_of(group_size) {
                return Err(format!("{} sample has invalid row width", sample.class));
            }
            let base = sample
                .tensor
                .strip_suffix(".weight")
                .ok_or_else(|| format!("{} is not a quantized weight", sample.tensor))?;
            let scales = tensors
                .get(&format!("{base}.scales"))
                .ok_or_else(|| format!("{} has no scales", sample.tensor))?;
            let biases = tensors
                .get(&format!("{base}.biases"))
                .ok_or_else(|| format!("{} has no biases", sample.tensor))?;
            let groups = columns / group_size;
            if scales.dtype != "BF16"
                || biases.dtype != "BF16"
                || scales.shape.last() != Some(&groups)
                || biases.shape != scales.shape
            {
                return Err(format!(
                    "{} has incompatible affine parameters",
                    sample.tensor
                ));
            }
            let packed_start =
                tensor_byte_offset(data_start, tensor, sample.row, packed_columns * 4)?;
            let scale_start = tensor_byte_offset(data_start, scales, sample.row, groups * 2)?;
            let bias_start = tensor_byte_offset(data_start, biases, sample.row, groups * 2)?;
            let packed = read_bytes(&path, packed_start, packed_columns * 4)?;
            let scales = read_bf16(&path, scale_start, groups)?;
            let biases = read_bf16(&path, bias_start, groups)?;
            let mask = (1_u32 << bits) - 1;
            for index in 0..sample.values {
                let word_index = (index / pack) as usize * 4;
                let word = u32::from_le_bytes([
                    packed[word_index],
                    packed[word_index + 1],
                    packed[word_index + 2],
                    packed[word_index + 3],
                ]);
                let quantized = (word >> ((index % pack) * u64::from(bits))) & mask;
                let group = (index / group_size) as usize;
                values.push(quantized as f32 * scales[group] + biases[group]);
            }
        }
        dtype => return Err(format!("{} sample uses unsupported {dtype}", sample.class)),
    }
    let mut hash = Sha256::new();
    for value in values {
        hash.update(value.to_le_bytes());
    }
    Ok(hex_digest(hash.finalize()))
}

fn tensor_byte_offset(
    data_start: u64,
    tensor: &Tensor,
    row: u64,
    row_bytes: u64,
) -> Result<u64, String> {
    let offset = row
        .checked_mul(row_bytes)
        .and_then(|offset| tensor.offsets[0].checked_add(offset))
        .and_then(|offset| data_start.checked_add(offset))
        .ok_or_else(|| "sample offset overflows".to_owned())?;
    let end = offset
        .checked_add(row_bytes)
        .ok_or_else(|| "sample range overflows".to_owned())?;
    if end > data_start.saturating_add(tensor.offsets[1]) {
        return Err("sample row exceeds its tensor".into());
    }
    Ok(offset)
}

fn read_bytes(path: &Path, offset: u64, length: u64) -> Result<Vec<u8>, String> {
    let mut file = File::open(path).map_err(|error| format!("{}: {error}", path.display()))?;
    file.seek(SeekFrom::Start(offset))
        .map_err(|error| format!("{}: {error}", path.display()))?;
    let mut bytes = vec![0_u8; length as usize];
    file.read_exact(&mut bytes)
        .map_err(|error| format!("{}: {error}", path.display()))?;
    Ok(bytes)
}

fn read_bf16(path: &Path, offset: u64, count: u64) -> Result<Vec<f32>, String> {
    Ok(read_bytes(path, offset, count * 2)?
        .chunks_exact(2)
        .map(|bytes| f32::from_bits(u32::from(u16::from_le_bytes([bytes[0], bytes[1]])) << 16))
        .collect())
}

fn validate_file_contract(manifest: &Manifest) -> Result<(), String> {
    validate_name(&manifest.tensor_inventory)?;
    if manifest.tensor_inventory_sha256.len() != 64
        || !manifest
            .tensor_inventory_sha256
            .bytes()
            .all(|byte| byte.is_ascii_hexdigit())
    {
        return Err("tensor inventory has an invalid SHA-256".into());
    }
    let mut paths = HashSet::new();
    let mut roles = BTreeMap::<&str, usize>::new();
    for artifact in &manifest.files {
        validate_name(&artifact.path)?;
        if !paths.insert(artifact.path.as_str()) {
            return Err(format!("duplicate artifact {}", artifact.path));
        }
        if artifact.size == 0
            || artifact.sha256.len() != 64
            || !artifact.sha256.bytes().all(|byte| byte.is_ascii_hexdigit())
        {
            return Err(format!("invalid SHA-256 for {}", artifact.path));
        }
        *roles.entry(&artifact.role).or_default() += 1;
    }
    for excluded in &manifest.excluded {
        validate_name(&excluded.path)?;
        if !paths.insert(excluded.path.as_str()) {
            return Err(format!("{} is both included and excluded", excluded.path));
        }
        if excluded.reason.is_empty() {
            return Err(format!("{} has no exclusion reason", excluded.path));
        }
        if excluded.size == 0
            || excluded.sha256.len() != 64
            || !excluded.sha256.bytes().all(|byte| byte.is_ascii_hexdigit())
        {
            return Err(format!("{} has invalid exclusion metadata", excluded.path));
        }
    }
    for (role, count) in [("core", 17), ("ple", 1), ("mtp", 1)] {
        if roles.get(role).copied() != Some(count) {
            return Err(format!("manifest must contain {count} {role} artifact(s)"));
        }
    }
    for required in [
        "model.safetensors.index.json",
        "config.json",
        "generation_config.json",
        "mtplx_runtime.json",
        "chat_template.jinja",
        "tokenizer.json",
        "tokenizer_config.json",
        "vocab.json",
        "merges.txt",
        "LICENSE",
    ] {
        if !manifest
            .files
            .iter()
            .any(|artifact| artifact.path == required)
        {
            return Err(format!("manifest is missing {required}"));
        }
    }
    if manifest.files.len() != 29 {
        return Err("manifest must contain exactly 29 selected files".into());
    }
    let excluded = manifest
        .excluded
        .iter()
        .map(|entry| entry.path.as_str())
        .collect::<BTreeSet<_>>();
    if excluded
        != [
            "model-vision.safetensors",
            "preprocessor_config.json",
            "processor_config.json",
            "video_preprocessor_config.json",
        ]
        .into_iter()
        .collect()
    {
        return Err("manifest must explicitly exclude the four vision artifacts".into());
    }
    Ok(())
}

fn validate_name(name: &str) -> Result<(), String> {
    let path = Path::new(name);
    if path.components().count() != 1 || path.file_name().is_none() {
        return Err(format!("artifact path must be one file name: {name}"));
    }
    Ok(())
}

fn validate_json_contract(
    source: &Path,
    file_name: &str,
    expected_values: &BTreeMap<String, Value>,
) -> Result<(), String> {
    let path = source.join(file_name);
    let config: Value = serde_json::from_slice(
        &fs::read(&path).map_err(|error| format!("{}: {error}", path.display()))?,
    )
    .map_err(|error| format!("{}: {error}", path.display()))?;
    for (pointer, expected) in expected_values {
        let actual = config
            .pointer(pointer)
            .ok_or_else(|| format!("{file_name} is missing {pointer}"))?;
        if actual != expected {
            return Err(format!(
                "{file_name} {pointer} is {actual}, expected {expected}"
            ));
        }
    }
    Ok(())
}

fn validate_index(manifest: &Manifest, source: &Path) -> Result<(), String> {
    let path = source.join("model.safetensors.index.json");
    let value: Value = serde_json::from_slice(
        &fs::read(&path).map_err(|error| format!("{}: {error}", path.display()))?,
    )
    .map_err(|error| format!("{}: {error}", path.display()))?;
    let map = value
        .get("weight_map")
        .and_then(Value::as_object)
        .ok_or_else(|| "model index has no weight_map object".to_owned())?;
    let core_files = manifest
        .files
        .iter()
        .filter(|artifact| artifact.role == "core")
        .map(|artifact| artifact.path.as_str())
        .collect::<HashSet<_>>();
    let indexed = map
        .iter()
        .filter_map(|(name, file)| {
            file.as_str()
                .filter(|file| core_files.contains(file))
                .map(|_| name.clone())
        })
        .collect::<BTreeSet<_>>();
    let actual = tensor_headers(manifest, source, Some("core"), true)?
        .keys()
        .cloned()
        .collect::<BTreeSet<_>>();
    if indexed != actual {
        return Err("model index and core shard tensor inventories differ".into());
    }
    let vision_count = map
        .values()
        .filter(|file| file.as_str() == Some("model-vision.safetensors"))
        .count();
    if vision_count == 0
        || !manifest
            .excluded
            .iter()
            .any(|entry| entry.path == "model-vision.safetensors")
    {
        return Err("vision tensors are not explicitly excluded".into());
    }
    Ok(())
}

fn build_inventory(manifest: &Manifest, source: &Path, complete: bool) -> Result<String, String> {
    tensor_headers(manifest, source, None, complete)?;
    let precision = precision_assignments(manifest, source, complete)?;
    let mut output = String::from(
        "file\tname\tdtype\tshape\tquant_bits\tgroup_size\tquant_mode\tdata_start\tdata_end\n",
    );
    for artifact in manifest
        .files
        .iter()
        .filter(|artifact| artifact.path.ends_with(".safetensors"))
    {
        let path = source.join(&artifact.path);
        let (data_start, tensors) = read_safetensors_header(&path, complete)?;
        for (name, tensor) in tensors {
            let start = data_start
                .checked_add(tensor.offsets[0])
                .ok_or_else(|| format!("{name} offset overflows"))?;
            let end = data_start
                .checked_add(tensor.offsets[1])
                .ok_or_else(|| format!("{name} offset overflows"))?;
            let quantization = precision.get(&name);
            output.push_str(&format!(
                "{}\t{name}\t{}\t{}\t{}\t{}\t{}\t{start}\t{end}\n",
                artifact.path,
                tensor.dtype,
                tensor
                    .shape
                    .iter()
                    .map(u64::to_string)
                    .collect::<Vec<_>>()
                    .join("x"),
                quantization.map_or_else(String::new, |value| value.bits.to_string()),
                quantization.map_or_else(String::new, |value| value.group_size.to_string()),
                quantization.map_or("", |value| value.mode.as_str()),
            ));
        }
    }
    Ok(output)
}

#[derive(Clone)]
struct Quantization {
    bits: u32,
    group_size: u64,
    mode: String,
}

fn precision_assignments(
    manifest: &Manifest,
    source: &Path,
    complete: bool,
) -> Result<BTreeMap<String, Quantization>, String> {
    let core = tensor_headers(manifest, source, Some("core"), complete)?;
    let config_path = source.join("config.json");
    let config: Value = serde_json::from_slice(
        &fs::read(&config_path).map_err(|error| format!("{}: {error}", config_path.display()))?,
    )
    .map_err(|error| format!("{}: {error}", config_path.display()))?;
    if config.get("quantization") != config.get("quantization_config") {
        return Err("config quantization maps differ".into());
    }
    let quantization = config
        .get("quantization")
        .and_then(Value::as_object)
        .ok_or_else(|| "config has no quantization object".to_owned())?;
    let mut assignments = BTreeMap::new();
    for (base, contract) in quantization {
        if matches!(base.as_str(), "bits" | "group_size" | "mode") {
            continue;
        }
        match contract {
            Value::Bool(false) => validate_unquantized(base, &core)?,
            Value::Object(contract) => {
                let bits = contract_u64(contract, "bits", base)? as u32;
                let group_size = contract_u64(contract, "group_size", base)?;
                let mode = contract
                    .get("mode")
                    .and_then(Value::as_str)
                    .ok_or_else(|| format!("{base} has no quantization mode"))?;
                add_quantized(
                    base,
                    Quantization {
                        bits,
                        group_size,
                        mode: mode.to_owned(),
                    },
                    &core,
                    &mut assignments,
                )?;
            }
            _ => return Err(format!("{base} has an invalid quantization contract")),
        }
    }
    reject_unassigned_quantized(&core, &assignments)?;

    let mtp = tensor_headers(manifest, source, Some("mtp"), complete)?;
    add_inferred_sidecar(&mtp, &mut assignments)?;
    let ple = tensor_headers(manifest, source, Some("ple"), complete)?;
    let ple_group = config
        .pointer("/mtplx_recipe/ngram/group_size")
        .and_then(Value::as_u64)
        .ok_or_else(|| "config has no PLE group size".to_owned())?;
    let ple_bits = config
        .pointer("/mtplx_recipe/ngram/bits")
        .and_then(Value::as_u64)
        .ok_or_else(|| "config has no PLE bit width".to_owned())? as u32;
    for tensor in ple.keys().filter(|name| name.ends_with(".weight")) {
        let base = tensor.trim_end_matches(".weight");
        add_quantized(
            base,
            Quantization {
                bits: ple_bits,
                group_size: ple_group,
                mode: "affine".into(),
            },
            &ple,
            &mut assignments,
        )?;
    }
    reject_unassigned_quantized(&ple, &assignments)?;
    Ok(assignments)
}

fn contract_u64(
    contract: &serde_json::Map<String, Value>,
    field: &str,
    base: &str,
) -> Result<u64, String> {
    contract
        .get(field)
        .and_then(Value::as_u64)
        .ok_or_else(|| format!("{base} has no quantization {field}"))
}

fn validate_unquantized(base: &str, tensors: &BTreeMap<String, Tensor>) -> Result<(), String> {
    if let Some(weight) = tensors.get(&format!("{base}.weight"))
        && weight.dtype != "BF16"
    {
        return Err(format!("{base} is unexpectedly quantized"));
    }
    if tensors.contains_key(&format!("{base}.scales"))
        || tensors.contains_key(&format!("{base}.biases"))
    {
        return Err(format!("{base} has unexpected affine parameters"));
    }
    Ok(())
}

fn add_inferred_sidecar(
    tensors: &BTreeMap<String, Tensor>,
    assignments: &mut BTreeMap<String, Quantization>,
) -> Result<(), String> {
    for (name, tensor) in tensors {
        if tensor.dtype != "U32" {
            continue;
        }
        let base = name
            .strip_suffix(".weight")
            .ok_or_else(|| format!("quantized sidecar tensor {name} is not a weight"))?;
        let scales = tensors
            .get(&format!("{base}.scales"))
            .ok_or_else(|| format!("{base} has no scales"))?;
        let packed = *tensor
            .shape
            .last()
            .ok_or_else(|| format!("{name} has no dimensions"))?;
        let groups = *scales
            .shape
            .last()
            .ok_or_else(|| format!("{base}.scales has no dimensions"))?;
        let group_size = if base.contains(".switch_mlp.") {
            32
        } else {
            64
        };
        let pack = groups
            .checked_mul(group_size)
            .and_then(|columns| columns.checked_div(packed))
            .filter(|pack| matches!(pack, 4 | 8 | 16))
            .ok_or_else(|| format!("{base} has an invalid packed width"))?;
        add_quantized(
            base,
            Quantization {
                bits: (32 / pack) as u32,
                group_size,
                mode: "affine".into(),
            },
            tensors,
            assignments,
        )?;
    }
    reject_unassigned_quantized(tensors, assignments)
}

fn add_quantized(
    base: &str,
    quantization: Quantization,
    tensors: &BTreeMap<String, Tensor>,
    assignments: &mut BTreeMap<String, Quantization>,
) -> Result<(), String> {
    if !matches!(quantization.bits, 2 | 4 | 8)
        || quantization.group_size == 0
        || quantization.mode != "affine"
    {
        return Err(format!("{base} has unsupported quantization"));
    }
    let names = [
        format!("{base}.weight"),
        format!("{base}.scales"),
        format!("{base}.biases"),
    ];
    let weight = tensors
        .get(&names[0])
        .ok_or_else(|| format!("{base} is missing its weight"))?;
    let scales = tensors
        .get(&names[1])
        .ok_or_else(|| format!("{base} is missing its scales"))?;
    let biases = tensors
        .get(&names[2])
        .ok_or_else(|| format!("{base} is missing its biases"))?;
    let packed = *weight
        .shape
        .last()
        .ok_or_else(|| format!("{base} has no dimensions"))?;
    let groups = *scales
        .shape
        .last()
        .ok_or_else(|| format!("{base}.scales has no dimensions"))?;
    let pack = u64::from(32 / quantization.bits);
    if weight.dtype != "U32"
        || scales.dtype != "BF16"
        || biases.dtype != "BF16"
        || scales.shape != biases.shape
        || weight.shape[..weight.shape.len() - 1] != scales.shape[..scales.shape.len() - 1]
        || packed.checked_mul(pack) != groups.checked_mul(quantization.group_size)
    {
        return Err(format!("{base} does not match its affine contract"));
    }
    for name in names {
        if assignments
            .insert(name.clone(), quantization.clone())
            .is_some()
        {
            return Err(format!("duplicate quantization assignment for {name}"));
        }
    }
    Ok(())
}

fn reject_unassigned_quantized(
    tensors: &BTreeMap<String, Tensor>,
    assignments: &BTreeMap<String, Quantization>,
) -> Result<(), String> {
    for (name, tensor) in tensors {
        let affine_parameter = name.ends_with(".scales") || name.ends_with(".biases");
        if (tensor.dtype == "U32" || affine_parameter) && !assignments.contains_key(name) {
            return Err(format!("{name} is unexpectedly quantized"));
        }
    }
    Ok(())
}

fn tensor_headers(
    manifest: &Manifest,
    source: &Path,
    role: Option<&str>,
    complete: bool,
) -> Result<BTreeMap<String, Tensor>, String> {
    let mut all = BTreeMap::new();
    for artifact in manifest.files.iter().filter(|artifact| {
        artifact.path.ends_with(".safetensors") && role.is_none_or(|role| artifact.role == role)
    }) {
        let (_, tensors) = read_safetensors_header(&source.join(&artifact.path), complete)?;
        for (name, tensor) in tensors {
            if all.insert(name.clone(), tensor).is_some() {
                return Err(format!("duplicate tensor {name}"));
            }
        }
    }
    Ok(all)
}

#[derive(Deserialize)]
struct Tensor {
    dtype: String,
    shape: Vec<u64>,
    #[serde(rename = "data_offsets")]
    offsets: [u64; 2],
}

fn read_safetensors_header(
    path: &Path,
    complete: bool,
) -> Result<(u64, BTreeMap<String, Tensor>), String> {
    let mut file = File::open(path).map_err(|error| format!("{}: {error}", path.display()))?;
    let size = file
        .metadata()
        .map_err(|error| format!("{}: {error}", path.display()))?
        .len();
    let mut length = [0_u8; 8];
    file.read_exact(&mut length)
        .map_err(|error| format!("{}: {error}", path.display()))?;
    let length = u64::from_le_bytes(length);
    if length == 0 || length > MAX_SAFETENSORS_HEADER {
        return Err(format!(
            "{} has invalid header size {length}",
            path.display()
        ));
    }
    let data_start = 8_u64
        .checked_add(length)
        .ok_or_else(|| format!("{} header size overflows", path.display()))?;
    let mut bytes = vec![0_u8; length as usize];
    file.read_exact(&mut bytes)
        .map_err(|error| format!("{}: {error}", path.display()))?;
    let mut value = serde_json::from_slice::<UniqueObject>(&bytes)
        .map_err(|error| format!("{}: {error}", path.display()))?;
    value.0.remove("__metadata__");
    let mut tensors = BTreeMap::new();
    for (name, value) in value.0 {
        let tensor: Tensor = serde_json::from_value(value)
            .map_err(|error| format!("{} tensor {name}: {error}", path.display()))?;
        if tensor.shape.is_empty() || tensor.offsets[0] > tensor.offsets[1] {
            return Err(format!(
                "{} tensor {name} has invalid layout",
                path.display()
            ));
        }
        let element_size = match tensor.dtype.as_str() {
            "BF16" => 2_u64,
            "I64" => 8_u64,
            "U32" => 4_u64,
            _ => {
                return Err(format!(
                    "{} tensor {name} has unsupported dtype",
                    path.display()
                ));
            }
        };
        let expected_bytes = tensor
            .shape
            .iter()
            .try_fold(element_size, |bytes, dimension| {
                bytes.checked_mul(*dimension)
            });
        if expected_bytes != tensor.offsets[1].checked_sub(tensor.offsets[0]) {
            return Err(format!(
                "{} tensor {name} has a shape/byte-length mismatch",
                path.display()
            ));
        }
        if complete && data_start.saturating_add(tensor.offsets[1]) > size {
            return Err(format!("{} tensor {name} exceeds its file", path.display()));
        }
        tensors.insert(name, tensor);
    }
    Ok((data_start, tensors))
}

struct UniqueObject(BTreeMap<String, Value>);

impl<'de> Deserialize<'de> for UniqueObject {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        struct UniqueObjectVisitor;

        impl<'de> Visitor<'de> for UniqueObjectVisitor {
            type Value = UniqueObject;

            fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
                formatter.write_str("a safetensors header object with unique tensor names")
            }

            fn visit_map<M>(self, mut map: M) -> Result<Self::Value, M::Error>
            where
                M: MapAccess<'de>,
            {
                let mut values = BTreeMap::new();
                while let Some((name, value)) = map.next_entry::<String, Value>()? {
                    if values.insert(name.clone(), value).is_some() {
                        return Err(serde::de::Error::custom(format!(
                            "duplicate tensor name {name}"
                        )));
                    }
                }
                Ok(UniqueObject(values))
            }
        }

        deserializer.deserialize_map(UniqueObjectVisitor)
    }
}

fn sha256_file(path: &Path) -> Result<String, String> {
    let mut file = File::open(path).map_err(|error| format!("{}: {error}", path.display()))?;
    let mut hash = Sha256::new();
    let mut buffer = vec![0_u8; 8 * 1024 * 1024];
    loop {
        let read = file
            .read(&mut buffer)
            .map_err(|error| format!("{}: {error}", path.display()))?;
        if read == 0 {
            break;
        }
        hash.update(&buffer[..read]);
    }
    Ok(hex_digest(hash.finalize()))
}

fn hex_digest(digest: impl AsRef<[u8]>) -> String {
    let mut output = String::with_capacity(64);
    for byte in digest.as_ref() {
        output.push_str(&format!("{byte:02x}"));
    }
    output
}

fn fetch(manifest_path: &Path, output: &Path) -> Result<(), String> {
    let (manifest, _) = load_manifest(manifest_path)?;
    validate_file_contract(&manifest)?;
    fs::create_dir_all(output).map_err(|error| format!("{}: {error}", output.display()))?;
    for artifact in &manifest.files {
        let destination = output.join(&artifact.path);
        if destination.exists() {
            verify_artifact(&destination, artifact)?;
            continue;
        }
        let partial = output.join(format!("{}.part", artifact.path));
        if partial.metadata().map_or(0, |metadata| metadata.len()) > artifact.size {
            File::create(&partial).map_err(|error| format!("{}: {error}", partial.display()))?;
        }
        if partial.metadata().map_or(0, |metadata| metadata.len()) == artifact.size {
            verify_artifact(&partial, artifact)?;
            fs::rename(&partial, &destination)
                .map_err(|error| format!("{}: {error}", destination.display()))?;
            println!("authenticated {}", artifact.path);
            continue;
        }
        let url = format!(
            "https://huggingface.co/{}/resolve/{}/{}",
            manifest.source.repository, manifest.source.revision, artifact.path
        );
        download(&url, &partial, true)?;
        verify_artifact(&partial, artifact)?;
        fs::rename(&partial, &destination)
            .map_err(|error| format!("{}: {error}", destination.display()))?;
        println!("fetched {}", artifact.path);
    }
    verify(manifest_path, output)?;
    Ok(())
}

fn verify_artifact(path: &Path, artifact: &Artifact) -> Result<(), String> {
    let size = path
        .metadata()
        .map_err(|error| format!("{}: {error}", path.display()))?
        .len();
    if size != artifact.size {
        return Err(format!(
            "{} has {size} bytes, expected {}",
            path.display(),
            artifact.size
        ));
    }
    let actual = sha256_file(path)?;
    if actual != artifact.sha256 {
        return Err(format!(
            "{} has SHA-256 {actual}, expected {}",
            path.display(),
            artifact.sha256
        ));
    }
    Ok(())
}

fn download(url: &str, partial: &Path, https_only: bool) -> Result<(), String> {
    let offset = partial.metadata().map_or(0, |metadata| metadata.len());
    let agent: ureq::Agent = ureq::Agent::config_builder()
        .https_only(https_only)
        .timeout_connect(Some(DOWNLOAD_TIMEOUT))
        .timeout_recv_response(Some(DOWNLOAD_TIMEOUT))
        .timeout_recv_body(Some(DOWNLOAD_TIMEOUT))
        .build()
        .into();
    let mut request = agent.get(url);
    if offset > 0 {
        request = request.header("Range", format!("bytes={offset}-"));
    }
    let mut response = request
        .call()
        .map_err(|error| format!("artifact download failed: {error}"))?;
    let status = response.status().as_u16();
    let append = offset > 0 && status == 206;
    if offset > 0 && !matches!(status, 200 | 206) {
        return Err(format!(
            "artifact server returned HTTP {status} while resuming at byte {offset}"
        ));
    }
    if append {
        let expected = format!("bytes {offset}-");
        let content_range = response
            .headers()
            .get("content-range")
            .and_then(|value| value.to_str().ok())
            .unwrap_or_default();
        if !content_range.starts_with(&expected) {
            return Err(format!(
                "artifact server returned an invalid Content-Range while resuming at byte {offset}"
            ));
        }
    }
    let mut file = OpenOptions::new()
        .create(true)
        .write(true)
        .append(append)
        .truncate(!append)
        .open(partial)
        .map_err(|error| format!("{}: {error}", partial.display()))?;
    let mut body = response.body_mut().as_reader();
    let mut buffer = vec![0_u8; 1024 * 1024];
    loop {
        let read = body
            .read(&mut buffer)
            .map_err(|error| format!("artifact download failed: {error}"))?;
        if read == 0 {
            break;
        }
        file.write_all(&buffer[..read])
            .map_err(|error| format!("{}: {error}", partial.display()))?;
    }
    Ok(())
}

fn materialize(manifest_path: &Path, source: &Path, output: &Path) -> Result<(), String> {
    verify(manifest_path, source)?;
    fs::create_dir(output).map_err(|error| format!("{}: {error}", output.display()))?;
    let (manifest, manifest_directory) = load_manifest(manifest_path)?;
    for artifact in &manifest.files {
        let from = source.join(&artifact.path);
        let to = output.join(&artifact.path);
        if fs::hard_link(&from, &to).is_err() {
            fs::copy(&from, &to).map_err(|error| format!("{}: {error}", to.display()))?;
        }
    }
    fs::copy(manifest_path, output.join("manifest.json"))
        .map_err(|error| format!("{}: {error}", output.display()))?;
    fs::copy(
        manifest_directory.join(&manifest.tensor_inventory),
        output.join(&manifest.tensor_inventory),
    )
    .map_err(|error| format!("{}: {error}", output.display()))?;
    verify(manifest_path, output)?;
    println!(
        "materialized {} files ({} bytes) in {}",
        manifest.files.len(),
        manifest
            .files
            .iter()
            .map(|artifact| artifact.size)
            .sum::<u64>(),
        output.display()
    );
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::TcpListener;
    use std::thread;

    fn safetensors_fixture(name: &str, header: &str, payload: &[u8]) -> PathBuf {
        let path = env::temp_dir().join(format!(
            "ds4-qwen38-{name}-{}-{}.safetensors",
            std::process::id(),
            std::thread::current().name().unwrap_or("test")
        ));
        let mut bytes = (header.len() as u64).to_le_bytes().to_vec();
        bytes.extend_from_slice(header.as_bytes());
        bytes.extend_from_slice(payload);
        fs::write(&path, bytes).unwrap();
        path
    }

    #[test]
    fn artifact_names_cannot_escape_the_output_directory() {
        assert!(validate_name("model.safetensors").is_ok());
        assert!(validate_name("../model.safetensors").is_err());
        assert!(validate_name("nested/model.safetensors").is_err());
        assert!(validate_name("").is_err());
    }

    #[test]
    fn safetensors_validation_rejects_duplicate_names_and_invalid_ranges() {
        let duplicate = safetensors_fixture(
            "duplicate",
            r#"{"x":{"dtype":"BF16","shape":[1],"data_offsets":[0,2]},"x":{"dtype":"BF16","shape":[1],"data_offsets":[0,2]}}"#,
            &[0, 0],
        );
        assert!(read_safetensors_header(&duplicate, true).is_err());
        fs::remove_file(duplicate).unwrap();

        let outside = safetensors_fixture(
            "outside",
            r#"{"x":{"dtype":"BF16","shape":[1],"data_offsets":[0,2]}}"#,
            &[0],
        );
        assert!(read_safetensors_header(&outside, true).is_err());
        fs::remove_file(outside).unwrap();
    }

    #[test]
    fn affine_sample_digest_uses_packed_values_scales_and_biases() {
        let header = r#"{"x.biases":{"dtype":"BF16","shape":[1,1],"data_offsets":[6,8]},"x.scales":{"dtype":"BF16","shape":[1,1],"data_offsets":[4,6]},"x.weight":{"dtype":"U32","shape":[1,1],"data_offsets":[0,4]}}"#;
        let mut payload = 0x7654_3210_u32.to_le_bytes().to_vec();
        payload.extend_from_slice(&0x4000_u16.to_le_bytes());
        payload.extend_from_slice(&0x3f80_u16.to_le_bytes());
        let path = safetensors_fixture("affine", header, &payload);
        let sample = Sample {
            class: "fixture".into(),
            file: path.file_name().unwrap().to_string_lossy().into_owned(),
            tensor: "x.weight".into(),
            row: 0,
            values: 8,
            bits: Some(4),
            group_size: Some(8),
            sha256: String::new(),
        };
        let mut expected = Sha256::new();
        for value in [1.0_f32, 3.0, 5.0, 7.0, 9.0, 11.0, 13.0, 15.0] {
            expected.update(value.to_le_bytes());
        }
        assert_eq!(
            sample_digest(path.parent().unwrap(), &sample).unwrap(),
            hex_digest(expected.finalize())
        );
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn artifact_download_resumes_an_authenticated_partial() {
        let content = b"authenticated artifact";
        let offset = 7;
        let directory = env::temp_dir().join(format!("ds4-qwen38-fetch-{}", std::process::id()));
        fs::create_dir_all(&directory).unwrap();
        let partial = directory.join("artifact.part");
        fs::write(&partial, &content[..offset]).unwrap();
        let listener = match TcpListener::bind("127.0.0.1:0") {
            Ok(listener) => listener,
            Err(error) if error.kind() == std::io::ErrorKind::PermissionDenied => {
                fs::remove_dir_all(directory).unwrap();
                return;
            }
            Err(error) => panic!("could not start test server: {error}"),
        };
        let address = listener.local_addr().unwrap();
        let server = thread::spawn(move || {
            let (mut connection, _) = listener.accept().unwrap();
            let mut request = [0_u8; 2048];
            let read = connection.read(&mut request).unwrap();
            let request = String::from_utf8_lossy(&request[..read]).to_ascii_lowercase();
            assert!(request.contains("range: bytes=7-"));
            let remaining = &content[offset..];
            write!(
                connection,
                "HTTP/1.1 206 Partial Content\r\nContent-Length: {}\r\nContent-Range: bytes {offset}-{}/{}\r\nConnection: close\r\n\r\n",
                remaining.len(),
                content.len() - 1,
                content.len(),
            )
            .unwrap();
            connection.write_all(remaining).unwrap();
        });

        download(&format!("http://{address}/artifact"), &partial, false).unwrap();
        server.join().unwrap();
        verify_artifact(
            &partial,
            &Artifact {
                path: "artifact".into(),
                role: "fixture".into(),
                size: content.len() as u64,
                sha256: "ef056ce05d5f86154a8498b5504a46bc4b633fb68ec795a3a1cfd0548f791e4c".into(),
            },
        )
        .unwrap();
        fs::remove_dir_all(directory).unwrap();
    }
}