"""Export actual 16-layer Qwen3.5-9B prefix, preserving its shared norm/head."""

import argparse
import copy
from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
import re
import shutil

LAYER = re.compile(r"^model\.language_model\.layers\.(\d+)\.")
READOUT = {
    "model.language_model.embed_tokens.weight",
    "model.language_model.norm.weight",
    "lm_head.weight",
}


def file_sha(path):
    h = hashlib.sha256()
    with Path(path).open("rb") as f:
        for block in iter(lambda: f.read(8 * 1024 * 1024), b""):
            h.update(block)
    return h.hexdigest()


def source_snapshot(source):
    """Stat all root model files; ignore unrelated nested evaluation directories."""
    return {
        p.name: {
            "bytes": p.stat().st_size,
            "mtime_ns": p.stat().st_mtime_ns,
            "inode": p.stat().st_ino,
        }
        for p in sorted(Path(source).iterdir())
        if p.is_file()
    }


def assert_source_unchanged(source, before):
    if source_snapshot(source) != before:
        raise ValueError(
            "source model files changed during export; refusing verified publication"
        )


def verify_artifact(path, depth):
    """Default launcher gate: rehash every weight, not only config and index."""
    path = Path(path)
    config = json.loads((path / "config.json").read_text())
    text = config.get("text_config", {})
    if (
        config.get("model_type") != "qwen3_5"
        or text.get("hidden_size") != 4096
        or text.get("vocab_size") != 248320
        or text.get("num_hidden_layers") != depth
        or len(text.get("layer_types", [])) != depth
        or (path / "adapter_config.json").exists()
    ):
        raise ValueError("backend architecture/depth mismatch")
    if depth == 32:
        check_source(path)
        return
    if depth != 16 or (path / "EXPORT_INCOMPLETE").exists():
        raise ValueError("only complete 16/32 exports accepted")
    manifest = json.loads((path / "export-low-provenance.json").read_text())
    if (
        manifest.get("status") != "verified"
        or manifest.get("actual_depth") != 16
        or manifest.get("all_retained_tensors_bit_identical") is not True
    ):
        raise ValueError("low export provenance is not verified")
    for name, field in (
        ("config.json", "export_config_sha256"),
        ("model.safetensors.index.json", "export_index_sha256"),
    ):
        if file_sha(path / name) != manifest[field]:
            raise ValueError("low config/index SHA mismatch")
    weights = json.loads((path / "model.safetensors.index.json").read_text())[
        "weight_map"
    ]
    select_keys(weights, 16, 16)
    entries = manifest["export_weight_files"]
    expected = {entry["file"]: entry for entry in entries}
    if len(expected) != len(entries) or set(expected) != set(weights.values()):
        raise ValueError("low weight manifest/index mismatch")
    for name, entry in expected.items():
        if Path(name).name != name or not name.endswith(".safetensors"):
            raise ValueError("invalid low shard path")
        if (path / name).stat().st_size != entry["bytes"] or file_sha(
            path / name
        ) != entry["sha256"]:
            raise ValueError(f"low weight SHA mismatch: {name}")


def tensor_sha(tensor):
    import torch

    raw = memoryview(tensor.contiguous().view(torch.uint8).numpy()).cast("B")
    h = hashlib.sha256()
    for start in range(0, len(raw), 8 * 1024 * 1024):
        h.update(raw[start : start + 8 * 1024 * 1024])
    return h.hexdigest()


def prefix_config(config, depth=16):
    result = copy.deepcopy(config)
    text = result["text_config"]
    if (
        not 0 < depth < text["num_hidden_layers"]
        or len(text["layer_types"]) != text["num_hidden_layers"]
    ):
        raise ValueError("invalid depth or layer_types")
    text["num_hidden_layers"] = depth
    text["layer_types"] = text["layer_types"][:depth]
    if "mlp_only_layers" in text:
        text["mlp_only_layers"] = [i for i in text["mlp_only_layers"] if i < depth]
    return result


