"""Show forced-alignment timestamps from Canary's main transcribe(timestamps=True).

Demonstrates the full Hypothesis structure produced by:
    AED text + CTC submodel acoustic alignment via Viterbi.

Output:
    /sweep/canary_timestamps_demo.json — full pretty-printed dict
    stdout — short summary
"""
import json
import os
import subprocess
import tempfile

import torch

from nemo.collections.asr.models import EncDecMultiTaskModel


WAV_SRC = "/sweep/case_c414cf52_15-01-39.m4a"
OUT_PATH = "/sweep/canary_timestamps_dosing.json"


def transcode(src: str) -> str:
    dst = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    dst.close()
    subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", src,
                    "-ac", "1", "-ar", "16000",
                    "-f", "wav", dst.name], check=True)
    return dst.name


def to_jsonable(o):
    if hasattr(o, "tolist"):
        return o.tolist()
    if hasattr(o, "item"):
        return o.item()
    return str(o)


def main():
    print("loading Canary 1B v2", flush=True)
    model = EncDecMultiTaskModel.from_pretrained("nvidia/canary-1b-v2")
    model = model.to("cuda").to(torch.bfloat16)
    model.eval()

    wav = transcode(WAV_SRC)
    print(f"running model.transcribe(timestamps=True) on {WAV_SRC}", flush=True)
    with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16):
        with torch.no_grad():
            hyps = model.transcribe([wav], source_lang="de", target_lang="de",
                                     pnc="yes", timestamps=True, batch_size=4)
    h = hyps[0]

    out = {"text": getattr(h, "text", None),
           "timestamp": getattr(h, "timestamp", None)}
    with open(OUT_PATH, "w") as f:
        json.dump(out, f, default=to_jsonable, ensure_ascii=False, indent=2)

    print(f"\ntext: {h.text}", flush=True)
    ts = h.timestamp if isinstance(h.timestamp, dict) else {}
    for k, v in ts.items():
        n = len(v) if hasattr(v, "__len__") else "?"
        print(f"timestamp[{k}]: {n} entries", flush=True)
    print(f"\nfull dump -> {OUT_PATH}", flush=True)


if __name__ == "__main__":
    main()