#!/usr/bin/env python3
"""Convert the frozen VibeThinker J Lens checkpoint to Safetensors.

The input path is intentionally required at runtime and is never copied into
the safetensors header or any generated public metadata. The conversion is
lossless: every FP16 bit pattern is compared after a safetensors round trip.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
from collections.abc import Mapping
from pathlib import Path
from typing import Any

import torch
from safetensors import safe_open

SOURCE_SHA256 = "f36a99447623e0d777c70951a9148a7a52e42e0df82942e22ce6326f63d8d664"
MODEL_ID = "WeiboAI/VibeThinker-3B"
MODEL_REVISION = "77bd2cced09193c8b9a59a32bd8577bbd1f3e01c"
SOURCE_LAYERS = tuple(range(0, 36, 2))
TARGET_LAYER = 35
D_MODEL = 2048
N_PROMPTS = 1000
EXPECTED_TOP_LEVEL_KEYS = {"J", "n_prompts", "source_layers", "d_model"}
FORBIDDEN_PUBLIC_FRAGMENTS = (
    os.sep.join(("", "Users", "")),
    os.sep.join(("", "Volumes", "")),
    os.sep.join(("", "workspace")),
)


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


def tensor_storage_bytes(tensor: torch.Tensor) -> bytes:
    """Return C-contiguous little-endian bytes, including FP16 bit patterns."""

    array = tensor.detach().cpu().contiguous().view(torch.int16).numpy()
    return array.astype("<i2", copy=False).tobytes(order="C")


def tensor_sha256(tensor: torch.Tensor) -> str:
    return hashlib.sha256(tensor_storage_bytes(tensor)).hexdigest()


def write_json(path: Path, value: Any) -> None:
    path.write_text(
        json.dumps(value, indent=2, sort_keys=True, ensure_ascii=True) + "\n",
        encoding="utf-8",
    )


def require(condition: bool, message: str) -> None:
    if not condition:
        raise ValueError(message)


def validate_source(checkpoint: Any) -> dict[int, torch.Tensor]:
    require(isinstance(checkpoint, Mapping), "checkpoint must be a mapping")
    require(
        set(checkpoint) == EXPECTED_TOP_LEVEL_KEYS,
        f"unexpected checkpoint keys: {sorted(checkpoint)}",
    )
    require(checkpoint["n_prompts"] == N_PROMPTS, "unexpected n_prompts")
    require(checkpoint["d_model"] == D_MODEL, "unexpected d_model")
    require(
        tuple(checkpoint["source_layers"]) == SOURCE_LAYERS,
        "unexpected source_layers",
    )

    matrices = checkpoint["J"]
    require(isinstance(matrices, Mapping), "J must be a layer-to-tensor mapping")
    require(set(matrices) == set(SOURCE_LAYERS), "unexpected J layer keys")

    validated: dict[int, torch.Tensor] = {}
    for layer in SOURCE_LAYERS:
        tensor = matrices[layer]
        require(isinstance(tensor, torch.Tensor), f"J[{layer}] is not a tensor")
        require(tensor.device.type == "cpu", f"J[{layer}] is not on CPU")
        require(tensor.dtype == torch.float16, f"J[{layer}] is not FP16")
        require(tuple(tensor.shape) == (D_MODEL, D_MODEL), f"J[{layer}] shape mismatch")
        require(tensor.is_contiguous(), f"J[{layer}] is not contiguous")
        require(bool(torch.isfinite(tensor).all()), f"J[{layer}] contains non-finite values")
        validated[layer] = tensor
    return validated


def public_header() -> dict[str, str]:
    return {
        "artifact_kind": "jacobian_lens",
        "d_model": str(D_MODEL),
        "format": "pt",
        "model_id": MODEL_ID,
        "model_revision": MODEL_REVISION,
        "n_prompts": str(N_PROMPTS),
        "schema_version": "1",
        "source_checkpoint_sha256": SOURCE_SHA256,
        "source_layers": json.dumps(SOURCE_LAYERS, separators=(",", ":")),
        "target_layer": str(TARGET_LAYER),
        "tensor_dtype": "float16",
        "tensor_key_pattern": "J.{source_layer}",
    }


def assert_public_header(metadata: Mapping[str, str]) -> None:
    encoded = json.dumps(dict(metadata), sort_keys=True)
    for fragment in FORBIDDEN_PUBLIC_FRAGMENTS:
        require(fragment not in encoded, f"private fragment found in safetensors header: {fragment}")


def save_deterministic_safetensors(
    tensors: Mapping[str, torch.Tensor],
    path: Path,
    metadata: Mapping[str, str],
) -> None:
    """Write the documented safetensors format with canonical key ordering.

    The upstream writer preserves tensor data exactly, but its Rust metadata
    map can serialize keys in a process-random order. Canonical JSON ordering
    makes the complete artifact reproducible byte for byte across runs.
    """

    offset = 0
    header: dict[str, Any] = {
        "__metadata__": {key: metadata[key] for key in sorted(metadata)}
    }
    for key in sorted(tensors):
        tensor = tensors[key]
        require(tensor.dtype == torch.float16, f"{key} is not FP16")
        nbytes = tensor.numel() * tensor.element_size()
        header[key] = {
            "dtype": "F16",
            "shape": list(tensor.shape),
            "data_offsets": [offset, offset + nbytes],
        }
        offset += nbytes

    encoded_header = json.dumps(
        header,
        ensure_ascii=False,
        separators=(",", ":"),
    ).encode("utf-8")
    padding = (-len(encoded_header)) % 8
    encoded_header += b" " * padding

    with path.open("wb") as handle:
        handle.write(len(encoded_header).to_bytes(8, byteorder="little", signed=False))
        handle.write(encoded_header)
        for key in sorted(tensors):
            handle.write(tensor_storage_bytes(tensors[key]))


def convert(source: Path, output_dir: Path, overwrite: bool) -> dict[str, Any]:
    source_digest = sha256_file(source)
    require(source_digest == SOURCE_SHA256, "source checkpoint SHA-256 mismatch")

    checkpoint = torch.load(source, map_location="cpu", weights_only=True)
    matrices = validate_source(checkpoint)
    tensors = {f"J.{layer}": matrices[layer] for layer in SOURCE_LAYERS}

    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / "model.safetensors"
    if output_path.exists() and not overwrite:
        raise FileExistsError(f"refusing to overwrite {output_path.name}; pass --overwrite")

    metadata = public_header()
    assert_public_header(metadata)
    temporary_path = output_dir / ".model.safetensors.tmp"
    save_deterministic_safetensors(tensors, temporary_path, metadata)
    os.replace(temporary_path, output_path)

    manifest_tensors: dict[str, Any] = {}
    exact_matches = 0
    with safe_open(output_path, framework="pt", device="cpu") as artifact:
        stored_metadata = artifact.metadata() or {}
        require(stored_metadata == metadata, "safetensors metadata changed during serialization")
        assert_public_header(stored_metadata)
        require(set(artifact.keys()) == set(tensors), "safetensors key set mismatch")

        for layer in SOURCE_LAYERS:
            key = f"J.{layer}"
            source_tensor = matrices[layer]
            output_tensor = artifact.get_tensor(key)
            require(output_tensor.dtype == source_tensor.dtype, f"{key} dtype mismatch")
            require(tuple(output_tensor.shape) == tuple(source_tensor.shape), f"{key} shape mismatch")
            require(output_tensor.is_contiguous(), f"{key} is not contiguous")
            require(torch.equal(output_tensor, source_tensor), f"{key} value mismatch")
            require(
                torch.equal(output_tensor.view(torch.int16), source_tensor.view(torch.int16)),
                f"{key} FP16 bit-pattern mismatch",
            )
            source_tensor_digest = tensor_sha256(source_tensor)
            output_tensor_digest = tensor_sha256(output_tensor)
            require(source_tensor_digest == output_tensor_digest, f"{key} byte hash mismatch")
            exact_matches += 1
            manifest_tensors[key] = {
                "dtype": "float16",
                "nbytes": source_tensor.numel() * source_tensor.element_size(),
                "numel": source_tensor.numel(),
                "sha256_c_contiguous_little_endian_bytes": source_tensor_digest,
                "shape": list(source_tensor.shape),
                "source_layer": layer,
            }

    output_digest = sha256_file(output_path)
    output_size = output_path.stat().st_size
    tensor_manifest = {
        "artifact": "model.safetensors",
        "artifact_sha256": output_digest,
        "artifact_size_bytes": output_size,
        "schema_version": 1,
        "tensor_count": len(manifest_tensors),
        "tensor_storage_bytes": sum(item["nbytes"] for item in manifest_tensors.values()),
        "tensors": manifest_tensors,
    }
    validation = {
        "artifact": "model.safetensors",
        "artifact_sha256": output_digest,
        "artifact_size_bytes": output_size,
        "checks": {
            "all_source_tensors_contiguous": True,
            "all_source_tensors_finite": True,
            "all_source_tensors_fp16": True,
            "all_source_tensors_shape_2048x2048": True,
            "roundtrip_all_tensor_byte_hashes_equal": True,
            "roundtrip_all_tensor_dtypes_equal": True,
            "roundtrip_all_tensor_shapes_equal": True,
            "roundtrip_all_tensor_values_equal": True,
            "roundtrip_all_tensor_bit_patterns_equal": True,
            "roundtrip_key_set_exact": True,
            "safetensors_header_public_safe": True,
            "safetensors_header_roundtrip_exact": True,
            "source_checkpoint_sha256_exact": True,
            "source_metadata_exact": True,
            "source_top_level_key_set_exact": True,
        },
        "exact_tensor_matches": exact_matches,
        "expected_tensor_matches": len(SOURCE_LAYERS),
        "ok": exact_matches == len(SOURCE_LAYERS),
        "schema_version": 1,
        "source_checkpoint_sha256": source_digest,
        "tensor_values_changed": 0,
    }
    write_json(output_dir / "tensor_manifest.json", tensor_manifest)
    write_json(output_dir / "validation.json", validation)
    (output_dir / "SHA256SUMS").write_text(
        f"{output_digest}  model.safetensors\n",
        encoding="ascii",
    )
    return validation


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", required=True, type=Path, help="Private source .pt checkpoint")
    parser.add_argument("--output-dir", required=True, type=Path, help="Public artifact directory")
    parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args()

    result = convert(args.source, args.output_dir, args.overwrite)
    print(json.dumps(result, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()