"""Export the final MOSS quantized topology to a self-verifying CPU bundle.

The deployment topology is deliberately fixed: Sherry decoder projections use
group size 128, Whisper audio linears use W4 RTN group size 128, and the tied
input embedding/output head uses W3 RTN group size 64.  Packed weights and fp16
scales live in ``weights.pt``; ``manifest.json`` describes every state tensor,
aliases the omitted tied ``lm_head.weight``, and reports payload sizes.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import re
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any, Mapping

import torch
import torch.nn as nn
import torch.nn.functional as F

from quantlib import (
    PackedRTN,
    PackedSTQ1,
    RTNEmbedding,
    RTNLinear,
    SherryLinear,
    TiedRTNLMHead,
    nm_quantize,
    pack,
    rtn_quantize,
    tie_rtn_lm_head,
    unpack,
    unpack_rtn,
    wrap_rtn,
    wrap_sherry,
)
from quantlib.checkpoint_audit import audit_quantized_student_state


SCHEMA = "moss-quantized-export-v1"
WEIGHTS_FILE = "weights.pt"
MANIFEST_FILE = "manifest.json"
SHERRY_GROUP_SIZE = 128
AUDIO_BITS = 4
AUDIO_GROUP_SIZE = 128
PROFILE = "final_w3"
EMBED_BITS = 3
EMBED_GRANULARITY = "per_group"
EMBED_GROUP_SIZE = 64
EMBED_CHECKPOINT_GROUP_SIZE: int | None = 64
EMBED_KIND = "rtn_w3"
EXPECTED_SHERRY_LAYERS = 28 * 7
EXPECTED_AUDIO_LAYERS = 24 * 6
CHECKPOINT_FORMATS = {"moss-sherry-kd-qat-v1", "moss-sherry-vespo-v1"}
EMBED_WEIGHT = "model.language_model.embed_tokens.weight"
LM_HEAD_WEIGHT = "lm_head.weight"
AUDIO_LINEAR_PATTERN = (
    r"re:^model\.whisper_encoder\.layers\.\d+\."
    r"(?:self_attn\.(?:q_proj|k_proj|v_proj|out_proj)|fc1|fc2)$"
)
SHERRY_LINEAR_RE = re.compile(
    r"^model\.language_model\.layers\.\d+\."
    r"(?:self_attn\.(?:q_proj|k_proj|v_proj|o_proj)|"
    r"mlp\.(?:gate_proj|up_proj|down_proj))$"
)
AUDIO_LINEAR_RE = re.compile(AUDIO_LINEAR_PATTERN.removeprefix("re:"))


def configure_profile(profile: str) -> None:
    """Select an explicit deployment embedding topology.

    The historical W3 profile remains the default for API compatibility.  W4
    checkpoints use the original per-channel embedding quantizer; its stored
    ``group_size`` is 128 but per-channel RTN uses the full embedding row as
    the effective group.
    """

    global PROFILE, EMBED_BITS, EMBED_GRANULARITY, EMBED_GROUP_SIZE
    global EMBED_CHECKPOINT_GROUP_SIZE, EMBED_KIND
    if profile == "final_w3":
        PROFILE = profile
        EMBED_BITS = 3
        EMBED_GRANULARITY = "per_group"
        EMBED_GROUP_SIZE = 64
        EMBED_CHECKPOINT_GROUP_SIZE = 64
        EMBED_KIND = "rtn_w3"
    elif profile == "e2e_w4":
        PROFILE = profile
        EMBED_BITS = 4
        EMBED_GRANULARITY = "per_channel"
        EMBED_GROUP_SIZE = 128
        EMBED_CHECKPOINT_GROUP_SIZE = None
        EMBED_KIND = "rtn_embed_w4"
    else:
        raise ValueError(f"unsupported export profile: {profile!r}")


def _tensor_bytes(tensor: torch.Tensor) -> int:
    return tensor.numel() * tensor.element_size()


def _file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _tensor_checksum(tensor: torch.Tensor) -> str:
    raw = (
        tensor.detach()
        .cpu()
        .contiguous()
        .reshape(-1)
        .view(torch.uint8)
        .numpy()
        .tobytes()
    )
    return hashlib.sha256(raw).hexdigest()


def _atomic_torch_save(path: Path, value: Any) -> None:
    temporary = path.with_suffix(path.suffix + ".tmp")
    torch.save(value, temporary)
    temporary.replace(path)


def _atomic_json(path: Path, value: Any) -> None:
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(
        json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    temporary.replace(path)


def _module_weight_name(module_name: str) -> str:
    return f"{module_name}.weight" if module_name else "weight"


def _validate_fixed_module(module_name: str, module: nn.Module) -> str:
    if isinstance(module, SherryLinear):
        if SHERRY_LINEAR_RE.fullmatch(module_name) is None:
            raise ValueError(f"Sherry module is outside decoder projections: {module_name!r}")
        if module.group_size != SHERRY_GROUP_SIZE:
            raise ValueError(
                f"Sherry module {module_name!r} must use group size {SHERRY_GROUP_SIZE}"
            )
        if float(module.eps.detach().cpu()) != 0.0:
            raise ValueError(f"Sherry module {module_name!r} must have eps=0 for export")
        return "sherry_stq1"
    if isinstance(module, RTNLinear):
        if AUDIO_LINEAR_RE.fullmatch(module_name) is None:
            raise ValueError(f"W4 RTN module is outside the audio encoder: {module_name!r}")
        expected = (AUDIO_BITS, "per_group", AUDIO_GROUP_SIZE)
        actual = (module.bits, module.granularity, module.group_size)
        if actual != expected:
            raise ValueError(f"audio RTN module {module_name!r} has {actual}, expected {expected}")
        return "rtn_w4"
    if isinstance(module, RTNEmbedding):
        expected = (EMBED_BITS, EMBED_GRANULARITY, EMBED_GROUP_SIZE)
        actual = (module.bits, module.granularity, module.group_size)
        if actual != expected:
            raise ValueError(f"embedding RTN module {module_name!r} has {actual}, expected {expected}")
        if _module_weight_name(module_name) != EMBED_WEIGHT:
            raise ValueError(f"quantized embedding must be {EMBED_WEIGHT!r}, got {module_name!r}")
        return EMBED_KIND
    raise TypeError(f"unsupported quantized module: {type(module).__name__}")


def build_export_plan(model: nn.Module) -> dict[str, Any]:
    """Classify unique state tensors and audit the fixed deployment topology."""

    quantized: dict[str, dict[str, Any]] = {}
    counts: defaultdict[str, int] = defaultdict(int)
    embedding: RTNEmbedding | None = None
    head: TiedRTNLMHead | None = None
    for module_name, module in model.named_modules():
        if isinstance(module, (SherryLinear, RTNLinear, RTNEmbedding)):
            kind = _validate_fixed_module(module_name, module)
            name = _module_weight_name(module_name)
            quantized[name] = {"kind": kind, "module_name": module_name, "module": module}
            counts[kind] += 1
            if isinstance(module, RTNEmbedding):
                if embedding is not None:
                    raise ValueError("exactly one RTN embedding is supported")
                embedding = module
        elif isinstance(module, TiedRTNLMHead):
            if module_name != "lm_head":
                raise ValueError(f"tied RTN head must be 'lm_head', got {module_name!r}")
            head = module

    if not quantized:
        raise ValueError("model contains no supported quantized modules")
    if embedding is None or head is None:
        raise ValueError("fixed topology requires an RTN embedding and TiedRTNLMHead")
    if head.weight is not embedding.weight or head.weight.data_ptr() != embedding.weight.data_ptr():
        raise ValueError("embedding and lm_head weights are not tied")

    state = model.state_dict()
    if EMBED_WEIGHT not in state or LM_HEAD_WEIGHT not in state:
        raise ValueError("tied embedding/head state keys are missing")
    if state[EMBED_WEIGHT].data_ptr() != state[LM_HEAD_WEIGHT].data_ptr():
        raise ValueError("embedding/head state tensors do not share storage")
    missing = sorted(set(quantized) - set(state))
    if missing:
        raise ValueError(f"quantized weights missing from state_dict: {missing}")

    component_names = [name for name in state if name != LM_HEAD_WEIGHT]
    return {
        "state": state,
        "quantized": quantized,
        "component_names": component_names,
        "aliases": {LM_HEAD_WEIGHT: EMBED_WEIGHT},
        "counts": dict(counts),
    }


def _canonical_fakequant(
    kind: str,
    module: nn.Module,
    canonical_device: torch.device,
) -> torch.Tensor:
    master = module.weight.detach().to(device=canonical_device, dtype=torch.float32)
    if kind == "sherry_stq1":
        return nm_quantize(master, SHERRY_GROUP_SIZE).to(torch.bfloat16).cpu()
    if kind == "rtn_w4":
        return rtn_quantize(master, AUDIO_BITS, "per_group", AUDIO_GROUP_SIZE).to(
            torch.bfloat16
        ).cpu()
    if kind == EMBED_KIND:
        return rtn_quantize(
            master, EMBED_BITS, EMBED_GRANULARITY, EMBED_GROUP_SIZE
        ).to(torch.bfloat16).cpu()
    raise ValueError(f"unknown quantized component kind: {kind}")


def _pack_rtn_codes(codes: torch.Tensor, bits: int) -> torch.Tensor:
    """Pack logical signed RTN codes without trying to infer their scale."""

    if bits == 4:
        if codes.numel() % 2:
            codes = F.pad(codes, (0, 1))
        nibbles = codes.to(torch.int16) & 0xF
        return (nibbles[0::2] | (nibbles[1::2] << 4)).to(torch.uint8)
    logical_codes = codes.numel()
    if logical_codes % 8:
        codes = F.pad(codes, (0, 8 - logical_codes % 8))
    values = (codes.to(torch.int16) & 0x7).reshape(-1, 8)
    payload = torch.empty((values.shape[0], 3), dtype=torch.int16)
    payload[:, 0] = values[:, 0] | (values[:, 1] << 3) | (values[:, 2] << 6)
    payload[:, 1] = (
        (values[:, 2] >> 2)
        | (values[:, 3] << 1)
        | (values[:, 4] << 4)
        | (values[:, 5] << 7)
    )
    payload[:, 2] = (
        (values[:, 5] >> 1) | (values[:, 6] << 2) | (values[:, 7] << 5)
    )
    return payload.reshape(-1)[: (logical_codes * 3 + 7) // 8].to(torch.uint8)


def _unpack_rtn_codes(data: torch.Tensor, logical_codes: int, bits: int) -> torch.Tensor:
    """Decode payload codes for range and canonical-padding verification."""

    payload = data.to(torch.int16)
    if bits == 4:
        values = torch.empty(payload.numel() * 2, dtype=torch.int16)
        values[0::2] = payload & 0xF
        values[1::2] = payload >> 4
        values = values[:logical_codes]
        return torch.where(values >= 8, values - 16, values)
    padding = (-payload.numel()) % 3
    if padding:
        payload = F.pad(payload, (0, padding))
    packed = payload.reshape(-1, 3)
    values = torch.empty((packed.shape[0], 8), dtype=torch.int16)
    values[:, 0] = packed[:, 0] & 0x7
    values[:, 1] = (packed[:, 0] >> 3) & 0x7
    values[:, 2] = ((packed[:, 0] >> 6) | (packed[:, 1] << 2)) & 0x7
    values[:, 3] = (packed[:, 1] >> 1) & 0x7
    values[:, 4] = (packed[:, 1] >> 4) & 0x7
    values[:, 5] = ((packed[:, 1] >> 7) | (packed[:, 2] << 1)) & 0x7
    values[:, 6] = (packed[:, 2] >> 2) & 0x7
    values[:, 7] = (packed[:, 2] >> 5) & 0x7
    values = values.reshape(-1)[:logical_codes]
    return torch.where(values >= 4, values - 8, values)


def _pack_rtn_master(
    weight: torch.Tensor,
    bits: int,
    group_size: int,
    granularity: str = "per_group",
    quantization_device: str | torch.device = "cpu",
) -> PackedRTN:
    """Pack codes and fp16 scales directly from an fp32 RTN master."""

    master = weight.detach().to(
        device=torch.device(quantization_device), dtype=torch.float32
    ).contiguous()
    rows, columns = master.shape
    actual_group_size = columns if granularity == "per_channel" else group_size
    number_of_groups = (columns + actual_group_size - 1) // actual_group_size
    padded_columns = number_of_groups * actual_group_size
    groups = F.pad(master, (0, padded_columns - columns)).reshape(
        rows, number_of_groups, actual_group_size
    )
    qmax = 2 ** (bits - 1) - 1
    scales = (groups.abs().amax(dim=-1) / qmax).to(torch.float16)
    safe_scales = torch.where(scales == 0, torch.ones_like(scales), scales).float()
    codes = torch.round(groups / safe_scales[..., None]).clamp(-qmax, qmax)
    logical_codes = codes.to(torch.int8).reshape(rows, padded_columns)[:, :columns]
    return PackedRTN(
        data=_pack_rtn_codes(logical_codes.reshape(-1).cpu(), bits),
        scales=scales.cpu().contiguous(),
        shape=(rows, columns),
        bits=bits,
        granularity=granularity,
        group_size=group_size,
        original_dtype="bfloat16",
    )


def _packed_component(
    name: str,
    kind: str,
    module: nn.Module,
    canonical_device: torch.device,
) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
    fakequant = _canonical_fakequant(kind, module, canonical_device)
    if kind == "sherry_stq1":
        packed = pack(fakequant, SHERRY_GROUP_SIZE)
        metadata = {
            "shape": list(packed.shape),
            "group_size": packed.group_size,
            "original_dtype": packed.original_dtype,
        }
    else:
        bits = AUDIO_BITS if kind == "rtn_w4" else EMBED_BITS
        granularity = "per_group" if kind == "rtn_w4" else EMBED_GRANULARITY
        group_size = AUDIO_GROUP_SIZE if kind == "rtn_w4" else EMBED_GROUP_SIZE
        # Preserve the integer codes selected on the evaluation device.
        # Recomputing them from the FP32 master on CPU can disagree with CUDA at
        # exact rounding boundaries, even though both paths are numerically
        # close.  Packing from the BF16 grid alone is also insufficient because
        # BF16 output rounding can obscure the original integer/scale pair.
        packed = _pack_rtn_master(
            module.weight,
            bits,
            group_size,
            granularity,
            canonical_device,
        )
        metadata = {
            "shape": list(packed.shape),
            "bits": packed.bits,
            "granularity": packed.granularity,
            "group_size": packed.group_size,
            "original_dtype": packed.original_dtype,
        }
    data_bytes = _tensor_bytes(packed.data)
    scale_bytes = _tensor_bytes(packed.scales)
    logical_values = packed.num_weights
    decoded = unpack(packed) if kind == "sherry_stq1" else unpack_rtn(packed)
    canonical_bf16 = fakequant.to(torch.bfloat16)
    if not torch.equal(decoded.to(torch.bfloat16), canonical_bf16):
        raise RuntimeError(f"packed decode differs from canonical BF16 fakequant for {name!r}")
    artifact = {"data": packed.data.cpu(), "scales": packed.scales.cpu()}
    description = {
        "name": name,
        "kind": kind,
        "logical_values": logical_values,
        "data_bytes": data_bytes,
        "scale_bytes": scale_bytes,
        "stored_bytes": data_bytes + scale_bytes,
        "effective_bits": packed.effective_bits,
        "packing_source": (
            f"{canonical_device.type}_fp32_master_to_bf16_eval_grid"
            if kind == "sherry_stq1"
            else f"{canonical_device.type}_fp32_master_to_bf16_eval_grid"
        ),
        "state_dtype": "bfloat16",
        "canonicalization": (
            "sherry_fp32_mask_bf16_grid_v1"
            if kind == "sherry_stq1"
            else "rtn_eval_device_grid_fp16_scale_bf16_state_v2"
        ),
        "data_sha256": _tensor_checksum(packed.data),
        "scales_sha256": _tensor_checksum(packed.scales),
        **metadata,
    }
    return artifact, description


def _plain_component(name: str, tensor: torch.Tensor) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
    cpu = tensor.detach().cpu().contiguous()
    if cpu.is_floating_point():
        cpu = cpu.to(torch.bfloat16)
        kind = "bf16"
    else:
        kind = "tensor"
    stored_bytes = _tensor_bytes(cpu)
    return {"value": cpu}, {
        "name": name,
        "kind": kind,
        "shape": list(cpu.shape),
        "dtype": str(cpu.dtype).removeprefix("torch."),
        "logical_values": cpu.numel(),
        "stored_bytes": stored_bytes,
        "effective_bits": cpu.element_size() * 8.0,
        "sha256": _tensor_checksum(cpu),
    }


def _summary(components: list[dict[str, Any]]) -> dict[str, Any]:
    logical_values = sum(int(component["logical_values"]) for component in components)
    stored_bytes = sum(int(component["stored_bytes"]) for component in components)
    packed_data_bytes = sum(int(component.get("data_bytes", 0)) for component in components)
    scale_bytes = sum(int(component.get("scale_bytes", 0)) for component in components)
    by_kind: dict[str, dict[str, int | float]] = {}
    for kind in sorted({str(component["kind"]) for component in components}):
        selected = [component for component in components if component["kind"] == kind]
        kind_values = sum(int(component["logical_values"]) for component in selected)
        kind_bytes = sum(int(component["stored_bytes"]) for component in selected)
        by_kind[kind] = {
            "components": len(selected),
            "logical_values": kind_values,
            "stored_bytes": kind_bytes,
            "effective_bits": 8.0 * kind_bytes / kind_values,
        }
    return {
        "components": len(components),
        "logical_values": logical_values,
        "stored_bytes": stored_bytes,
        "packed_data_bytes": packed_data_bytes,
        "scale_bytes": scale_bytes,
        "plain_tensor_bytes": stored_bytes - packed_data_bytes - scale_bytes,
        "effective_bits": 8.0 * stored_bytes / logical_values,
        "by_kind": by_kind,
    }


@torch.no_grad()
def export_model(
    model: nn.Module,
    output_dir: str | Path,
    *,
    source: Mapping[str, Any] | None = None,
    canonical_device: str | torch.device = "cpu",
) -> dict[str, Any]:
    """Export a rebuilt model and return its complete JSON manifest."""

    destination = Path(output_dir).expanduser().resolve()
    destination.mkdir(parents=True, exist_ok=True)
    quant_device = torch.device(canonical_device)
    if quant_device.type not in {"cpu", "cuda"}:
        raise ValueError(f"unsupported canonical quantization device: {quant_device}")
    if quant_device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA canonical export requested but CUDA is unavailable")
    plan = build_export_plan(model)
    artifacts: dict[str, dict[str, torch.Tensor]] = {}
    components: list[dict[str, Any]] = []
    for name in plan["component_names"]:
        if name in plan["quantized"]:
            specification = plan["quantized"][name]
            artifact, component = _packed_component(
                name,
                specification["kind"],
                specification["module"],
                quant_device,
            )
        else:
            artifact, component = _plain_component(name, plan["state"][name])
        artifacts[name] = artifact
        components.append(component)

    manifest = {
        "schema": SCHEMA,
        "weights_file": WEIGHTS_FILE,
        "topology": {
            "profile": PROFILE,
            "canonical_quantization_device": str(quant_device),
            "sherry_group_size": SHERRY_GROUP_SIZE,
            "audio": {"bits": AUDIO_BITS, "granularity": "per_group", "group_size": AUDIO_GROUP_SIZE},
            "embedding": {"bits": EMBED_BITS, "granularity": EMBED_GRANULARITY, "group_size": EMBED_GROUP_SIZE},
            "tied_lm_head": True,
            "counts": plan["counts"],
        },
        "source": dict(source or {}),
        "aliases": plan["aliases"],
        "components": components,
        "summary": _summary(components),
    }
    _atomic_torch_save(
        destination / WEIGHTS_FILE,
        {"schema": SCHEMA, "components": artifacts},
    )
    _atomic_json(destination / MANIFEST_FILE, manifest)
    return manifest


def _read_bundle(export_dir: str | Path) -> tuple[Path, dict[str, Any], dict[str, Any]]:
    directory = Path(export_dir).expanduser().resolve()
    manifest = json.loads((directory / MANIFEST_FILE).read_text(encoding="utf-8"))
    if manifest.get("schema") != SCHEMA:
        raise ValueError(f"unsupported export schema: {manifest.get('schema')!r}")
    weights_name = manifest.get("weights_file")
    if weights_name != WEIGHTS_FILE:
        raise ValueError(f"unexpected weights file: {weights_name!r}")
    document = torch.load(directory / weights_name, map_location="cpu", weights_only=True)
    if not isinstance(document, dict) or document.get("schema") != SCHEMA:
        raise ValueError("weights artifact schema does not match manifest")
    artifacts = document.get("components")
    if not isinstance(artifacts, dict):
        raise ValueError("weights artifact has no component mapping")
    return directory, manifest, artifacts


def _check_checksum(tensor: torch.Tensor, expected: str, label: str) -> None:
    actual = _tensor_checksum(tensor)
    if actual != expected:
        raise ValueError(f"checksum mismatch for {label}: expected {expected}, got {actual}")


def _decode_packed(component: dict[str, Any], artifact: dict[str, torch.Tensor]) -> torch.Tensor:
    name = component["name"]
    if set(artifact) != {"data", "scales"}:
        raise ValueError(f"packed component {name!r} has invalid artifact fields")
    data, scales = artifact["data"], artifact["scales"]
    _check_checksum(data, component["data_sha256"], f"{name}.data")
    _check_checksum(scales, component["scales_sha256"], f"{name}.scales")
    shape = tuple(int(value) for value in component["shape"])
    if component["kind"] == "sherry_stq1":
        if int(component["group_size"]) != SHERRY_GROUP_SIZE:
            raise ValueError(f"invalid Sherry group size for packed component {name!r}")
        packed = PackedSTQ1(
            data=data,
            scales=scales,
            shape=shape,
            group_size=int(component["group_size"]),
            original_dtype=component["original_dtype"],
        )
        decoded = unpack(packed)
        repacked = pack(decoded, packed.group_size)
    else:
        expected = (
            (AUDIO_BITS, "per_group", AUDIO_GROUP_SIZE)
            if component["kind"] == "rtn_w4"
            else (EMBED_BITS, EMBED_GRANULARITY, EMBED_GROUP_SIZE)
        )
        actual = (
            int(component["bits"]),
            component["granularity"],
            int(component["group_size"]),
        )
        if actual != expected:
            raise ValueError(
                f"invalid RTN metadata for packed component {name!r}: {actual}"
            )
        packed = PackedRTN(
            data=data,
            scales=scales,
            shape=shape,
            bits=int(component["bits"]),
            granularity=component["granularity"],
            group_size=int(component["group_size"]),
            original_dtype=component["original_dtype"],
        )
        decoded = unpack_rtn(packed)
        codes = _unpack_rtn_codes(data, packed.num_weights, packed.bits)
        qmax = 2 ** (packed.bits - 1) - 1
        if bool(torch.any(codes < -qmax)) or bool(torch.any(codes > qmax)):
            raise ValueError(f"out-of-range RTN code in packed component {name!r}")
        canonical_data = _pack_rtn_codes(codes.to(torch.int8), packed.bits)
        if not torch.equal(canonical_data, data):
            raise ValueError(f"non-canonical RTN padding in packed component {name!r}")
        repacked = packed
    data_bytes = _tensor_bytes(data)
    scale_bytes = _tensor_bytes(scales)
    if (
        int(component["logical_values"]) != packed.num_weights
        or int(component["data_bytes"]) != data_bytes
        or int(component["scale_bytes"]) != scale_bytes
        or int(component["stored_bytes"]) != data_bytes + scale_bytes
    ):
        raise ValueError(f"packed size metadata mismatch for component {name!r}")
    if not torch.equal(repacked.data, data) or not torch.equal(repacked.scales, scales):
        raise ValueError(f"packed component {name!r} is not canonical")
    if not math.isclose(float(component["effective_bits"]), packed.effective_bits):
        raise ValueError(f"effective_bits mismatch for packed component {name!r}")
    return decoded.to(torch.bfloat16)


def _load_verified(export_dir: str | Path) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
    directory, manifest, artifacts = _read_bundle(export_dir)
    descriptions = manifest.get("components")
    if not isinstance(descriptions, list):
        raise ValueError("manifest components must be a list")
    names = [component.get("name") for component in descriptions]
    if len(names) != len(set(names)) or set(names) != set(artifacts):
        raise ValueError("manifest and artifact component names do not match uniquely")

    state: dict[str, torch.Tensor] = {}
    packed_verified = 0
    for component in descriptions:
        name = component["name"]
        artifact = artifacts[name]
        kind = component["kind"]
        if kind in {"sherry_stq1", "rtn_w4", "rtn_w3", "rtn_embed_w4"}:
            state[name] = _decode_packed(component, artifact)
            packed_verified += 1
        elif kind in {"bf16", "tensor"}:
            if set(artifact) != {"value"}:
                raise ValueError(f"plain component {name!r} has invalid artifact fields")
            value = artifact["value"]
            _check_checksum(value, component["sha256"], name)
            if list(value.shape) != component["shape"]:
                raise ValueError(f"shape mismatch for component {name!r}")
            if _tensor_bytes(value) != int(component["stored_bytes"]):
                raise ValueError(f"stored byte mismatch for component {name!r}")
            if kind == "bf16" and value.dtype != torch.bfloat16:
                raise ValueError(f"BF16 component {name!r} has dtype {value.dtype}")
            state[name] = value
        else:
            raise ValueError(f"unknown component kind: {kind!r}")

    aliases = manifest.get("aliases")
    if aliases != {LM_HEAD_WEIGHT: EMBED_WEIGHT}:
        raise ValueError("manifest does not contain the required tied lm_head alias")
    for alias, target in aliases.items():
        if alias in state or target not in state:
            raise ValueError(f"invalid state alias {alias!r} -> {target!r}")
        state[alias] = state[target]

    calculated_summary = _summary(descriptions)
    if calculated_summary != manifest.get("summary"):
        raise ValueError("manifest package summary is inconsistent")
    report = {
        **calculated_summary,
        "packed_tensors_verified": packed_verified,
        "state_tensors": len(state),
        "weights_file_bytes": (directory / WEIGHTS_FILE).stat().st_size,
        "manifest_file_bytes": (directory / MANIFEST_FILE).stat().st_size,
    }
    report["package_file_bytes"] = report["weights_file_bytes"] + report["manifest_file_bytes"]
    return state, report


def load_export(export_dir: str | Path, *, verify: bool = True) -> dict[str, torch.Tensor]:
    """Reload a bundle into CPU state tensors, recreating the tied head alias."""

    if not verify:
        raise ValueError("unverified loading is intentionally unsupported")
    state, _ = _load_verified(export_dir)
    return state


class CanonicalSTQLinear(nn.Linear):
    """Native packed Sherry compute without QAT or a fused bias epilogue.

    The evaluated QAT path caches the BF16 STQ weight, performs ``F.linear``
    without a bias, and then adds the BF16 bias as a separate operation.  A
    plain ``nn.Linear`` may fuse that addition into the matrix multiplication,
    which is mathematically equivalent but not bit-exact in BF16.  Deployment
    uses this lightweight native layer to preserve the evaluated operation
    order while avoiding any second fake-quantization.
    """

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        output = F.linear(input, self.weight, None)
        if self.bias is not None:
            output = output + self.bias.to(input.dtype)
        return output


def _restore_sherry_compute_semantics(
    model: nn.Module, export_dir: str | Path
) -> int:
    manifest = json.loads(
        (Path(export_dir).expanduser().resolve() / MANIFEST_FILE).read_text(
            encoding="utf-8"
        )
    )
    names = [
        component["name"]
        for component in manifest.get("components", [])
        if component.get("kind") == "sherry_stq1"
    ]
    for weight_name in names:
        if not weight_name.endswith(".weight"):
            raise ValueError(f"Sherry component is not a weight tensor: {weight_name!r}")
        module_name = weight_name.removesuffix(".weight")
        native = model.get_submodule(module_name)
        if not isinstance(native, nn.Linear):
            raise TypeError(
                f"native Sherry target {module_name!r} is not Linear: "
                f"{type(native).__name__}"
            )
        replacement = CanonicalSTQLinear(
            native.in_features,
            native.out_features,
            bias=native.bias is not None,
            device=native.weight.device,
            dtype=native.weight.dtype,
        )
        parent_name, _, leaf = module_name.rpartition(".")
        parent = model.get_submodule(parent_name) if parent_name else model
        setattr(parent, leaf, replacement)
    return len(names)


def load_export_into_native_model(
    model: nn.Module,
    export_dir: str | Path,
    *,
    return_verification_report: bool = False,
) -> nn.Module | tuple[nn.Module, dict[str, Any]]:
    """Load decoded canonical weights into native layers without re-quantizing.

    The packed payload already represents the final BF16 compute grid.  Loading
    those values into QAT wrappers would treat them as latent masters and apply
    RTN/Sherry a second time.  Deployment emulation therefore uses the native
    MOSS Linear/Embedding modules, preserves Sherry's evaluated separate-bias
    operation order, and restores the manifest's tied alias.
    """

    decoded, verification_report = _load_verified(export_dir)
    restored_sherry = _restore_sherry_compute_semantics(model, export_dir)
    expected_sherry = int(
        json.loads(
            (Path(export_dir).expanduser().resolve() / MANIFEST_FILE).read_text(
                encoding="utf-8"
            )
        )
        .get("topology", {})
        .get("counts", {})
        .get("sherry_stq1", -1)
    )
    if restored_sherry != expected_sherry:
        raise ValueError(
            f"restored Sherry compute layer count {restored_sherry} != "
            f"manifest count {expected_sherry}"
        )
    native_state = {
        name: value for name, value in decoded.items() if not name.endswith(".eps")
    }
    expected = set(model.state_dict())
    found = set(native_state)
    if found != expected:
        raise ValueError(
            f"native export state mismatch: missing={sorted(expected - found)[:20]} "
            f"extra={sorted(found - expected)[:20]}"
        )
    model.load_state_dict(native_state, strict=True)
    embedding = model.get_input_embeddings()
    head = model.get_output_embeddings()
    if embedding is None or head is None or not hasattr(head, "weight"):
        raise RuntimeError("native model does not expose tied embedding/head weights")
    if tuple(embedding.weight.shape) != tuple(head.weight.shape):
        raise RuntimeError("native embedding/head shapes cannot be tied")
    head.weight = embedding.weight
    if head.weight is not embedding.weight:
        raise RuntimeError("native deployment embedding/head tie restoration failed")
    if return_verification_report:
        return model, verification_report
    return model


def verify_export(export_dir: str | Path) -> dict[str, Any]:
    """Verify checksums, metadata, canonical repacking, aliases, and byte reports."""

    _, report = _load_verified(export_dir)
    return report


def _checkpoint_state(
    path: Path,
    *,
    legacy_w4_checkpoint_sha256: str | None = None,
) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
    document = torch.load(path, map_location="cpu", weights_only=False)
    if isinstance(document, dict) and isinstance(document.get("student"), dict):
        checkpoint_format = document.get("format")
        if checkpoint_format not in CHECKPOINT_FORMATS:
            raise ValueError(f"unsupported KD/VESPO checkpoint format: {checkpoint_format!r}")
        state = document["student"]
        topology = document.get("config")
        source = {
            "checkpoint": str(path.resolve()),
            "checkpoint_format": checkpoint_format,
            "checkpoint_version": document.get("version"),
            "checkpoint_step": int(document.get("step", 0)),
        }
    elif isinstance(document, dict) and all(torch.is_tensor(value) for value in document.values()):
        state = document
        metadata_path = path.with_name("metadata.json")
        if not metadata_path.is_file():
            raise ValueError(
                "bare VESPO student state requires sibling metadata.json to prove topology"
            )
        topology = json.loads(metadata_path.read_text(encoding="utf-8"))
        if topology.get("schema") != "vespo-snapshot-v1":
            raise ValueError("bare student state metadata is not a VESPO snapshot")
        if topology.get("dtype") != "mixed_fp32_quant_bf16_residual":
            raise ValueError(
                f"legacy/rounded VESPO snapshot dtype is not allowed: "
                f"{topology.get('dtype')!r}"
            )
        source = {
            "checkpoint": str(path.resolve()),
            "checkpoint_format": "vespo-snapshot-v1",
            "checkpoint_version": topology.get("version")
            or (path.parent.name if path.name == "student.pt" else path.stem),
            "checkpoint_step": int(topology.get("step", 0)),
        }
    else:
        raise ValueError(f"unsupported KD/VESPO checkpoint format: {path}")
    source["checkpoint_sha256"] = _file_sha256(path)
    source["checkpoint_bytes"] = path.stat().st_size
    if not all(isinstance(key, str) and torch.is_tensor(value) for key, value in state.items()):
        raise ValueError("checkpoint student state must map string keys to tensors")
    if not isinstance(topology, dict):
        raise ValueError("checkpoint does not describe its quantized topology")
    found_topology = (
        topology.get("e2e"),
        topology.get("embed_bits"),
        topology.get("embed_group_size"),
    )
    expected_topology = (True, EMBED_BITS, EMBED_CHECKPOINT_GROUP_SIZE)
    legacy_w4 = (
        PROFILE == "e2e_w4"
        and found_topology == (True, None, None)
        and legacy_w4_checkpoint_sha256 is not None
        and legacy_w4_checkpoint_sha256 == source["checkpoint_sha256"]
    )
    if found_topology != expected_topology and not legacy_w4:
        raise ValueError(
            f"checkpoint topology {found_topology} does not match final {expected_topology}"
        )
    source["checkpoint_topology"] = {
        "e2e": True,
        "embed_bits": EMBED_BITS,
        "embed_group_size": EMBED_CHECKPOINT_GROUP_SIZE,
    }
    source["legacy_w4_topology_assertion"] = legacy_w4
    return state, source


def rebuild_cpu_model(
    model_dir: str | Path,
    checkpoint: str | Path,
    *,
    moss_repo: str | Path | None = None,
    legacy_w4_checkpoint_sha256: str | None = None,
) -> tuple[nn.Module, dict[str, Any]]:
    """Rebuild the exact final topology on CPU and apply a KD/VESPO state."""

    model_path = Path(model_dir).expanduser().resolve()
    checkpoint_path = Path(checkpoint).expanduser().resolve()
    repository_path = (
        Path(moss_repo).expanduser().resolve()
        if moss_repo is not None
        else model_path.parents[1] / "MOSS-Transcribe-Diarize"
    )
    if repository_path.is_dir():
        repository = str(repository_path)
        if repository not in sys.path:
            sys.path.insert(0, repository)
    from transformers import AutoModelForCausalLM

    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        trust_remote_code=True,
        dtype=torch.bfloat16,
    )
    sherry = wrap_sherry(model, group_size=SHERRY_GROUP_SIZE)
    audio = wrap_rtn(
        model,
        [AUDIO_LINEAR_PATTERN],
        bits=AUDIO_BITS,
        granularity="per_group",
        group_size=AUDIO_GROUP_SIZE,
    )
    embedding = wrap_rtn(
        model,
        ["model.language_model.embed_tokens"],
        bits=EMBED_BITS,
        granularity=EMBED_GRANULARITY,
        group_size=EMBED_GROUP_SIZE,
    )
    if len(sherry) != EXPECTED_SHERRY_LAYERS:
        raise RuntimeError(f"expected {EXPECTED_SHERRY_LAYERS} Sherry layers, got {len(sherry)}")
    if len(audio) != EXPECTED_AUDIO_LAYERS:
        raise RuntimeError(f"expected {EXPECTED_AUDIO_LAYERS} audio W4 layers, got {len(audio)}")
    if embedding != ["model.language_model.embed_tokens"]:
        raise RuntimeError(f"expected the tied quantized embedding, got {embedding}")
    tie_rtn_lm_head(model)

    state, source = _checkpoint_state(
        checkpoint_path,
        legacy_w4_checkpoint_sha256=legacy_w4_checkpoint_sha256,
    )
    model_state = model.state_dict()
    unexpected = sorted(set(state) - set(model_state))
    if unexpected:
        raise RuntimeError(f"unexpected checkpoint keys: {unexpected[:20]}")
    bad_shapes = [
        (name, tuple(value.shape), tuple(model_state[name].shape))
        for name, value in state.items()
        if tuple(value.shape) != tuple(model_state[name].shape)
    ]
    if bad_shapes:
        raise RuntimeError(f"checkpoint tensor shape mismatch: {bad_shapes[:10]}")
    required_quantized = {
        *(_module_weight_name(name) for name in sherry),
        *(_module_weight_name(name) for name in audio),
        EMBED_WEIGHT,
    }
    missing_quantized = sorted(required_quantized - set(state))
    if missing_quantized:
        raise RuntimeError(
            f"checkpoint is missing final quantized weights: {missing_quantized[:20]}"
        )
    if EMBED_WEIGHT in state and LM_HEAD_WEIGHT in state:
        if not torch.equal(state[EMBED_WEIGHT], state[LM_HEAD_WEIGHT]):
            raise RuntimeError("checkpoint embedding and lm_head tensors disagree")
    audit_quantized_student_state(
        state,
        e2e=True,
        required_quantized_keys={*required_quantized, LM_HEAD_WEIGHT},
    )
    incompatible = model.load_state_dict(state, strict=False)
    if incompatible.unexpected_keys:
        raise RuntimeError(f"unexpected checkpoint keys: {incompatible.unexpected_keys[:20]}")
    for module in model.modules():
        if isinstance(module, SherryLinear):
            module.set_eps(0.0)
    model.cpu().eval()
    tied_embedding, tied_head = tie_rtn_lm_head(model)
    if tied_head.weight is not tied_embedding.weight:
        raise RuntimeError("embedding/lm_head tie was lost during CPU rebuild")
    source["model_dir"] = str(model_path)
    source["missing_checkpoint_keys"] = len(incompatible.missing_keys)
    return model, source


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--model-dir", type=Path)
    parser.add_argument("--checkpoint", type=Path)
    parser.add_argument("--moss-repo", type=Path)
    parser.add_argument(
        "--legacy-w4-checkpoint-sha256",
        help=(
            "explicit SHA256 assertion required only for old W4 checkpoints whose "
            "config predates embed_bits metadata"
        ),
    )
    parser.add_argument(
        "--profile",
        choices=("final_w3", "e2e_w4"),
        default="final_w3",
        help="explicit deployment embedding topology",
    )
    parser.add_argument(
        "--canonical-device",
        choices=("cpu", "cuda"),
        default="cpu",
        help="device whose fake-quant grid must match deployment evaluation",
    )
    parser.add_argument(
        "--verify-only",
        action="store_true",
        help="verify and reload an existing output bundle without loading MOSS",
    )
    return parser


def main(argv: list[str] | None = None) -> None:
    args = _parser().parse_args(argv)
    configure_profile(args.profile)
    if args.verify_only:
        report = verify_export(args.output_dir)
    else:
        if args.model_dir is None or args.checkpoint is None:
            raise SystemExit("--model-dir and --checkpoint are required unless --verify-only is set")
        model, source = rebuild_cpu_model(
            args.model_dir,
            args.checkpoint,
            moss_repo=args.moss_repo,
            legacy_w4_checkpoint_sha256=args.legacy_w4_checkpoint_sha256,
        )
        manifest = export_model(
            model,
            args.output_dir,
            source=source,
            canonical_device=args.canonical_device,
        )
        report = {"summary": manifest["summary"], "verification": verify_export(args.output_dir)}
    print(json.dumps(report, ensure_ascii=False, sort_keys=True))


if __name__ == "__main__":
    main()