franken_ocr 0.10.0

//! Exact production tensor censuses for the pinned Unlimited-OCR checkpoints.
//!
//! The generic [`super::weights::Weights::census`] API checks names only.  The
//! production contract is stronger: the verified v1 source shard has exactly
//! 2,710 tensors, and the conservative `.focrq` recipe fixes each tensor's
//! logical shape and storage dtype.  This module keeps that complete manifest
//! embedded in the binary and applies the same bounded, named comparison to
//! both the header probe and the post-load directory. This is an exact schema
//! census; payload identity is established separately by artifact SHA-256. The
//! manifest distinguishes the safetensors index's tensor-payload `total_size`
//! from the full shard size (8-byte length prefix + JSON header + payload).
//!
//! Unlimited-OCR-2 (plan D22) is v1's census plus an embedded delta
//! (`unlimited_ocr_2_manifest.json`, generated by `scripts/gen_uocr2_census.py`
//! from the pinned checkpoint): the 108 tensors it adds, its one I64 tensor
//! (verified and skipped), the SHA-256 of the resulting 2,819-entry list, which
//! this module re-derives, and the SHA-256 of every vision tensor, which equals
//! v1's.

use std::collections::BTreeMap;
use std::sync::OnceLock;

use serde::Deserialize;

use super::model_arch::{self, ModelArch, TensorCensus};
use super::shards::{Int64Content, Int64Rule};
use super::weights::{DType, TensorRecord, Weights};
use crate::error::{FocrError, FocrResult};
use crate::quant::convert::{UNLIMITED_OCR_INT8_RECIPE_ID, UNLIMITED_OCR_WASM_INT4_RECIPE_ID};
use crate::quant::recipe::{Recipe, WasmInt4Policy, classify_wasm_experts_int4};

const MANIFEST_JSON: &str = include_str!("unlimited_ocr_manifest.json");
const MANIFEST_SHA256: &str = "24ff1cfffe71eec6f07bfae8e8eb12b342877bccc43ad6ae4a4a4ceffb76edd3";
const MANIFEST_SCHEMA: &str = "franken_ocr.unlimited_ocr_tensor_manifest.v1";
const HF_COMMIT: &str = "3a7f4dbbbffcc6f9282712c5b0d7cc31b3812da5";
const SOURCE_FILE: &str = "model-00001-of-000001.safetensors";
const SOURCE_SHA256: &str = "2bc48a7a110061ea58fff65d3169367eebe3aee371ca6968dc2219c1b2855fc6";
const SOURCE_INDEX_SHA256: &str =
    "354be1f2dcfb72ebb385e25465522ce5413a77c36f3b35fec088a3162a11af99";
const VERIFIED_ARTIFACT_SHA256: &str =
    "573340710167697891bf52dfa4cbb5d0a02a68f3011c01f8ef83fd34622fb592";
const SOURCE_INDEX_TOTAL_SIZE: u64 = 6_672_212_480;
const SOURCE_FILE_SIZE: u64 = 6_672_547_120;
const TENSOR_COUNT: usize = 2_710;
const DIFF_PREVIEW_LIMIT: usize = 6;

const MANIFEST_2_JSON: &str = include_str!("unlimited_ocr_2_manifest.json");
const MANIFEST_2_SHA256: &str = "cbfd207d3f66c803e484a49f2ed754916e567aaa7b7c65125759e70d54e77dd0";
const MANIFEST_2_SCHEMA: &str = "franken_ocr.unlimited_ocr_2_tensor_delta.v1";
const V2_HF_COMMIT: &str = "5632329a32e103d5f1add270681e0f5c1233c6e4";
/// The pinned Unlimited-OCR-2 shards in index order: name and SHA-256.
const V2_SOURCE_SHARDS: [(&str, &str); 2] = [
    (
        "model-00001-of-00002.safetensors",
        "4c40b740f209fd7682040d4ebf19409208acde5c401ccd2dde426d0f41542f18",
    ),
    (
        "model-00002-of-00002.safetensors",
        "9490b82fa47f89e5110eb5577f8ab1b6694c12121905dbc086ddb10207a4ba0c",
    ),
];
/// The byte size of each pinned v2 shard, in index order.
const V2_SOURCE_SHARD_SIZES: [u64; 2] = [4_998_430_312, 1_689_875_264];
/// The `source_sha256` an artifact converted from the pinned v2 shards
/// carries: SHA-256 over the shard digests in index order (D16).
pub(super) const V2_SOURCE_SHA256: &str =
    "da73bc2203a906313f8448dd92be7d0866853234c8e7087f79e202257be6a5f3";
const V2_SOURCE_INDEX_SHA256: &str =
    "dc60a31ecb77eb03e166345d18cbe25bc71d2a19cc6fc089d8a9ac1c7f209671";
const V2_SOURCE_INDEX_TOTAL_SIZE: u64 = 6_687_958_536;
/// Every v2 tensor, the I64 one included.
const V2_TENSOR_COUNT: usize = 2_819;
/// The non-I64 tensors v2 adds over v1: per layer, the SR-SWA `dash_router`
/// and `indexer` projections and norms.
const V2_ADDED_COUNT: usize = 108;
/// SHA-256 of the canonical 2,819-entry list, one `{name}\t{dtype}\t{dims}\n`
/// line per tensor in byte order of names.
const V2_CENSUS_SHA256: &str = "61d7135fd616b1851c6bdb7d91b60ec2553b28b46f4393f7cba4be78a883bc50";
/// The vision tower, projector, `image_newline` and `view_seperator` tensors.
const V2_VISION_TENSOR_COUNT: usize = 476;
/// v2's I64 tensor: the CLIP embedding's constant `position_ids` (D8).
static V2_INT64_RULES: [Int64Rule; 1] = [Int64Rule {
    name: "model.vision_model.embeddings.position_ids",
    shape: &[1, 257],
    content: Int64Content::Arange,
}];

/// The pinned source shards of `census`'s checkpoint, in index order: each
/// shard's file name and SHA-256. A single-entry list is also the shard name
/// the model resolver looks for in a directory without an index.
pub(crate) fn source_shards(census: TensorCensus) -> &'static [(&'static str, &'static str)] {
    match census {
        TensorCensus::UnlimitedOcrV1 => &[(SOURCE_FILE, SOURCE_SHA256)],
        TensorCensus::UnlimitedOcrV2 => &V2_SOURCE_SHARDS,
    }
}

/// The I64 tensors `census` accepts only to verify and skip (D8).
pub(crate) fn int64_rules(census: TensorCensus) -> &'static [Int64Rule] {
    match census {
        TensorCensus::UnlimitedOcrV1 => &[],
        TensorCensus::UnlimitedOcrV2 => &V2_INT64_RULES,
    }
}

/// Refuse a conversion source whose shard count differs from the pinned
/// source's, before any shard is read.
#[cfg_attr(not(feature = "native"), allow(dead_code))] // sole non-test caller is the CLI convert path
pub(crate) fn validate_conversion_shard_count(
    census: TensorCensus,
    count: usize,
) -> FocrResult<()> {
    let pinned = source_shards(census);
    if count == pinned.len() {
        return Ok(());
    }
    Err(FocrError::FormatMismatch(format!(
        "{} conversion requires its pinned {}-shard source checkpoint; this source has {count} \
         shard(s)",
        census_model_id(census),
        pinned.len()
    )))
}

/// Refuse conversion inputs whose output the production artifact loader would
/// later reject: shard `position`'s SHA-256 must be the pinned one. The
/// conservative format is bound to the exact pinned source shards, not merely
/// a same-shaped safetensors serialization.
#[cfg_attr(not(feature = "native"), allow(dead_code))] // sole non-test caller is the CLI convert path
pub(crate) fn validate_conversion_source_shard(
    census: TensorCensus,
    position: usize,
    actual: &[u8; 32],
) -> FocrResult<()> {
    let actual = actual
        .iter()
        .map(|byte| format!("{byte:02x}"))
        .collect::<String>();
    let pinned = source_shards(census);
    let Some(&(name, expected)) = pinned.get(position) else {
        return validate_conversion_shard_count(census, position + 1);
    };
    if actual == expected {
        return Ok(());
    }
    Err(FocrError::FormatMismatch(format!(
        "{} conversion requires pinned source shard {name} SHA-256 {expected}; got {actual}. \
         Supply the canonical checkpoint or use raw compatible weights directly",
        census_model_id(census)
    )))
}

/// The model id a census describes.
fn census_model_id(census: TensorCensus) -> &'static str {
    match census {
        TensorCensus::UnlimitedOcrV1 => model_arch::UNLIMITED_OCR_V1_ID,
        TensorCensus::UnlimitedOcrV2 => model_arch::UNLIMITED_OCR_V2_ID,
    }
}

#[derive(Debug, Deserialize)]
struct UnlimitedOcrManifest {
    schema: String,
    model_id: String,
    hf_commit: String,
    source_safetensors: String,
    source_sha256: String,
    source_index_total_size: u64,
    source_file_size: u64,
    source_index_sha256: String,
    tensor_count: usize,
    quant_recipe: String,
    verified_conservative_artifact_sha256: String,
    tensors: BTreeMap<String, ManifestTensor>,
}

