"""Mini-sweep: does chunk_len == audio_len fix Path 2 for short audios?

If yes  → Path 2 is OK as long as audio fits in exactly one chunk.
         The bug is in the multi-chunk merging or pad/split logic.
If no   → Path 2 has a structural problem with single-chunk inputs
         independent of length matching. Confirms "different inference
         pipeline" theory.
"""
import copy
import json
import os
import subprocess
import tempfile
import time
import wave

import torch
from omegaconf import OmegaConf

from nemo.collections.asr.models import ASRModel
from nemo.collections.asr.parts.utils.streaming_utils import FrameBatchMultiTaskAED
from nemo.collections.asr.parts.utils.transcribe_utils import get_buffered_pred_feat_multitaskAED


AUDIO_DIR = "/sweep/threshold_audios"
RESULTS_PATH = "/sweep/chunklen_results.jsonl"


def transcode(src):
    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 wav_dur(p):
    with wave.open(p, "rb") as w:
        return w.getnframes() / float(w.getframerate())


def write_manifest(wav, dur):
    f = tempfile.NamedTemporaryFile(suffix=".json", mode="w", delete=False)
    json.dump({"audio_filepath": wav, "duration": dur, "taskname": "asr",
               "source_lang": "de", "target_lang": "de", "pnc": "yes"}, f)
    f.write("\n")
    f.close()
    return f.name


def run_path2(model, cfg, wav, dur, chunk_len):
    manifest = write_manifest(wav, dur)
    feature_stride = cfg.preprocessor["window_stride"]
    model_stride = feature_stride * 8
    fa = FrameBatchMultiTaskAED(asr_model=model, frame_len=chunk_len,
                                 total_buffer=chunk_len, batch_size=8)
    t0 = time.monotonic()
    with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16):
        with torch.no_grad():
            hyps = get_buffered_pred_feat_multitaskAED(
                fa, cfg.preprocessor, model_stride, model.device,
                manifest=manifest, filepaths=None, timestamps=False,
            )
    elapsed = time.monotonic() - t0
    os.unlink(manifest)
    return (getattr(hyps[0], "text", None) or str(hyps[0])).strip(), elapsed


def run_path1(model, wav):
    t0 = time.monotonic()
    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", batch_size=4)
    elapsed = time.monotonic() - t0
    return (getattr(hyps[0], "text", None) or str(hyps[0])).strip(), elapsed


def main():
    open(RESULTS_PATH, "w").close()
    print("loading model bf16", flush=True)
    model = ASRModel.from_pretrained("nvidia/canary-1b-v2").to("cuda").to(torch.bfloat16)
    model.eval()
    cfg = copy.deepcopy(model._cfg)
    OmegaConf.set_struct(cfg.preprocessor, False)
    cfg.preprocessor.dither = 0.0
    cfg.preprocessor.pad_to = 0
    OmegaConf.set_struct(cfg.preprocessor, True)
    print("model loaded", flush=True)

    audios = ["trim_20s.m4a", "trim_30s.m4a"]
    chunk_lens = {
        "trim_20s.m4a": [21.0, 25.0, 40.0],   # exact, slightly larger, default
        "trim_30s.m4a": [31.0, 35.0, 40.0],
    }

    def record(d):
        with open(RESULTS_PATH, "a") as f:
            f.write(json.dumps(d, ensure_ascii=False) + "\n")

    for name in audios:
        src = os.path.join(AUDIO_DIR, name)
        wav = transcode(src)
        dur = wav_dur(wav)
        print(f"\n=== {name} ({dur:.1f}s) ===", flush=True)

        # baseline path1
        text, t = run_path1(model, wav)
        print(f"path1: {len(text)} chars in {t:.1f}s", flush=True)
        record({"audio": name, "duration": dur, "path": "path1",
                "chunk_len": None, "chars": len(text),
                "duration_infer": round(t, 2), "text": text})

        for cl in chunk_lens[name]:
            text, t = run_path2(model, cfg, wav, dur, cl)
            print(f"path2 chunk_len={cl}: {len(text)} chars in {t:.1f}s", flush=True)
            record({"audio": name, "duration": dur, "path": "path2",
                    "chunk_len": cl, "chars": len(text),
                    "duration_infer": round(t, 2), "text": text})

    print("done", flush=True)


if __name__ == "__main__":
    main()