# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Build the model bundle `anuvani_diar` loads (no pickles, the Indic model only) from the frozen release bundle.

  python -s tools/export_bundle.py --out artifacts/bundle

Layout:
  MANIFEST.json                  version, the model's clustering config and decoder (`modes.indic`), licences, the
                                 sha256 of every file, and the sha256 of the reference checkpoints they were made from
  segmentation/config.json       SegmentationNet arguments
  segmentation/model.safetensors the release's hub/pytorch_model.bin state dict, tensors unchanged
  voiceprint/model.safetensors   the release's embedding/pytorch_model.bin Lightning state dict, tensors unchanged
  decoder.json                   the release's decoder, without its internal recipe label

Every tensor is checked equal (bit for bit, same dtype and shape) to the release checkpoint it came from, and each new
network is loaded with strict=True.
"""
import argparse
import hashlib
import json
import sys
from pathlib import Path

import torch
from safetensors.torch import load_file, save_file

PROJ = Path(__file__).resolve().parents[1]
REL = PROJ.parent / "anuvani_diar_v1" / "bundle"
# the published model's version (the release's checkpoints, unchanged)
MODEL_VERSION = "1.0.0"
SEG_ARGS = dict(wavlm_src="wavlm_large_s80_md", wavlm_layer_num=25, wavlm_feat_dim=1024, attention_in=256,
                ffn_hidden=1024, num_head=4, num_layer=4, kernel_size=31, dropout=0.1, use_posi=False,
                output_activate_function=False, num_classes=16)


def sha256(p):
    h = hashlib.sha256()
    with open(p, "rb") as f:
        for b in iter(lambda: f.read(1 << 24), b""):
            h.update(b)
    return h.hexdigest()


def same(a, b):
    return a.dtype == b.dtype and a.shape == b.shape and torch.equal(a.view(torch.uint8) if a.dtype == torch.bool else a, b)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", type=Path, required=True)
    a = ap.parse_args()
    sys.path.insert(0, str(PROJ / "python"))
    from anuvani_diar.nn.seg import SegmentationNet
    from anuvani_diar.nn.voiceprint import VoiceprintNet
    rel = json.loads((REL / "MANIFEST.json").read_text())
    for f, want in rel["files"].items():
        if sha256(REL / f) != want:
            raise SystemExit(f"release bundle file {f} does not match its MANIFEST sha256")
    out = a.out
    (out / "segmentation").mkdir(parents=True, exist_ok=True)
    (out / "voiceprint").mkdir(parents=True, exist_ok=True)

    # a plain state dict (599 tensors)
    seg_state = torch.load(REL / "hub/pytorch_model.bin", map_location="cpu")
    save_file({k: v.contiguous() for k, v in seg_state.items()}, str(out / "segmentation/model.safetensors"),
              metadata={"format": "pt", "source": "reference/bundle/hub/pytorch_model.bin"})
    (out / "segmentation/config.json").write_text(json.dumps(dict(architecture="SegmentationNet", model_args=SEG_ARGS), indent=1) + "\n")
    ck = torch.load(REL / "embedding/pytorch_model.bin", map_location="cpu")
    vp_state = {k: v for k, v in ck["state_dict"].items()}
    save_file({k: v.contiguous() for k, v in vp_state.items()}, str(out / "voiceprint/model.safetensors"),
              metadata={"format": "pt", "source": "reference/bundle/embedding/pytorch_model.bin (pyannote/wespeaker-voxceleb-resnet34-LM)"})
    decoder = (REL / rel["modes"]["indic"]["decoder"]).read_text()
    label = '{"recipe": "cd1", '
    if not decoder.startswith(label):
        raise SystemExit("the release's decoder does not start with its recipe label")
    (out / "decoder.json").write_text("{" + decoder[len(label):])
    if json.loads((out / "decoder.json").read_text()) != {k: v for k, v in json.loads(decoder).items() if k != "recipe"}:
        raise SystemExit("decoder.json differs from the release's decoder beyond the recipe label")

    # every tensor round-trips bit for bit, and both networks load strictly
    for name, state in (("segmentation", seg_state), ("voiceprint", vp_state)):
        back = load_file(str(out / name / "model.safetensors"))
        assert back.keys() == state.keys(), name
        bad = [k for k in state if not same(back[k], state[k])]
        assert not bad, (name, bad[:5])
    seg = SegmentationNet.from_pretrained(out / "segmentation")
    vp = VoiceprintNet.from_pretrained(out / "voiceprint")
    n_seg = sum(p.numel() for p in seg.parameters())
    n_vp = sum(p.numel() for p in vp.parameters())

    files = {}
    for p in sorted(out.rglob("*")):
        if p.is_file() and p.name != "MANIFEST.json":
            files[str(p.relative_to(out))] = sha256(p)
    manifest = dict(
        version=MODEL_VERSION, model="Anuvani-Diarization",
        modes={"indic": dict(clustering=rel["modes"]["indic"]["clustering"], decoder="decoder.json")},
        networks=dict(segmentation=dict(dir="segmentation", parameters=n_seg, source_sha256=rel["files"]["hub/pytorch_model.bin"]),
                      voiceprint=dict(dir="voiceprint", parameters=n_vp, source_sha256=rel["files"]["embedding/pytorch_model.bin"])),
        licence=rel["licence"], files=files)
    (out / "MANIFEST.json").write_text(json.dumps(manifest, indent=1) + "\n")
    print(json.dumps(dict(out=str(out), segmentation_parameters=n_seg, voiceprint_parameters=n_vp, files=files), indent=1))


if __name__ == "__main__":
    main()