#[derive(Debug, Clone, Deserialize)]
struct ManifestTensor {
    shape: Vec<usize>,
    source_dtype: DType,
    storage_dtype: DType,
}

static MANIFEST: OnceLock<Result<UnlimitedOcrManifest, String>> = OnceLock::new();

fn manifest() -> FocrResult<&'static UnlimitedOcrManifest> {
    MANIFEST
        .get_or_init(parse_and_verify_manifest)
        .as_ref()
        .map_err(|message| {
            FocrError::FormatMismatch(format!(
                "embedded Unlimited-OCR tensor manifest is invalid: {message}"
            ))
        })
}

fn parse_and_verify_manifest() -> Result<UnlimitedOcrManifest, String> {
    use sha2::Digest as _;

    let manifest_sha256 = format!("{:x}", sha2::Sha256::digest(MANIFEST_JSON.as_bytes()));
    if manifest_sha256 != MANIFEST_SHA256 {
        return Err(format!(
            "embedded JSON sha256 is {manifest_sha256}, expected {MANIFEST_SHA256}"
        ));
    }
    let manifest: UnlimitedOcrManifest =
        serde_json::from_str(MANIFEST_JSON).map_err(|error| error.to_string())?;

    let metadata_checks = [
        ("schema", manifest.schema.as_str(), MANIFEST_SCHEMA),
        (
            "model_id",
            manifest.model_id.as_str(),
            model_arch::UNLIMITED_OCR_V1_ID,
        ),
        ("hf_commit", manifest.hf_commit.as_str(), HF_COMMIT),
        (
            "source_safetensors",
            manifest.source_safetensors.as_str(),
            SOURCE_FILE,
        ),
        (
            "source_sha256",
            manifest.source_sha256.as_str(),
            SOURCE_SHA256,
        ),
        (
            "source_index_sha256",
            manifest.source_index_sha256.as_str(),
            SOURCE_INDEX_SHA256,
        ),
        (
            "quant_recipe",
            manifest.quant_recipe.as_str(),
            UNLIMITED_OCR_INT8_RECIPE_ID,
        ),
        (
            "verified_conservative_artifact_sha256",
            manifest.verified_conservative_artifact_sha256.as_str(),
            VERIFIED_ARTIFACT_SHA256,
        ),
    ];
    for (field, actual, expected) in metadata_checks {
        if actual != expected {
            return Err(format!(
                "{field} is {actual:?}, expected pinned value {expected:?}"
            ));
        }
    }
    if manifest.source_index_total_size != SOURCE_INDEX_TOTAL_SIZE {
        return Err(format!(
            "source_index_total_size is {}, expected {SOURCE_INDEX_TOTAL_SIZE}",
            manifest.source_index_total_size
        ));
    }
    if manifest.source_file_size != SOURCE_FILE_SIZE {
        return Err(format!(
            "source_file_size is {}, expected {SOURCE_FILE_SIZE}",
            manifest.source_file_size
        ));
    }
    if manifest.tensor_count != TENSOR_COUNT || manifest.tensors.len() != TENSOR_COUNT {
        return Err(format!(
            "tensor_count declares {} and directory contains {}, expected {TENSOR_COUNT}",
            manifest.tensor_count,
            manifest.tensors.len()
        ));
    }

    let recipe = Recipe::validated_default();
    for (name, tensor) in &manifest.tensors {
        if tensor.shape.is_empty() || tensor.shape.contains(&0) {
            return Err(format!(
                "tensor {name:?} has invalid shape {:?}",
                tensor.shape
            ));
        }
        if tensor.source_dtype != DType::BF16 {
            return Err(format!(
                "tensor {name:?} source dtype is {:?}, expected BF16",
                tensor.source_dtype
            ));
        }
        let expected_storage = if recipe.is_quantized(name) {
            DType::QInt8PerChan
        } else {
            DType::BF16
        };
        if tensor.storage_dtype != expected_storage {
            return Err(format!(
                "tensor {name:?} storage dtype is {:?}, recipe requires {expected_storage:?}",
                tensor.storage_dtype
            ));
        }
    }
    Ok(manifest)
}

/// The embedded Unlimited-OCR-2 delta (`unlimited_ocr_2_manifest.json`).
#[derive(Debug, Deserialize)]
struct UnlimitedOcr2Delta {
    schema: String,
    model_id: String,
    base_model_id: String,
    hf_commit: String,
    source_index_sha256: String,
    source_index_total_size: u64,
    source_shards: Vec<DeltaShard>,
    tokenizer_sha256: String,
    tensor_count: usize,
    census_sha256: String,
    quant_recipe: String,
    verified_conservative_artifact_sha256: Option<String>,
    added_tensors: BTreeMap<String, DeltaTensor>,
    int64_tensors: BTreeMap<String, DeltaInt64>,
    vision_identity_sha256: BTreeMap<String, String>,
}

#[derive(Debug, Deserialize)]
struct DeltaShard {
    name: String,
    size: u64,
    sha256: String,
}

#[derive(Debug, Deserialize)]
struct DeltaTensor {
    shape: Vec<usize>,
    source_dtype: DType,
}

#[derive(Debug, Deserialize)]
struct DeltaInt64 {
    shape: Vec<usize>,
    content: String,
}

/// The resolved v2 census: v1's tensors plus the delta's, without the I64 one
/// (raw loads verify and drop it before the census sees the directory).
struct UnlimitedOcr2Census {
    tensors: BTreeMap<String, ManifestTensor>,
    /// The full-tensor SHA-256 of each vision tensor, equal to v1's bytes
    /// (checked against real bytes by the source-gated tests).
    #[cfg_attr(not(test), allow(dead_code))]
    vision_identity_sha256: BTreeMap<String, String>,
}

static MANIFEST_2: OnceLock<Result<UnlimitedOcr2Census, String>> = OnceLock::new();

fn v2_census() -> FocrResult<&'static UnlimitedOcr2Census> {
    MANIFEST_2
        .get_or_init(parse_and_verify_v2_delta)
        .as_ref()
        .map_err(|message| {
            FocrError::FormatMismatch(format!(
                "embedded Unlimited-OCR-2 tensor delta is invalid: {message}"
            ))
        })
}

/// Whether `name` is a vision tower, projector, newline or separator tensor:
/// the set v2 reuses from v1 byte for byte.
fn is_vision_identity_tensor(name: &str) -> bool {
    [
        "model.sam_model.",
        "model.vision_model.",
        "model.projector.",
    ]
    .iter()
    .any(|prefix| name.starts_with(prefix))
        || name == "model.image_newline"
        || name == "model.view_seperator"
}

fn is_sha256_hex(text: &str) -> bool {
    text.len() == 64 && text.bytes().all(|b| matches!(b, b'0'..=b'9' | b'a'..=b'f'))
}

/// The SHA-256 of a census in canonical form: one `{name}\t{dtype}\t{dims}\n`
/// line per tensor, in byte order of names.
fn canonical_census_sha256<'a>(
    entries: impl IntoIterator<Item = (&'a str, &'a str, &'a [usize])>,
) -> String {
    use sha2::Digest as _;
    use std::fmt::Write as _;

    let mut lines: Vec<String> = entries
        .into_iter()
        .map(|(name, dtype, shape)| {
            let mut line = format!("{name}\t{dtype}\t");
            for (axis, dim) in shape.iter().enumerate() {
                let _ = write!(line, "{}{dim}", if axis == 0 { "" } else { "," });
            }
            line.push('\n');
            line
        })
        .collect();
    lines.sort_unstable();
    format!("{:x}", sha2::Sha256::digest(lines.concat().as_bytes()))
}

fn parse_and_verify_v2_delta() -> Result<UnlimitedOcr2Census, String> {
    verify_v2_delta(MANIFEST_2_JSON, MANIFEST_2_SHA256)
}

