import argparse
import hashlib
import sys

import torch

from modeling import load_model

VERIFICATION = None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="model.safetensors")
    ap.add_argument("--bpe", default="bpe256.model")
    ap.add_argument("--capture", action="store_true")
    args = ap.parse_args()

    model, bpe, _ = load_model(args.ckpt, bpe=args.bpe)

    g = torch.Generator().manual_seed(0)
    wav = torch.randn(1, 16000, generator=g) * 0.01

    with torch.no_grad():
        offline = model(wav)
    assert offline.dim() == 3 and offline.size(-1) == bpe.vocab, (
        f"Bad output shape {tuple(offline.shape)} (expected vocab {bpe.vocab})"
    )
    print(f"Output shape ok {tuple(offline.shape)}")

    with torch.no_grad():
        chunks = list(model.stream_logits(wav, frame_chunk=100, reset_frames=0))
    streamed = torch.cat(chunks, dim=1)
    n = min(streamed.size(1), offline.size(1))
    diff = (streamed[:, :n] - offline[:, :n]).abs().max().item()

    print(f"Streaming vs offline max abs diff {diff:.2e}")
    assert diff < 1e-3, "Streaming output diverged from the offline forward"

    h = hashlib.sha256(
        torch.round(offline.flatten() * 1e3).to(torch.int64).numpy().tobytes()
    ).hexdigest()
    if args.capture:
        print(f'VERIFICATION = "{h}"')
        return

    if VERIFICATION is None:
        print("[WARNING] VERIFICATION is a placeholder; value check skipped.")
        return

    if h != VERIFICATION:
        print(f"[FAILURE] verification mismatch: {h} != {VERIFICATION}")
        sys.exit(1)

    print("Verification value ok")


if __name__ == "__main__":
    main()