"""Threshold sweep — varying audio length around the 40 s Encoder-Window.

Compares three transcription paths on artificially trimmed audios:
  - path1_default: model.transcribe([wav]) — what main.py uses for ≤ 25 s
  - path1_bs1:     model.transcribe([wav], batch_size=1) — NeMo ≥ 2.5 auto-chunking
  - path2:         FrameBatchMultiTaskAED with chunk_len=40 — what main.py
                   uses for > 25 s

Audios are 20, 30, 38 and 45 s slices of the 80.9 s source. The 25 s
threshold in main.py sits in the middle, so this sweep tells us:
  - Does Path 1 truncate already at 30/38 s, or only > 40 s?
  - Does the threshold need to move?
"""
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"
OUT_DIR = "/sweep"
RESULTS_PATH = f"{OUT_DIR}/threshold_results.jsonl"
LOG_PATH = f"{OUT_DIR}/threshold_log.txt"


def log(msg: str) -> None:
    line = f"{time.strftime('%H:%M:%S')}  {msg}"
    print(line, flush=True)
    with open(LOG_PATH, "a") as f:
        f.write(line + "\n")


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


def vram_mb() -> float:
    return torch.cuda.max_memory_allocated() / (1024 * 1024) if torch.cuda.is_available() else 0


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


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


def run_path2(model, model_cfg, wav, dur, chunk_len=40.0):
    manifest = write_manifest(wav, dur)
    feature_stride = model_cfg.preprocessor["window_stride"]
    model_stride_secs = feature_stride * 8
    frame_asr = FrameBatchMultiTaskAED(
        asr_model=model, frame_len=chunk_len,
        total_buffer=chunk_len, batch_size=8,
    )
    torch.cuda.reset_peak_memory_stats()
    t0 = time.monotonic()
    with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16):
        with torch.no_grad():
            hyps = get_buffered_pred_feat_multitaskAED(
                frame_asr, model_cfg.preprocessor, model_stride_secs,
                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, vram_mb()


def main():
    open(LOG_PATH, "w").close()
    open(RESULTS_PATH, "w").close()

    log("loading model bf16")
    model = ASRModel.from_pretrained("nvidia/canary-1b-v2")
    model = model.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)
    log(f"model loaded, vram={vram_mb():.0f}MB")

    audios = sorted(os.listdir(AUDIO_DIR))
    wavs = {}
    for name in audios:
        if not name.endswith(".m4a"):
            continue
        src = os.path.join(AUDIO_DIR, name)
        wav = transcode(src)
        dur = wav_duration(wav)
        wavs[name] = (wav, dur)
        log(f"  {name} → {dur:.1f}s")

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

    for name, (wav, dur) in wavs.items():
        # path 1 default (no batch_size override → triggers no auto-chunking
        # because input is a list, not a single string)
        log(f"path1_default on {name}")
        try:
            text, t, vram = run_path1(model, wav, batch_size=4)
            record({"audio": name, "duration_audio": dur, "path": "path1_default",
                    "batch_size": 4, "duration_infer": round(t, 2),
                    "vram_peak_mb": round(vram, 0), "chars": len(text), "text": text})
        except Exception as e:
            log(f"  FAIL: {e}")
            record({"audio": name, "path": "path1_default", "error": str(e)})

        # path 1 with batch_size=1 (auto-chunking trigger per NeMo docs)
        log(f"path1_bs1 on {name}")
        try:
            text, t, vram = run_path1(model, wav, batch_size=1)
            record({"audio": name, "duration_audio": dur, "path": "path1_bs1",
                    "batch_size": 1, "duration_infer": round(t, 2),
                    "vram_peak_mb": round(vram, 0), "chars": len(text), "text": text})
        except Exception as e:
            log(f"  FAIL: {e}")
            record({"audio": name, "path": "path1_bs1", "error": str(e)})

        # path 2 (always chunked, chunk_len=40)
        log(f"path2 on {name}")
        try:
            text, t, vram = run_path2(model, cfg, wav, dur, chunk_len=40.0)
            record({"audio": name, "duration_audio": dur, "path": "path2",
                    "chunk_len": 40.0, "batch_size": 8,
                    "duration_infer": round(t, 2),
                    "vram_peak_mb": round(vram, 0), "chars": len(text), "text": text})
        except Exception as e:
            log(f"  FAIL: {e}")
            record({"audio": name, "path": "path2", "error": str(e)})

    log("done")


if __name__ == "__main__":
    main()