/// Parse `json` as the v2 delta and check it against `pinned_sha256` and every
/// pinned fact, re-deriving the census from v1's.
fn verify_v2_delta(json: &str, pinned_sha256: &str) -> Result<UnlimitedOcr2Census, String> {
    use sha2::Digest as _;

    let manifest_sha256 = format!("{:x}", sha2::Sha256::digest(json.as_bytes()));
    if manifest_sha256 != pinned_sha256 {
        return Err(format!(
            "embedded JSON sha256 is {manifest_sha256}, expected {pinned_sha256}"
        ));
    }
    let delta: UnlimitedOcr2Delta =
        serde_json::from_str(json).map_err(|error| error.to_string())?;
    let metadata_checks = [
        ("schema", delta.schema.as_str(), MANIFEST_2_SCHEMA),
        (
            "model_id",
            delta.model_id.as_str(),
            model_arch::UNLIMITED_OCR_V2_ID,
        ),
        (
            "base_model_id",
            delta.base_model_id.as_str(),
            model_arch::UNLIMITED_OCR_V1_ID,
        ),
        ("hf_commit", delta.hf_commit.as_str(), V2_HF_COMMIT),
        (
            "source_index_sha256",
            delta.source_index_sha256.as_str(),
            V2_SOURCE_INDEX_SHA256,
        ),
        (
            "tokenizer_sha256",
            delta.tokenizer_sha256.as_str(),
            model_arch::UNLIMITED_OCR_V2_TOKENIZER_SHA256,
        ),
        (
            "census_sha256",
            delta.census_sha256.as_str(),
            V2_CENSUS_SHA256,
        ),
        (
            "quant_recipe",
            delta.quant_recipe.as_str(),
            crate::quant::convert::UNLIMITED_OCR_2_INT8_RECIPE_ID,
        ),
    ];
    for (field, actual, expected) in metadata_checks {
        if actual != expected {
            return Err(format!(
                "{field} is {actual:?}, expected pinned value {expected:?}"
            ));
        }
    }
    if delta.source_index_total_size != V2_SOURCE_INDEX_TOTAL_SIZE {
        return Err(format!(
            "source_index_total_size is {}, expected {V2_SOURCE_INDEX_TOTAL_SIZE}",
            delta.source_index_total_size
        ));
    }
    let shards_match = delta.source_shards.len() == V2_SOURCE_SHARDS.len()
        && delta
            .source_shards
            .iter()
            .zip(V2_SOURCE_SHARDS.iter().zip(V2_SOURCE_SHARD_SIZES))
            .all(|(shard, (&(name, sha256), size))| {
                shard.name == name && shard.sha256 == sha256 && shard.size == size
            });
    if !shards_match {
        return Err(format!(
            "source_shards {:?} are not the pinned {V2_SOURCE_SHARDS:?}",
            delta.source_shards
        ));
    }
    if let Some(sha256) = &delta.verified_conservative_artifact_sha256
        && !is_sha256_hex(sha256)
    {
        return Err(format!(
            "verified_conservative_artifact_sha256 {sha256:?} is not a SHA-256"
        ));
    }

    let base = manifest().map_err(|error| format!("its v1 base census is unavailable: {error}"))?;
    let recipe = Recipe::validated_default();
    let mut tensors = base.tensors.clone();
    for (name, tensor) in &delta.added_tensors {
        if tensor.shape.is_empty() || tensor.shape.contains(&0) {
            return Err(format!(
                "added tensor {name:?} has invalid shape {:?}",
                tensor.shape
            ));
        }
        if tensor.source_dtype != DType::BF16 {
            return Err(format!(
                "added tensor {name:?} source dtype is {:?}, expected BF16",
                tensor.source_dtype
            ));
        }
        // v2's conservative recipe is the same validated classification as
        // v1's; the SR-SWA router and indexer are kept BF16 by name (U2-C3).
        let storage_dtype = if recipe.is_quantized(name) {
            DType::QInt8PerChan
        } else {
            DType::BF16
        };
        let entry = ManifestTensor {
            shape: tensor.shape.clone(),
            source_dtype: tensor.source_dtype,
            storage_dtype,
        };
        if tensors.insert(name.clone(), entry).is_some() {
            return Err(format!("added tensor {name:?} is already in the v1 census"));
        }
    }
    if delta.added_tensors.len() != V2_ADDED_COUNT {
        return Err(format!(
            "added_tensors lists {}, expected {V2_ADDED_COUNT}",
            delta.added_tensors.len()
        ));
    }
    let int64_match = delta.int64_tensors.len() == V2_INT64_RULES.len()
        && V2_INT64_RULES.iter().all(|rule| {
            delta.int64_tensors.get(rule.name).is_some_and(|tensor| {
                tensor.shape == rule.shape
                    && tensor.content == "arange"
                    && rule.content == Int64Content::Arange
                    && !tensors.contains_key(rule.name)
            })
        });
    if !int64_match {
        return Err(format!(
            "int64_tensors {:?} are not the census rules {V2_INT64_RULES:?}",
            delta.int64_tensors
        ));
    }
    let vision: Vec<&String> = base
        .tensors
        .keys()
        .filter(|name| is_vision_identity_tensor(name))
        .collect();
    if vision.len() != V2_VISION_TENSOR_COUNT
        || !delta
            .vision_identity_sha256
            .keys()
            .eq(vision.iter().copied())
    {
        return Err(format!(
            "vision_identity_sha256 lists {} tensors; it must list exactly the {} v1 vision \
             tensors",
            delta.vision_identity_sha256.len(),
            vision.len()
        ));
    }
    if let Some((name, sha256)) = delta
        .vision_identity_sha256
        .iter()
        .find(|(_, sha256)| !is_sha256_hex(sha256))
    {
        return Err(format!(
            "vision tensor {name:?} has hash {sha256:?}, not a SHA-256"
        ));
    }

    // The census is exact: v1 plus the delta re-derives the canonical list of
    // the pinned checkpoint's headers.
    let total = tensors.len() + delta.int64_tensors.len();
    if total != V2_TENSOR_COUNT || delta.tensor_count != V2_TENSOR_COUNT {
        return Err(format!(
            "tensor_count declares {} and v1 plus the delta has {total}, expected \
             {V2_TENSOR_COUNT}",
            delta.tensor_count
        ));
    }
    let mut source_entries = Vec::with_capacity(tensors.len());
    for (name, tensor) in &tensors {
        let dtype = match tensor.source_dtype {
            DType::BF16 => "BF16",
            DType::F16 => "F16",
            DType::F32 => "F32",
            other => return Err(format!("{name:?} has non-source dtype {other:?}")),
        };
        source_entries.push((name.as_str(), dtype, tensor.shape.as_slice()));
    }
    let derived = canonical_census_sha256(
        source_entries.into_iter().chain(
            delta
                .int64_tensors
                .iter()
                .map(|(name, tensor)| (name.as_str(), "I64", tensor.shape.as_slice())),
        ),
    );
    if derived != V2_CENSUS_SHA256 {
        return Err(format!(
            "v1 plus the delta derives census sha256 {derived}, expected {V2_CENSUS_SHA256}"
        ));
    }
    Ok(UnlimitedOcr2Census {
        tensors,
        vision_identity_sha256: delta.vision_identity_sha256,
    })
}

/// The exact non-I64 tensor directory `census` describes.
fn census_tensors(census: TensorCensus) -> FocrResult<&'static BTreeMap<String, ManifestTensor>> {
    match census {
        TensorCensus::UnlimitedOcrV1 => Ok(&manifest()?.tensors),
        TensorCensus::UnlimitedOcrV2 => Ok(&v2_census()?.tensors),
    }
}

/// An Unlimited-OCR-2 `.focrq` must come from the pinned shards (D16) and
/// declare exactly the v2 conservative recipe (U2-C3). The wasm int4 recipe,
/// and every other id, is refused: v2 has no int4 recipe yet (U2-H1). The
/// caller then holds every tensor to the census's storage dtypes.
fn validate_v2_focrq_identity(source_sha256: &str, quant_recipe: Option<&str>) -> FocrResult<()> {
    if source_sha256 != V2_SOURCE_SHA256 {
        return Err(FocrError::FormatMismatch(format!(
            "Unlimited-OCR-2 .focrq source_sha256 is {source_sha256:?}, expected the pinned \
             two-shard source {V2_SOURCE_SHA256:?}"
        )));
    }
    let expected = crate::quant::convert::UNLIMITED_OCR_2_INT8_RECIPE_ID;
    if quant_recipe != Some(expected) {
        return Err(FocrError::FormatMismatch(format!(
            "Unlimited-OCR-2 .focrq packing_manifest.quant_recipe must be exactly {expected:?}, \
             got {quant_recipe:?}; Unlimited-OCR-2 has no int4 or wasm recipe"
        )));
    }
    Ok(())
}

#[derive(Clone, Copy)]
enum CensusSurface {
    SourceSafetensors,
    ConservativeFocrq,
    /// The explicitly-tagged NON-DEFAULT wasm experts-int4 recipe (bd-4l71):
    /// same 2,710-tensor name/shape census, but expert FFN tensors stored
    /// QInt4PerGroup and attention/lm_head/embed stored QInt8PerChan.
    WasmInt4Focrq,
}

impl CensusSurface {
    fn label(self) -> &'static str {
        match self {
            Self::SourceSafetensors => "source safetensors",
            Self::ConservativeFocrq => "conservative .focrq",
            Self::WasmInt4Focrq => "wasm experts-int4 .focrq",
        }
    }

    fn expected_dtype(self, name: &str, tensor: &ManifestTensor) -> DType {
        match self {
            Self::SourceSafetensors => tensor.source_dtype,
            Self::ConservativeFocrq => tensor.storage_dtype,
            // The wasm quantized set is a superset of the conservative one, so
            // every KeepHighPrecision name stores exactly what the conservative
            // manifest stores (BF16) — the storage dtype cross-checks that.
            Self::WasmInt4Focrq => match classify_wasm_experts_int4(name) {
                WasmInt4Policy::ExpertInt4 => DType::QInt4PerGroup,
                WasmInt4Policy::Int8 => DType::QInt8PerChan,
                WasmInt4Policy::KeepHighPrecision => tensor.storage_dtype,
            },
        }
    }
}

#[derive(Default)]
struct NamedDiff {
    total: usize,
    examples: Vec<String>,
}

impl NamedDiff {
    fn push(&mut self, example: String) {
        self.total += 1;
        if self.examples.len() < DIFF_PREVIEW_LIMIT {
            self.examples.push(example);
        }
    }

    fn append_to(&self, output: &mut String, label: &str) {
        if self.total == 0 {
            return;
        }
        use std::fmt::Write;
        let _ = write!(
            output,
            "; {label} {}: [{}",
            self.total,
            self.examples.join(", ")
        );
        if self.total > self.examples.len() {
            let _ = write!(output, ", ... {} more", self.total - self.examples.len());
        }
        output.push(']');
    }
}