def select_keys(keys, depth, full_depth):
    kept, removed, seen = [], [], set()
    for key in sorted(keys):
        match = LAYER.match(key)
        if match:
            i = int(match.group(1))
            if not 0 <= i < full_depth:
                raise ValueError("layer outside full depth")
            seen.add(i)
            (kept if i < depth else removed).append(key)
        elif key in READOUT or key.startswith("model.visual."):
            kept.append(key)
        else:
            raise ValueError(f"unexpected weight namespace: {key}")
    if seen != set(range(full_depth)) or not READOUT <= set(kept):
        raise ValueError("missing decoder, embedding, final norm or lm_head")
    return kept, removed


def check_source(source):
    config = json.loads((source / "config.json").read_text())
    text = config.get("text_config", {})
    if (
        config.get("model_type") != "qwen3_5"
        or config.get("architectures") != ["Qwen3_5ForConditionalGeneration"]
        or text.get("num_hidden_layers") != 32
        or text.get("hidden_size") != 4096
        or text.get("vocab_size") != 248320
        or text.get("intermediate_size") != 12288
        or config.get("tie_word_embeddings") is not False
        or text.get("tie_word_embeddings") is not False
        or (source / "adapter_config.json").exists()
    ):
        raise ValueError("only merged, untied dense Qwen3.5-9B full32 accepted")
    depth = json.loads((source / "depth_config.json").read_text())
    if depth.get("exit_depth") != 16 or depth.get("full_depth") != 32:
        raise ValueError("source must declare trained exits 16/32")
    if text.get("layer_types") != (["linear_attention"] * 3 + ["full_attention"]) * 8:
        raise ValueError("unexpected Qwen3.5-9B hybrid schedule")
    weight_map = json.loads((source / "model.safetensors.index.json").read_text())[
        "weight_map"
    ]
    kept, removed = select_keys(weight_map, 16, 32)
    manifest = json.loads((source / "release-manifest.json").read_text())
    expected = {r["path"]: r["sha256"] for r in manifest["files"]}
    source_files = []
    for name in sorted(set(weight_map.values())):
        if Path(name).name != name or not name.endswith(".safetensors"):
            raise ValueError("weight shard must be a simple safetensors filename")
        observed = file_sha(source / name)
        if expected.get(name) != observed:
            raise ValueError(f"source weight SHA mismatch: {name}")
        source_files.append(
            {"file": name, "sha256": observed, "bytes": (source / name).stat().st_size}
        )
    for name in (
        "config.json",
        "depth_config.json",
        "model.safetensors.index.json",
        "tokenizer.json",
        "tokenizer_config.json",
        "chat_template.jinja",
    ):
        if name not in expected or file_sha(source / name) != expected[name]:
            raise ValueError(f"source metadata SHA mismatch: {name}")
    return config, weight_map, kept, removed, source_files


def copy_weights(source, output, weight_map, kept):
    from safetensors import safe_open
    from safetensors.torch import save_file

    shards = [
        n
        for n in sorted(set(weight_map.values()))
        if any(weight_map[k] == n for k in kept)
    ]
    output_map, files, manifests = {}, [], []
    total = 0
    for ordinal, name in enumerate(shards, 1):
        selected = sorted(k for k in kept if weight_map[k] == name)
        with safe_open(source / name, framework="pt", device="cpu") as src:
            if set(src.keys()) != {k for k, v in weight_map.items() if v == name}:
                raise ValueError("source shard/index mismatch")
            tensors = {k: src.get_tensor(k) for k in selected}
            total += sum(t.numel() * t.element_size() for t in tensors.values())
            out_name = f"model-low16-{ordinal:05d}-of-{len(shards):05d}.safetensors"
            save_file(tensors, output / out_name, metadata={"format": "pt"})
            with safe_open(output / out_name, framework="pt", device="cpu") as dst:
                if set(dst.keys()) != set(selected):
                    raise ValueError("exported keys mismatch")
                for key in selected:
                    original, exported = tensors[key], dst.get_tensor(key)
                    checksum = tensor_sha(original)
                    if (
                        original.dtype != exported.dtype
                        or original.shape != exported.shape
                        or checksum != tensor_sha(exported)
                    ):
                        raise ValueError(f"exported tensor changed: {key}")
                    manifests.append(
                        {
                            "name": key,
                            "shape": list(original.shape),
                            "dtype": str(original.dtype),
                            "sha256": checksum,
                        }
                    )
                    output_map[key] = out_name
            files.append(
                {
                    "file": out_name,
                    "sha256": file_sha(output / out_name),
                    "bytes": (output / out_name).stat().st_size,
                }
            )
        print(
            json.dumps(
                {"verified_export_shard": out_name, "tensor_count": len(selected)}
            ),
            flush=True,
        )
    if set(output_map) != set(kept):
        raise ValueError("export index incomplete")
    (output / "model.safetensors.index.json").write_text(
        json.dumps(
            {"metadata": {"total_size": total}, "weight_map": output_map}, indent=2
        )
        + "\n"
    )
    return files, manifests


