#!/usr/bin/env python3
"""Convert DinoVision's runtime decoder.bin into named SafeTensors."""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path

import numpy as np
from safetensors.numpy import load_file, save_file


STAGES = (256, 128, 64, 32)
BLEND_STAGES = 2


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


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--training", required=True, type=Path)
    parser.add_argument("--input", type=Path)
    parser.add_argument("--output", required=True, type=Path)
    args = parser.parse_args()

    training = json.loads(args.training.read_text(encoding="utf-8"))
    source = args.input or args.training.parent / "decoder.bin"
    source_hash = sha256(source)
    if source_hash != training["decoder_sha256"]:
        raise SystemExit(
            f"{source} has SHA-256 {source_hash}; training record expects "
            f"{training['decoder_sha256']}"
        )

    flat = np.fromfile(source, dtype="<f4")
    offset = 0
    tensors: dict[str, np.ndarray] = {}

    def take(name: str, shape: tuple[int, ...]) -> None:
        nonlocal offset
        count = int(np.prod(shape))
        end = offset + count
        if end > len(flat):
            raise SystemExit(f"{source} ends inside tensor {name}")
        tensors[name] = flat[offset:end].reshape(shape).copy()
        offset = end

    def block(name: str, in_channels: int, out_channels: int) -> None:
        take(f"{name}.weight", (out_channels, in_channels, 3, 3))
        take(f"{name}.bias", (out_channels,))
        take(f"{name}.norm.weight", (out_channels,))
        take(f"{name}.norm.bias", (out_channels,))

    in_channels = 384
    for index, out_channels in enumerate(STAGES):
        block(f"dec.{index}", in_channels, out_channels)
        if index < BLEND_STAGES:
            block(f"dec.{index}b", out_channels, out_channels)
        in_channels = out_channels
    take("dec.out.weight", (3, in_channels, 3, 3))
    take("dec.out.bias", (3,))
    if offset != len(flat):
        raise SystemExit(f"{source} has {len(flat) - offset} trailing f32 values")

    metadata = {
        "format": "dinovision-decoder-v1",
        "source_decoder_sha256": source_hash,
        "model_sha256": training["model_sha256"],
        "dataset_manifest_sha256": training["dataset_manifest_sha256"],
        "encoder_layers": str(training["encoder_layers"]),
        "image_size": str(training["image_size"]),
        "seed": str(training["seed"]),
        "objective": training["objective"],
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    save_file(tensors, args.output, metadata=metadata)
    loaded = load_file(args.output)
    for name, expected in tensors.items():
        if name not in loaded or not np.array_equal(expected, loaded[name]):
            raise SystemExit(f"SafeTensors roundtrip changed tensor {name}")
    print(
        f"wrote {args.output}: {len(tensors)} named tensors, "
        f"SHA-256 {sha256(args.output)}"
    )


if __name__ == "__main__":
    main()