fn validate_records<'a>(
    census: TensorCensus,
    records: impl IntoIterator<Item = (&'a str, &'a TensorRecord)>,
    surface: CensusSurface,
) -> FocrResult<()> {
    let expected_tensors = census_tensors(census)?;
    let actual: BTreeMap<&str, &TensorRecord> = records.into_iter().collect();
    let mut missing = NamedDiff::default();
    let mut unexpected = NamedDiff::default();
    let mut wrong_shape = NamedDiff::default();
    let mut wrong_dtype = NamedDiff::default();

    for (name, expected) in expected_tensors {
        let Some(record) = actual.get(name.as_str()) else {
            missing.push(name.clone());
            continue;
        };
        if record.shape != expected.shape {
            wrong_shape.push(format!(
                "{name} expected {:?} got {:?}",
                expected.shape, record.shape
            ));
        }
        let expected_dtype = surface.expected_dtype(name, expected);
        if record.dtype != expected_dtype {
            wrong_dtype.push(format!(
                "{name} expected {expected_dtype:?} got {:?}",
                record.dtype
            ));
        }
    }
    for name in actual.keys() {
        if !expected_tensors.contains_key(*name) {
            unexpected.push((*name).to_owned());
        }
    }

    if missing.total == 0
        && unexpected.total == 0
        && wrong_shape.total == 0
        && wrong_dtype.total == 0
    {
        return Ok(());
    }

    let mut message = format!(
        "{} {} tensor census failed: expected {} tensors, found {}",
        census_model_id(census),
        surface.label(),
        expected_tensors.len(),
        actual.len()
    );
    missing.append_to(&mut message, "missing");
    unexpected.append_to(&mut message, "unexpected");
    wrong_shape.append_to(&mut message, "wrong shape");
    wrong_dtype.append_to(&mut message, "wrong dtype");
    Err(FocrError::FormatMismatch(message))
}

/// Identify a raw safetensors checkpoint, which declares no model id, by the
/// exact embedded census its directory matches. A shard matching none is
/// refused with every census's diff, never assumed to be the default model.
pub(super) fn identify_raw<'a, I>(records: I) -> FocrResult<&'static dyn ModelArch>
where
    I: IntoIterator<Item = (&'a str, &'a TensorRecord)> + Clone,
{
    let mut failures = Vec::new();
    for &arch in model_arch::registry() {
        let Some(census) = arch.tensor_census() else {
            continue;
        };
        match validate_records(census, records.clone(), CensusSurface::SourceSafetensors) {
            Ok(()) => return Ok(arch),
            Err(FocrError::FormatMismatch(diff)) => failures.push(format!("{}: {diff}", arch.id())),
            Err(other) => return Err(other),
        }
    }
    Err(FocrError::FormatMismatch(format!(
        "raw safetensors checkpoint matches no registered model's tensor census [{}]",
        failures.join("; ")
    )))
}

/// Validate identity, recipe, and complete tensor schema of a conservative
/// Unlimited-OCR `.focrq` bounded header for `census`'s model.
pub(super) fn validate_focrq_header(
    census: TensorCensus,
    records: &BTreeMap<String, TensorRecord>,
    source_sha256: &str,
    quant_recipe: Option<&str>,
) -> FocrResult<()> {
    if census == TensorCensus::UnlimitedOcrV2 {
        validate_v2_focrq_identity(source_sha256, quant_recipe)?;
        return validate_records(
            census,
            records.iter().map(|(name, record)| (name.as_str(), record)),
            CensusSurface::ConservativeFocrq,
        );
    }
    let manifest = manifest()?;
    if source_sha256 != manifest.source_sha256 {
        return Err(FocrError::FormatMismatch(format!(
            "Unlimited-OCR .focrq source_sha256 is {source_sha256:?}, expected pinned source {:?}",
            manifest.source_sha256
        )));
    }
    let surface = if quant_recipe == Some(manifest.quant_recipe.as_str()) {
        CensusSurface::ConservativeFocrq
    } else if quant_recipe == Some(UNLIMITED_OCR_WASM_INT4_RECIPE_ID) {
        // The explicitly-tagged non-default wasm recipe (bd-4l71): same pinned
        // source + name/shape census, wasm storage dtypes. Never selected by
        // the default resolver — explicit path / from_weights only.
        CensusSurface::WasmInt4Focrq
    } else {
        return Err(FocrError::FormatMismatch(format!(
            "Unlimited-OCR .focrq packing_manifest.quant_recipe must be exactly {:?} (or the \
             explicitly non-default {UNLIMITED_OCR_WASM_INT4_RECIPE_ID:?}), got {quant_recipe:?}",
            manifest.quant_recipe
        )));
    };
    validate_records(
        census,
        records.iter().map(|(name, record)| (name.as_str(), record)),
        surface,
    )
}

/// Re-run the complete schema check against the directory produced by the real
/// mmap/owned loader, immediately before the production forward receives it.
/// The bounded header path separately verifies the recipe string on the same
/// open descriptor; this post-load check proves name/shape/dtype parity survived
/// the parser handoff.
pub(super) fn validate_loaded_weights(weights: &Weights) -> FocrResult<()> {
    let Some(census) =
        model_arch::arch_by_id(weights.model_id()).and_then(|arch| arch.tensor_census())
    else {
        return Ok(());
    };
    let surface = if weights.is_focrq() {
        if census == TensorCensus::UnlimitedOcrV2 {
            validate_v2_focrq_identity(weights.source_sha256(), weights.quant_recipe())?;
            CensusSurface::ConservativeFocrq
        } else if weights.source_sha256() != SOURCE_SHA256 {
            return Err(FocrError::FormatMismatch(format!(
                "Unlimited-OCR .focrq source_sha256 is {:?}, expected pinned source {SOURCE_SHA256:?}",
                weights.source_sha256()
            )));
        } else if weights.quant_recipe() == Some(UNLIMITED_OCR_WASM_INT4_RECIPE_ID) {
            CensusSurface::WasmInt4Focrq
        } else {
            CensusSurface::ConservativeFocrq
        }
    } else {
        CensusSurface::SourceSafetensors
    };
    let mut records = Vec::with_capacity(weights.len());
    for name in weights.names() {
        let Some(record) = weights.record(name) else {
            return Err(FocrError::FormatMismatch(format!(
                "weights directory lost tensor {name:?} while running the production census"
            )));
        };
        records.push((name, record));
    }
    validate_records(census, records, surface)
}

/// The census's full-tensor SHA-256 of each v2 vision tensor.
#[cfg(test)]
pub(super) fn v2_vision_identity_sha256() -> &'static BTreeMap<String, String> {
    &v2_census()
        .expect("embedded v2 delta verifies")
        .vision_identity_sha256
}

/// A source tensor as `(name, source dtype, shape)`.
#[cfg(test)]
pub(super) type SourceTensor = (String, DType, Vec<usize>);

/// An I64 source tensor as `(name, shape)`; its content is `arange`.
#[cfg(test)]
pub(super) type Int64SourceTensor = (&'static str, &'static [usize]);

/// Every tensor in the pinned v1 source checkpoint, for tests that need a
/// census-exact header without the weights.
#[cfg(test)]
pub(super) fn v1_source_tensors() -> Vec<SourceTensor> {
    manifest()
        .expect("embedded production manifest parses")
        .tensors
        .iter()
        .map(|(name, tensor)| (name.clone(), tensor.source_dtype, tensor.shape.clone()))
        .collect()
}

/// [`v1_source_tensors`] for the v2 census, plus its I64 tensors.
#[cfg(test)]
pub(super) fn v2_source_tensors() -> (Vec<SourceTensor>, Vec<Int64SourceTensor>) {
    let tensors = v2_census()
        .expect("embedded v2 delta verifies")
        .tensors
        .iter()
        .map(|(name, tensor)| (name.clone(), tensor.source_dtype, tensor.shape.clone()))
        .collect();
    let int64 = V2_INT64_RULES
        .iter()
        .map(|rule| (rule.name, rule.shape))
        .collect();
    (tensors, int64)
}

#[cfg(test)]
mod tests {
    use std::io::Read;
    use std::path::PathBuf;

    use sha2::{Digest, Sha256};

    use super::*;

    fn sha256_file(path: &std::path::Path) -> String {
        let mut file = std::fs::File::open(path).expect("open real model artifact");
        let mut digest = Sha256::new();
        let mut buffer = vec![0u8; 16 * 1024 * 1024];
        loop {
            let read = file.read(&mut buffer).expect("hash real model artifact");
            if read == 0 {
                break;
            }
            digest.update(&buffer[..read]);
        }
        format!("{:x}", digest.finalize())
    }

    fn synthetic_directory(surface: CensusSurface) -> BTreeMap<String, TensorRecord> {
        synthetic_census_directory(TensorCensus::UnlimitedOcrV1, surface)
    }