def export(source, output):
    from safetensors import safe_open

    source, output = Path(source).resolve(), Path(output).resolve()
    if output.exists() or source in output.parents:
        raise ValueError("output must be a new directory outside source model")
    before = source_snapshot(source)
    config, weight_map, kept, removed, source_files = check_source(source)
    for name in sorted(set(weight_map.values())):
        with safe_open(source / name, framework="pt", device="cpu") as handle:
            if set(handle.keys()) != {k for k, v in weight_map.items() if v == name}:
                raise ValueError("source index/shard mismatch")
            for key, shape in {
                "lm_head.weight": (248320, 4096),
                "model.language_model.embed_tokens.weight": (248320, 4096),
                "model.language_model.norm.weight": (4096,),
            }.items():
                if (
                    weight_map[key] == name
                    and tuple(handle.get_slice(key).get_shape()) != shape
                ):
                    raise ValueError(f"unexpected shared readout shape: {key}")
    output.mkdir(parents=True, exist_ok=False)
    (output / "EXPORT_INCOMPLETE").write_text(
        "Only provenance status=verified denotes success.\n"
    )
    (output / "config.json").write_text(
        json.dumps(prefix_config(config), indent=2) + "\n"
    )
    for name in (
        "tokenizer.json",
        "tokenizer_config.json",
        "chat_template.jinja",
        "generation_config.json",
        "special_tokens_map.json",
        "added_tokens.json",
        "vocab.json",
        "merges.txt",
        "preprocessor_config.json",
        "processor_config.json",
        "video_preprocessor_config.json",
        "LICENSE",
    ):
        if (source / name).is_file():
            shutil.copy2(source / name, output / name)
    files, tensors = copy_weights(source, output, weight_map, kept)
    provenance = {
        "format": "openjev.vllm-low-prefix.v1",
        "status": "verified",
        "created_utc": datetime.now(timezone.utc).isoformat(),
        "source": str(source),
        "source_full_depth": 32,
        "actual_depth": 16,
        "dynamic_continuation": False,
        "selection": "original decoder 0..15 + shared embed/final norm/lm_head + all visual tensors",
        "source_config_sha256": file_sha(source / "config.json"),
        "source_depth_config": json.loads((source / "depth_config.json").read_text()),
        "source_release_manifest_sha256": file_sha(source / "release-manifest.json"),
        "source_index_sha256": file_sha(source / "model.safetensors.index.json"),
        "source_weight_files": source_files,
        "export_config_sha256": file_sha(output / "config.json"),
        "export_index_sha256": file_sha(output / "model.safetensors.index.json"),
        "export_weight_files": files,
        "retained_tensors": tensors,
        "removed_tensors": removed,
        "all_retained_tensors_bit_identical": True,
        "vllm_numerical_parity": "not established by CPU export; requires independent GPU comparison",
        "script_sha256": file_sha(__file__),
    }
    assert_source_unchanged(source, before)
    provenance["source_files_stat_before_after"] = before
    provenance["source_files_unchanged_during_export"] = True
    (output / "export-low-provenance.json").write_text(
        json.dumps(provenance, indent=2) + "\n"
    )
    (output / "EXPORT_INCOMPLETE").unlink()
    return provenance


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--source", type=Path, required=True)
    p.add_argument("--output", type=Path, required=True)
    args = p.parse_args()
    result = export(args.source, args.output)
    print(
        json.dumps(
            {
                "status": result["status"],
                "output": str(args.output),
                "actual_depth": 16,
                "retained_tensors": len(result["retained_tensors"]),
            }
        ),
        flush=True,
    )


if __name__ == "__main__":
    main()