    fn synthetic_census_directory(
        census: TensorCensus,
        surface: CensusSurface,
    ) -> BTreeMap<String, TensorRecord> {
        census_tensors(census)
            .expect("embedded production census parses")
            .iter()
            .map(|(name, tensor)| {
                (
                    name.clone(),
                    TensorRecord {
                        dtype: surface.expected_dtype(name, tensor),
                        shape: tensor.shape.clone(),
                        byte_offset: 0,
                        byte_len: 0,
                        scales_offset: 0,
                        scales_len: 0,
                        group_size: 0,
                        tier: 0,
                    },
                )
            })
            .collect()
    }

    fn identify(directory: &BTreeMap<String, TensorRecord>) -> FocrResult<&'static dyn ModelArch> {
        identify_raw(
            directory
                .iter()
                .map(|(name, record)| (name.as_str(), record)),
        )
    }

    #[test]
    fn raw_checkpoints_are_identified_by_census_not_by_the_default() {
        let source = synthetic_directory(CensusSurface::SourceSafetensors);
        assert_eq!(
            identify(&source).expect("the v1 census matches").id(),
            model_arch::UNLIMITED_OCR_V1_ID
        );
        let got = model_arch::arch_by_id("got-ocr2").expect("got-ocr2 registered");
        let flip = model_arch::test_default::flip(got);
        assert_eq!(
            identify(&source).expect("still v1").id(),
            model_arch::UNLIMITED_OCR_V1_ID,
            "the default model plays no part"
        );
        drop(flip);

        // A shard no census describes (v1 less one tensor, plus a foreign one)
        // is refused with every census's diff, never assumed to be a model.
        let mut foreign = source.clone();
        foreign.remove("lm_head.weight");
        let mut router = source["model.norm.weight"].clone();
        router.shape = vec![32, 1280];
        foreign.insert("model.layers.0.self_attn.router.weight".into(), router);
        let message = identify(&foreign)
            .expect_err("matches no census")
            .to_string();
        assert!(
            message.contains("matches no registered model's tensor census"),
            "{message}"
        );
        for arch in model_arch::registry() {
            if arch.tensor_census().is_some() {
                assert!(message.contains(&format!("{}: ", arch.id())), "{message}");
            }
        }
        assert!(message.contains("missing 1: [lm_head.weight]"), "{message}");
        assert!(
            message.contains("unexpected 1: [model.layers.0.self_attn.router.weight]"),
            "{message}"
        );
    }

    #[test]
    fn v2_census_is_v1_plus_the_pinned_delta() {
        let v1 = &manifest().expect("v1 census").tensors;
        let v2 = v2_census().expect("embedded v2 delta verifies");
        assert_eq!(v2.tensors.len() + V2_INT64_RULES.len(), V2_TENSOR_COUNT);
        for (name, tensor) in v1 {
            let same = &v2.tensors[name];
            assert_eq!(
                (&same.shape, same.source_dtype, same.storage_dtype),
                (&tensor.shape, tensor.source_dtype, tensor.storage_dtype),
                "{name}"
            );
        }
        let added: Vec<&String> = v2.tensors.keys().filter(|n| !v1.contains_key(*n)).collect();
        assert_eq!(added.len(), V2_ADDED_COUNT);
        for layer in 0..12 {
            for module in ["dash_router", "indexer"] {
                for (part, shape) in [
                    ("q_proj.weight", &[128, 1280][..]),
                    ("k_proj.weight", &[128, 1280]),
                    ("q_norm.weight", &[128]),
                    ("k_norm.weight", &[128]),
                ] {
                    let name = format!("model.layers.{layer}.self_attn.{module}.{part}");
                    assert_eq!(v2.tensors[&name].shape, shape, "{name}");
                }
            }
            let head = format!("model.layers.{layer}.self_attn.dash_router.head_cls");
            assert_eq!(v2.tensors[&head].shape, [128], "{head}");
        }
        assert!(
            added
                .iter()
                .all(|name| v2.tensors[*name].source_dtype == DType::BF16)
        );
        // No router or indexer tensor is ever quantized (U2-C3's recipe).
        assert!(
            added
                .iter()
                .all(|name| v2.tensors[*name].storage_dtype == DType::BF16)
        );
        assert_eq!(v2.vision_identity_sha256.len(), V2_VISION_TENSOR_COUNT);
        assert!(
            v2.vision_identity_sha256
                .contains_key("model.view_seperator")
        );
        assert!(
            v2.vision_identity_sha256
                .contains_key("model.projector.layers.weight")
        );

        assert_eq!(
            source_shards(TensorCensus::UnlimitedOcrV2),
            &V2_SOURCE_SHARDS
        );
        assert_eq!(int64_rules(TensorCensus::UnlimitedOcrV2), &V2_INT64_RULES);
        assert!(int64_rules(TensorCensus::UnlimitedOcrV1).is_empty());
        assert_eq!(
            census_model_id(TensorCensus::UnlimitedOcrV2),
            model_arch::UNLIMITED_OCR_V2_ID
        );

        // The multi-shard source convention over the pinned digests (D16).
        let mut concat = Vec::new();
        for (_, sha256) in V2_SOURCE_SHARDS {
            concat.extend(
                (0..32).map(|i| u8::from_str_radix(&sha256[i * 2..i * 2 + 2], 16).unwrap()),
            );
        }
        assert_eq!(V2_SOURCE_SHA256, format!("{:x}", Sha256::digest(&concat)));
        println!(
            "{{\"event\":\"assert\",\"test\":\"v2_census_is_v1_plus_the_pinned_delta\",\
             \"assertion\":\"v2 census = 2710 v1 + 108 added + 1 I64; 476 vision hashes\",\
             \"pass\":true}}"
        );
    }

    #[test]
    fn v2_delta_refuses_every_tampered_field() {
        let pristine: serde_json::Value =
            serde_json::from_str(MANIFEST_2_JSON).expect("embedded delta is JSON");
        let first_added = pristine["added_tensors"]
            .as_object()
            .and_then(|tensors| tensors.keys().next().cloned())
            .expect("an added tensor");
        let first_vision = pristine["vision_identity_sha256"]
            .as_object()
            .and_then(|hashes| hashes.keys().next().cloned())
            .expect("a vision hash");
        let flip_hex = |text: &str| {
            let mut flipped = text.to_owned();
            flipped.replace_range(0..1, if text.starts_with('0') { "1" } else { "0" });
            flipped
        };
        type Mutation = Box<dyn Fn(&mut serde_json::Value)>;
        let cases: Vec<(&str, Mutation, &str)> = vec![
            (
                "model id",
                Box::new(|v| v["model_id"] = "unlimited-ocr".into()),
                "model_id",
            ),
            (
                "base id",
                Box::new(|v| v["base_model_id"] = "got-ocr2".into()),
                "base_model_id",
            ),
            (
                "schema",
                Box::new(|v| v["schema"] = "x.v2".into()),
                "schema",
            ),
            (
                "hf commit",
                Box::new(move |v| {
                    let commit = flip_hex(v["hf_commit"].as_str().unwrap());
                    v["hf_commit"] = commit.into();
                }),
                "hf_commit",
            ),
            (
                "tokenizer",
                Box::new(move |v| {
                    let sha = flip_hex(v["tokenizer_sha256"].as_str().unwrap());
                    v["tokenizer_sha256"] = sha.into();
                }),
                "tokenizer_sha256",
            ),
            (
                "index size",
                Box::new(|v| v["source_index_total_size"] = 1.into()),
                "source_index_total_size",
            ),
            (
                "shard digest",
                Box::new(move |v| {
                    let sha = flip_hex(v["source_shards"][1]["sha256"].as_str().unwrap());
                    v["source_shards"][1]["sha256"] = sha.into();
                }),
                "source_shards",
            ),
            (
                "shard order",
                Box::new(|v| {
                    v["source_shards"].as_array_mut().unwrap().reverse();
                }),
                "source_shards",
            ),
            (
                "recipe",
                Box::new(|v| v["quant_recipe"] = UNLIMITED_OCR_INT8_RECIPE_ID.into()),
                "quant_recipe",
            ),
            (
                "artifact pin",
                Box::new(|v| {
                    v["verified_conservative_artifact_sha256"] = "not-a-digest".into();
                }),
                "verified_conservative_artifact_sha256",
            ),
            (
                "dropped tensor",
                Box::new({
                    let name = first_added.clone();
                    move |v| {
                        v["added_tensors"].as_object_mut().unwrap().remove(&name);
                    }
                }),
                "added_tensors lists 107",
            ),
            (
                "re-added v1 tensor",
                Box::new(|v| {
                    v["added_tensors"]["lm_head.weight"] =
                        serde_json::json!({"shape": [129280, 1280], "source_dtype": "BF16"});
                }),
                "already in the v1 census",
            ),
            (
                "added dtype",
                Box::new({
                    let name = first_added.clone();
                    move |v| v["added_tensors"][&name]["source_dtype"] = "F32".into()
                }),
                "expected BF16",
            ),
            (
                "added shape",
                Box::new({
                    let name = first_added.clone();
                    move |v| v["added_tensors"][&name]["shape"] = serde_json::json!([129])
                }),
                "derives census sha256",
            ),
            (
                "census pin",
                Box::new(move |v| {
                    let sha = flip_hex(v["census_sha256"].as_str().unwrap());
                    v["census_sha256"] = sha.into();
                }),
                "census_sha256",
            ),
            (
                "tensor count",
                Box::new(|v| v["tensor_count"] = 2_818.into()),
                "tensor_count",
            ),
            (
                "int64 content",
                Box::new(|v| {
                    v["int64_tensors"]["model.vision_model.embeddings.position_ids"]["content"] =
                        "zeros".into();
                }),
                "int64_tensors",
            ),
            (
                "int64 extra",
                Box::new(|v| {
                    v["int64_tensors"]["model.other_ids"] =
                        serde_json::json!({"shape": [1, 257], "content": "arange"});
                }),
                "int64_tensors",
            ),
            (
                "vision dropped",
                Box::new({
                    let name = first_vision.clone();
                    move |v| {
                        v["vision_identity_sha256"]
                            .as_object_mut()
                            .unwrap()
                            .remove(&name);
                    }
                }),
                "vision_identity_sha256 lists 475",
            ),
            (
                "vision foreign",
                Box::new(|v| {
                    v["vision_identity_sha256"]["lm_head.weight"] = "00".repeat(32).into();
                }),
                "vision_identity_sha256 lists 477",
            ),
            (
                "vision hash",
                Box::new({
                    let name = first_vision.clone();
                    move |v| v["vision_identity_sha256"][&name] = "XYZ".into()
                }),
                "not a SHA-256",
            ),
        ];
        verify_v2_delta(MANIFEST_2_JSON, MANIFEST_2_SHA256).expect("pristine delta verifies");
        let mismatch = verify_v2_delta(MANIFEST_2_JSON, &"00".repeat(32))
            .err()
            .expect("a wrong pin is refused");
        assert!(mismatch.contains("embedded JSON sha256"), "{mismatch}");
        for (label, mutate, expected) in cases {
            let mut value = pristine.clone();
            mutate(&mut value);
            let json = serde_json::to_string_pretty(&value).expect("re-serialize");
            let pin = format!("{:x}", Sha256::digest(json.as_bytes()));
            let error = verify_v2_delta(&json, &pin)
                .err()
                .unwrap_or_else(|| panic!("{label}: the tampered delta verified"));
            assert!(error.contains(expected), "{label}: {error}");
            println!(
                "{{\"event\":\"assert\",\"test\":\"v2_delta_refuses_every_tampered_field\",\
                 \"assertion\":\"tampered {label} refused\",\"pass\":true}}"
            );
        }
    }

    #[test]
    fn raw_v2_checkpoints_are_identified_by_the_v2_census() {
        let v2 = synthetic_census_directory(
            TensorCensus::UnlimitedOcrV2,
            CensusSurface::SourceSafetensors,
        );
        assert_eq!(
            identify(&v2).expect("the v2 census matches").id(),
            model_arch::UNLIMITED_OCR_V2_ID
        );
        let v1 = synthetic_directory(CensusSurface::SourceSafetensors);
        assert_eq!(
            identify(&v1).expect("v1 is still v1").id(),
            model_arch::UNLIMITED_OCR_V1_ID
        );

        // v2 less one SR-SWA tensor is neither model: v1 sees 107 extras and
        // v2 one missing tensor.
        let mut partial = v2.clone();
        partial.remove("model.layers.5.self_attn.indexer.k_norm.weight");
        let message = identify(&partial)
            .expect_err("matches no census")
            .to_string();
        assert!(
            message.contains(
                "unlimited-ocr-2: unlimited-ocr-2 source safetensors tensor census failed"
            ),
            "{message}"
        );
        assert!(
            message.contains("missing 1: [model.layers.5.self_attn.indexer.k_norm.weight]"),
            "{message}"
        );
        assert!(message.contains("unexpected 107"), "{message}");
    }

    /// U2-C3's loader gate: a v2 `.focrq` passes only with the pinned two-shard
    /// source, exactly the v2 recipe id, and the recipe's storage dtypes; each
    /// refusal is a FormatMismatch (exit 7) naming the field or tensor.
    #[test]
    fn v2_focrq_headers_hold_source_recipe_and_storage_dtypes() {
        let v2_recipe = crate::quant::convert::UNLIMITED_OCR_2_INT8_RECIPE_ID;
        let directory = synthetic_census_directory(
            TensorCensus::UnlimitedOcrV2,
            CensusSurface::ConservativeFocrq,
        );
        validate_focrq_header(
            TensorCensus::UnlimitedOcrV2,
            &directory,
            V2_SOURCE_SHA256,
            Some(v2_recipe),
        )
        .expect("the v2 recipe's own artifact passes");

        let refuse =
            |directory: &BTreeMap<String, TensorRecord>, source: &str, recipe: Option<&str>| {
                let error =
                    validate_focrq_header(TensorCensus::UnlimitedOcrV2, directory, source, recipe)
                        .expect_err("refused");
                assert_eq!(error.exit_code(), 7, "{error}");
                error.to_string()
            };
        let wrong = refuse(&directory, SOURCE_SHA256, Some(v2_recipe));
        assert!(wrong.contains("pinned two-shard source"), "{wrong}");
        for recipe in [
            None,
            Some(UNLIMITED_OCR_INT8_RECIPE_ID),
            Some(UNLIMITED_OCR_WASM_INT4_RECIPE_ID),
        ] {
            let message = refuse(&directory, V2_SOURCE_SHA256, recipe);
            assert!(
                message.contains("quant_recipe must be exactly") && message.contains(v2_recipe),
                "{message}"
            );
        }

        // A quantized router, an int8 attention projection, and a BF16 expert
        // each violate the recipe, by name.
        for (name, dtype) in [
            (
                "model.layers.4.self_attn.dash_router.q_proj.weight",
                DType::QInt8PerChan,
            ),
            (
                "model.layers.9.self_attn.indexer.k_proj.weight",
                DType::QInt8PerChan,
            ),
            (
                "model.layers.2.self_attn.q_proj.weight",
                DType::QInt8PerChan,
            ),
            ("model.layers.6.mlp.experts.11.up_proj.weight", DType::BF16),
        ] {
            let mut tampered = directory.clone();
            tampered.get_mut(name).expect("census tensor").dtype = dtype;
            let message = refuse(&tampered, V2_SOURCE_SHA256, Some(v2_recipe));
            assert!(
                message.contains("wrong dtype 1") && message.contains(name),
                "{message}"
            );
        }
    }

    /// U2-C3: `focr convert` picks the arch by census on the bounded headers
    /// and refuses an explicit `--model-id` that another census matches (exit
    /// 2, naming both), all before any payload byte or shard digest is read.
    #[test]
    fn convert_arch_is_the_census_match_and_cross_checks_model_id() {
        use crate::native_engine::{conversion_arch, shards::SafetensorsHeader};
        let header = |census, drop: Option<&str>| {
            let mut records = synthetic_census_directory(census, CensusSurface::SourceSafetensors);
            if let Some(name) = drop {
                records.remove(name).expect("census tensor");
            }
            SafetensorsHeader {
                records,
                int64: BTreeMap::new(),
            }
        };
        let arch = |id| model_arch::arch_by_id(id).expect("registered");
        let (v1_id, v2_id) = (
            model_arch::UNLIMITED_OCR_V1_ID,
            model_arch::UNLIMITED_OCR_V2_ID,
        );
        let v1 = header(TensorCensus::UnlimitedOcrV1, None);
        let v2 = header(TensorCensus::UnlimitedOcrV2, None);

        assert_eq!(
            conversion_arch(None, &v1).expect("v1 by census").id(),
            v1_id
        );
        assert_eq!(
            conversion_arch(None, &v2).expect("v2 by census").id(),
            v2_id
        );
        for (explicit, source) in [(v1_id, &v1), (v2_id, &v2)] {
            let picked = conversion_arch(Some(arch(explicit)), source).expect("agreeing id");
            assert_eq!(picked.id(), explicit);
        }
        for (explicit, source, found) in [
            (v1_id, &v2, v2_id),
            (v2_id, &v1, v1_id),
            ("got-ocr2", &v2, v2_id),
        ] {
            let error = conversion_arch(Some(arch(explicit)), source).expect_err("disagreeing id");
            assert_eq!(error.exit_code(), 2, "{error}");
            let message = error.to_string();
            assert!(
                message.contains(&format!("--model-id {explicit:?} disagrees"))
                    && message.contains(&format!("{found:?}'s")),
                "{message}"
            );
        }

        // A source no census matches: only an explicit id can name it.
        let partial = header(
            TensorCensus::UnlimitedOcrV2,
            Some("model.layers.5.self_attn.indexer.k_norm.weight"),
        );
        let error = conversion_arch(None, &partial).expect_err("no census, no id");
        assert_eq!(error.exit_code(), 2, "{error}");
        assert!(error.to_string().contains("pass --model-id"), "{error}");
        let error = conversion_arch(Some(arch(v2_id)), &partial).expect_err("census diff");
        assert_eq!(error.exit_code(), 7, "{error}");
        assert!(
            error.to_string().contains("indexer.k_norm.weight"),
            "{error}"
        );
        assert_eq!(
            conversion_arch(Some(arch("got-ocr2")), &partial)
                .expect("a zoo id has no census to disagree with")
                .id(),
            "got-ocr2"
        );
    }

    /// U2-C3's recipe over the full v2 census: exactly 2,148 int8 FFN/expert
    /// matrices and 670 high-precision tensors (the I64 `position_ids` is the
    /// one skipped). No lever ever quantizes a router or indexer tensor, and
    /// the attention lever moves exactly the 48 decoder projections.
    #[test]
    fn v2_recipe_classifies_the_census_exactly() {
        let v2 = &v2_census().expect("embedded v2 census").tensors;
        assert_eq!(v2.len() + V2_INT64_RULES.len(), V2_TENSOR_COUNT);
        let count = |dtype| v2.values().filter(|t| t.storage_dtype == dtype).count();
        assert_eq!(
            (count(DType::QInt8PerChan), count(DType::BF16)),
            (2_148, 670)
        );
        let selectors: Vec<&String> = v2
            .keys()
            .filter(|n| n.contains(".dash_router.") || n.contains(".indexer."))
            .collect();
        assert_eq!(selectors.len(), V2_ADDED_COUNT);
        let baseline = Recipe::validated_default();
        for (attn, lmhead) in [(false, false), (true, false), (false, true), (true, true)] {
            let recipe = Recipe::new(attn, lmhead);
            assert!(
                selectors.iter().all(|n| !recipe.is_quantized(n)),
                "a router/indexer tensor quantized at attn={attn} lmhead={lmhead}"
            );
            let moved: Vec<&String> = v2
                .keys()
                .filter(|n| recipe.is_quantized(n) != baseline.is_quantized(n))
                .collect();
            let attention = moved
                .iter()
                .filter(|n| n.ends_with("_proj.weight") && n.contains(".self_attn."))
                .count();
            assert_eq!(attention, if attn { 48 } else { 0 }, "attn={attn}");
            assert_eq!(moved.len(), attention + usize::from(lmhead), "{moved:?}");
        }
    }

    /// The exact attention matcher leaves v1's recipe output unchanged: over
    /// the v1 census, every lever quantizes the same set the old suffix match
    /// did (no v1 tensor is a router or indexer).
    #[test]
    fn v1_recipe_output_is_unchanged_by_the_exact_attention_matcher() {
        fn legacy_is_attention(name: &str) -> bool {
            name.contains(".self_attn.")
                && [
                    ".q_proj.weight",
                    ".k_proj.weight",
                    ".v_proj.weight",
                    ".o_proj.weight",
                ]
                .iter()
                .any(|leaf| name.ends_with(leaf))
        }
        use crate::quant::recipe::{GatedKind, QuantPolicy, classify};
        let v1 = &manifest().expect("v1 census").tensors;
        // The vision check ran before the attention one, so the old gated set
        // was the suffix match minus the vision tower.
        let legacy: Vec<&String> = v1
            .keys()
            .filter(|n| legacy_is_attention(n))
            .filter(|n| !n.starts_with("model.vision_model.") && !n.starts_with("model.sam_model."))
            .collect();
        let exact: Vec<&String> = v1
            .keys()
            .filter(|n| classify(n).policy == QuantPolicy::Gated(GatedKind::Attention))
            .collect();
        assert_eq!(legacy.len(), 48);
        assert_eq!(exact, legacy);
        for (attn, lmhead) in [(false, false), (true, false), (false, true), (true, true)] {
            let recipe = Recipe::new(attn, lmhead);
            let quantized = v1.keys().filter(|n| recipe.is_quantized(n)).count();
            assert_eq!(
                quantized,
                2_148 + if attn { 48 } else { 0 } + usize::from(lmhead),
                "attn={attn} lmhead={lmhead}"
            );
        }
    }

    #[test]
    fn manifest_matches_pinned_census_and_recipe_exactly() {
        let manifest = manifest().expect("embedded production manifest parses");
        assert_eq!(manifest.source_index_sha256, SOURCE_INDEX_SHA256);
        assert_eq!(manifest.source_index_total_size, SOURCE_INDEX_TOTAL_SIZE);
        assert_eq!(manifest.source_file_size, SOURCE_FILE_SIZE);
        assert_eq!(manifest.tensors.len(), TENSOR_COUNT);
        for landmark in [
            "lm_head.weight",
            "model.embed_tokens.weight",
            "model.layers.11.mlp.experts.63.down_proj.weight",
            "model.vision_model.transformer.layers.23.self_attn.out_proj.weight",
        ] {
            assert!(
                manifest.tensors.contains_key(landmark),
                "embedded production census omits {landmark}"
            );
        }

        let int8_count = manifest
            .tensors
            .values()
            .filter(|tensor| tensor.storage_dtype == DType::QInt8PerChan)
            .count();
        let bf16_count = manifest
            .tensors
            .values()
            .filter(|tensor| tensor.storage_dtype == DType::BF16)
            .count();
        assert_eq!((int8_count, bf16_count), (2_148, 562));
    }

    #[test]
    fn census_rejects_missing_extra_rename_shape_and_dtype_with_names() {
        let valid = synthetic_directory(CensusSurface::ConservativeFocrq);
        validate_records(
            TensorCensus::UnlimitedOcrV1,
            valid.iter().map(|(name, record)| (name.as_str(), record)),
            CensusSurface::ConservativeFocrq,
        )
        .expect("exact synthetic manifest must pass");

        let mut missing = synthetic_directory(CensusSurface::ConservativeFocrq);
        missing.remove("lm_head.weight");
        let error = validate_records(
            TensorCensus::UnlimitedOcrV1,
            missing.iter().map(|(name, record)| (name.as_str(), record)),
            CensusSurface::ConservativeFocrq,
        )
        .expect_err("missing tensor must fail");
        assert!(error.to_string().contains("missing 1: [lm_head.weight]"));

        let mut extra = synthetic_directory(CensusSurface::ConservativeFocrq);
        let record = extra["lm_head.weight"].clone();
        extra.insert("unexpected.weight".into(), record);
        let error = validate_records(
            TensorCensus::UnlimitedOcrV1,
            extra.iter().map(|(name, record)| (name.as_str(), record)),
            CensusSurface::ConservativeFocrq,
        )
        .expect_err("extra tensor must fail");
        assert!(error.to_string().contains("unexpected.weight"));

        let mut renamed = synthetic_directory(CensusSurface::ConservativeFocrq);
        let record = renamed.remove("lm_head.weight").expect("manifest lm_head");
        renamed.insert("lm_heads.weight".into(), record);
        let error = validate_records(
            TensorCensus::UnlimitedOcrV1,
            renamed.iter().map(|(name, record)| (name.as_str(), record)),
            CensusSurface::ConservativeFocrq,
        )
        .expect_err("rename must fail as a bounded missing/extra diff");
        let text = error.to_string();
        assert!(text.contains("lm_head.weight"));
        assert!(text.contains("lm_heads.weight"));

        let mut wrong_shape = synthetic_directory(CensusSurface::ConservativeFocrq);
        wrong_shape
            .get_mut("lm_head.weight")
            .expect("manifest lm_head")
            .shape[0] -= 1;
        let error = validate_records(
            TensorCensus::UnlimitedOcrV1,
            wrong_shape
                .iter()
                .map(|(name, record)| (name.as_str(), record)),
            CensusSurface::ConservativeFocrq,
        )
        .expect_err("wrong shape must fail");
        assert!(error.to_string().contains("wrong shape 1: [lm_head.weight"));

        let mut wrong_dtype = synthetic_directory(CensusSurface::ConservativeFocrq);
        wrong_dtype
            .get_mut("lm_head.weight")
            .expect("manifest lm_head")
            .dtype = DType::QInt8PerChan;
        let error = validate_records(
            TensorCensus::UnlimitedOcrV1,
            wrong_dtype
                .iter()
                .map(|(name, record)| (name.as_str(), record)),
            CensusSurface::ConservativeFocrq,
        )
        .expect_err("wrong dtype must fail");
        assert!(error.to_string().contains("wrong dtype 1: [lm_head.weight"));
    }

    #[test]
    fn focrq_identity_rejects_wrong_source_and_recipe() {
        let directory = synthetic_directory(CensusSurface::ConservativeFocrq);
        validate_focrq_header(
            TensorCensus::UnlimitedOcrV1,
            &directory,
            SOURCE_SHA256,
            Some(UNLIMITED_OCR_INT8_RECIPE_ID),
        )
        .expect("pinned identity must pass");

        let source_error = validate_focrq_header(
            TensorCensus::UnlimitedOcrV1,
            &directory,
            &"00".repeat(32),
            Some(UNLIMITED_OCR_INT8_RECIPE_ID),
        )
        .expect_err("wrong source must fail");
        assert!(source_error.to_string().contains("source_sha256"));

        let recipe_error = validate_focrq_header(
            TensorCensus::UnlimitedOcrV1,
            &directory,
            SOURCE_SHA256,
            Some("legacy-full-int8"),
        )
        .expect_err("wrong recipe must fail");
        assert!(recipe_error.to_string().contains("quant_recipe"));
    }

    /// bd-4l71: the explicitly-tagged wasm experts-int4 recipe is accepted with
    /// its OWN storage dtypes (experts QInt4PerGroup, attn/lm_head/embed
    /// QInt8PerChan, rest BF16) — and each surface rejects the other's dtypes.
    #[test]
    fn wasm_int4_focrq_surface_accepts_wasm_dtypes_and_rejects_cross_recipe() {
        let wasm = synthetic_directory(CensusSurface::WasmInt4Focrq);
        // Sanity: the wasm directory really stores the three storage classes.
        assert_eq!(
            wasm["model.layers.11.mlp.experts.63.down_proj.weight"].dtype,
            DType::QInt4PerGroup
        );
        assert_eq!(wasm["lm_head.weight"].dtype, DType::QInt8PerChan);
        assert_eq!(wasm["model.embed_tokens.weight"].dtype, DType::QInt8PerChan);
        assert_eq!(
            wasm["model.layers.0.self_attn.q_proj.weight"].dtype,
            DType::QInt8PerChan
        );
        assert_eq!(wasm["model.norm.weight"].dtype, DType::BF16);
        assert_eq!(
            wasm["model.layers.5.mlp.gate.weight"].dtype,
            DType::BF16,
            "router gate must stay high precision even under the wasm recipe"
        );

        validate_focrq_header(
            TensorCensus::UnlimitedOcrV1,
            &wasm,
            SOURCE_SHA256,
            Some(crate::quant::convert::UNLIMITED_OCR_WASM_INT4_RECIPE_ID),
        )
        .expect("wasm directory + wasm recipe id must pass");

        // Wasm dtypes under the CONSERVATIVE recipe id must fail (the default
        // recipe is untouched)...
        let error = validate_focrq_header(
            TensorCensus::UnlimitedOcrV1,
            &wasm,
            SOURCE_SHA256,
            Some(UNLIMITED_OCR_INT8_RECIPE_ID),
        )
        .expect_err("wasm dtypes must fail the conservative surface");
        assert!(error.to_string().contains("wrong dtype"));

        // ...and conservative dtypes under the wasm recipe id must fail too.
        let conservative = synthetic_directory(CensusSurface::ConservativeFocrq);
        let error = validate_focrq_header(
            TensorCensus::UnlimitedOcrV1,
            &conservative,
            SOURCE_SHA256,
            Some(crate::quant::convert::UNLIMITED_OCR_WASM_INT4_RECIPE_ID),
        )
        .expect_err("conservative dtypes must fail the wasm surface");
        assert!(error.to_string().contains("wrong dtype"));
    }

    #[test]
    fn real_conservative_artifact_matches_manifest_when_configured() {
        let Some(path) = std::env::var_os("FOCR_TEST_UNLIMITED_RECIPE_MODEL").map(PathBuf::from)
        else {
            eprintln!(
                "skip-with-SUCCESS: set FOCR_TEST_UNLIMITED_RECIPE_MODEL to the verified .focrq"
            );
            return;
        };

        assert_eq!(sha256_file(&path), VERIFIED_ARTIFACT_SHA256);

        let (file, file_len) = super::super::open_model_file(&path).expect("open model descriptor");
        super::super::validate_model_header_from_reader(&file, file_len)
            .expect("bounded real header must match the production census");
        let weights = Weights::load_opened(file, &path).expect("mmap real artifact");
        validate_loaded_weights(&weights).expect("real loaded directory must match the census");
        assert_eq!(weights.len(), TENSOR_COUNT);
    }

    /// U2-C3 round trip on the real v2 artifact (model-gated). The converted
    /// `.focrq` passes the bounded header probe, the post-load census and the
    /// recipe check. Given the source shards too, every high-precision tensor
    /// is byte-equal to its source, and a sample of the int8 matrices (dense
    /// layer 0, expert 0 and the shared experts of every MoE layer) dequantizes
    /// to within half a quantization step of the source.
    #[test]
    fn real_v2_artifact_round_trips_when_configured() {
        let Some(path) = std::env::var_os("FOCR_TEST_UNLIMITED_2_RECIPE_MODEL").map(PathBuf::from)
        else {
            eprintln!(
                "skip-with-SUCCESS: set FOCR_TEST_UNLIMITED_2_RECIPE_MODEL to the converted v2 \
                 .focrq (and FOCR_TEST_UNLIMITED_2_SOURCE_DIR to its shards for the tensor check)"
            );
            return;
        };
        let (file, file_len) = super::super::open_model_file(&path).expect("open v2 artifact");
        let arch = super::super::validate_model_header_from_reader(&file, file_len)
            .expect("bounded v2 header must pass the census and recipe");
        assert_eq!(arch.id(), model_arch::UNLIMITED_OCR_V2_ID);
        let weights = Weights::load_opened(file, &path).expect("load v2 artifact");
        validate_loaded_weights(&weights).expect("loaded v2 directory must match the census");
        super::super::validate_unlimited_ocr_quant_recipe(&weights)
            .expect("loaded v2 directory must hold the recipe");
        assert_eq!(weights.len(), V2_TENSOR_COUNT - V2_INT64_RULES.len());
        assert_eq!(
            weights.quant_recipe(),
            Some(crate::quant::convert::UNLIMITED_OCR_2_INT8_RECIPE_ID)
        );
        assert_eq!(weights.source_sha256(), V2_SOURCE_SHA256);

        let Some(dir) = std::env::var_os("FOCR_TEST_UNLIMITED_2_SOURCE_DIR").map(PathBuf::from)
        else {
            eprintln!(
                "skip-with-SUCCESS: FOCR_TEST_UNLIMITED_2_SOURCE_DIR unset; tensors not compared"
            );
            return;
        };
        let sampled = |name: &str| {
            name.starts_with("model.layers.0.mlp.")
                || name.contains(".mlp.experts.0.")
                || name.contains(".mlp.shared_experts.")
        };
        let (mut kept, mut int8, mut checked) = (0usize, 0usize, 0usize);
        for (shard, _) in V2_SOURCE_SHARDS {
            let bytes = std::fs::read(dir.join(shard)).expect("read source shard");
            let source = Weights::from_source_shard(bytes, arch).expect("parse source shard");
            for name in source.names() {
                let (src, out) = (
                    source.tensor(name).expect("source tensor"),
                    weights.tensor(name).expect("artifact tensor"),
                );
                assert_eq!(src.shape, out.shape, "{name}");
                if out.dtype != DType::QInt8PerChan {
                    assert_eq!((src.dtype, src.data), (out.dtype, out.data), "{name}");
                    kept += 1;
                    continue;
                }
                int8 += 1;
                if !sampled(name) {
                    continue;
                }
                let (w, q) = (
                    source.mat(name).expect("source f32"),
                    weights.qint8(name).expect("artifact int8"),
                );
                let deq = weights.mat(name).expect("artifact dequant");
                for row in 0..q.n {
                    let step = q.scales[row];
                    let span = row * q.k..(row + 1) * q.k;
                    for (a, b) in w.data[span.clone()].iter().zip(&deq.data[span]) {
                        assert!(
                            (a - b).abs() <= step * 0.500_01 + a.abs() * 1e-6,
                            "{name} row {row}: source {a} vs dequantized {b} (step {step})"
                        );
                    }
                }
                checked += 1;
            }
        }
        assert_eq!((int8, kept), (2_148, 670));
        assert_eq!(checked, 3 + 11 * 3 + 11 * 3);
    }

    #[test]
    fn real_source_shard_matches_manifest_when_configured() {
        let Some(path) = std::env::var_os("FOCR_TEST_UNLIMITED_SOURCE_MODEL").map(PathBuf::from)
        else {
            eprintln!(
                "skip-with-SUCCESS: set FOCR_TEST_UNLIMITED_SOURCE_MODEL to the verified shard"
            );
            return;
        };

        assert_eq!(sha256_file(&path), SOURCE_SHA256);
        assert_eq!(
            std::fs::metadata(&path)
                .expect("stat real source shard")
                .len(),
            SOURCE_FILE_SIZE,
            "full source shard size includes its safetensors prefix and JSON header"
        );
        let (file, file_len) =
            super::super::open_model_file(&path).expect("open source descriptor");
        super::super::validate_model_header_from_reader(&file, file_len)
            .expect("bounded source header must match the production census");
        let weights = Weights::load_opened(file, &path).expect("mmap real source shard");
        validate_loaded_weights(&weights).expect("real source directory must match the census");
        assert_eq!(weights.len(), TENSOR_COUNT);
    }

    #[test]
    fn conversion_source_identity_fails_before_emitting_an_unloadable_artifact() {
        let mut expected = [0u8; 32];
        for (index, byte) in expected.iter_mut().enumerate() {
            *byte = u8::from_str_radix(&SOURCE_SHA256[index * 2..index * 2 + 2], 16)
                .expect("source SHA constant is hex");
        }
        let census = TensorCensus::UnlimitedOcrV1;
        validate_conversion_shard_count(census, 1).expect("the pinned source is one shard");
        validate_conversion_source_shard(census, 0, &expected).expect("pinned source accepted");
        let error = validate_conversion_source_shard(census, 0, &[0u8; 32])
            .expect_err("reserialized source must fail before conversion output");
        let message = error.to_string();
        assert!(message.contains(SOURCE_SHA256));
        assert!(message.contains(SOURCE_FILE));
        assert!(message.contains(&"0".repeat(64)));
        let error = validate_conversion_shard_count(census, 2)
            .expect_err("a resharded source is not the pinned checkpoint");
        assert!(error.to_string().contains("1-shard"), "{error}");
        let error = validate_conversion_source_shard(census, 1, &expected)
            .expect_err("a second shard has no pin");
        assert!(error.to_string().contains("2 shard(s)"), "{error}");
    }
}