"""Canary parameter sweep — direct NeMo calls, no HTTP layer.

Runs many parameter combinations against a fixed set of audios and writes
JSONL + a human-readable summary. Designed to run inside the doctate-canary
Docker image (or its base) with the model volume mounted at /models.

Outputs:
  /sweep/results.jsonl  — one JSON object per run
  /sweep/log.txt        — chronological log
"""
import copy
import json
import os
import subprocess
import sys
import tempfile
import time
import wave
from contextlib import contextmanager

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/audios"
OUT_DIR = "/sweep"
RESULTS_PATH = f"{OUT_DIR}/results.jsonl"
LOG_PATH = f"{OUT_DIR}/log.txt"

# Subset for most axes — one short, one medium, one long (>40 s each).
SUBSET = [
    "2026-04-27T15-00-48Z.m4a",  # 48.4 s
    "2026-04-27T14-55-41Z.m4a",  # 65.9 s
    "2026-04-27T14-58-17Z.m4a",  # 80.9 s
]

# Single-audio subset for speed/VRAM-only axes.
SINGLE = ["2026-04-27T14-58-17Z.m4a"]  # 80.9 s

DTYPE_MAP = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}


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


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


@contextmanager
def vram_track():
    torch.cuda.reset_peak_memory_stats()
    try:
        yield
    finally:
        pass


def load_model(precision: str):
    log(f"loading model precision={precision}")
    t0 = time.monotonic()
    model = ASRModel.from_pretrained("nvidia/canary-1b-v2")
    model = model.to("cuda")
    if precision != "fp32":
        model = model.to(DTYPE_MAP[precision])
    model.eval()
    log(f"model loaded in {time.monotonic() - t0:.1f}s, vram={vram_mb():.0f}MB")
    return model


def prep_cfg(model):
    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)
    return cfg


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


def run_buffered(model, model_cfg, wav_path: str, duration: float,
                 chunk_len: float, total_buffer: float, batch_size: int,
                 precision: str, pnc: str = "yes", timestamps: bool = False):
    manifest = write_manifest(wav_path, duration, pnc=pnc)
    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=total_buffer,
        batch_size=batch_size,
    )
    amp_dtype = DTYPE_MAP[precision]
    use_amp = precision != "fp32"
    t0 = time.monotonic()
    with vram_track():
        with torch.amp.autocast("cuda", enabled=use_amp, dtype=amp_dtype):
            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=timestamps,
                )
    elapsed = time.monotonic() - t0
    peak_vram = vram_mb()
    os.unlink(manifest)
    text = (getattr(hyps[0], "text", None) or str(hyps[0])).strip()
    return text, elapsed, peak_vram


def run_transcribe_path1(model, wav_path: str, batch_size: int, pnc: str,
                          timestamps: bool, precision: str):
    """Path 1 — direct .transcribe() with batch_size=1 enables auto-chunking
    in NeMo >= 2.5 (per docstring of speech_to_text_aed_chunked_infer.py)."""
    amp_dtype = DTYPE_MAP[precision]
    use_amp = precision != "fp32"
    t0 = time.monotonic()
    with vram_track():
        with torch.amp.autocast("cuda", enabled=use_amp, dtype=amp_dtype):
            with torch.no_grad():
                hyps = model.transcribe(
                    [wav_path],
                    source_lang="de", target_lang="de",
                    pnc=pnc, timestamps=timestamps,
                    batch_size=batch_size,
                )
    elapsed = time.monotonic() - t0
    peak_vram = vram_mb()
    text = (getattr(hyps[0], "text", None) or str(hyps[0])).strip()
    return text, elapsed, peak_vram


def write_result(record: dict) -> None:
    with open(RESULTS_PATH, "a") as f:
        f.write(json.dumps(record, ensure_ascii=False) + "\n")


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

    log("=== preparing audios ===")
    wavs = {}
    for name in set(SUBSET + SINGLE):
        src = os.path.join(AUDIO_DIR, name)
        wav = transcode_to_wav16k(src)
        dur = wav_duration(wav)
        wavs[name] = (wav, dur)
        log(f"  {name} → {wav} ({dur:.1f}s)")

    # ===== precision axis (B.2) — separate model loads =====
    for precision in ["bf16", "fp16", "fp32"]:
        try:
            torch.cuda.empty_cache()
            model = load_model(precision)
            cfg = prep_cfg(model)
            for name in SINGLE:
                wav, dur = wavs[name]
                axis = f"B2_precision={precision}"
                log(f"running {axis} on {name}")
                try:
                    text, t, vram = run_buffered(
                        model, cfg, wav, dur,
                        chunk_len=40.0, total_buffer=40.0, batch_size=8,
                        precision=precision,
                    )
                    write_result({
                        "axis": "precision", "precision": precision,
                        "chunk_len": 40.0, "overlap": 0, "batch_size": 8,
                        "decoding": "greedy", "pnc": "yes", "timestamps": False,
                        "audio": name, "duration_audio": dur,
                        "duration_infer": round(t, 2), "vram_peak_mb": round(vram, 0),
                        "chars": len(text), "text": text,
                    })
                except Exception as e:
                    log(f"  FAIL {axis} on {name}: {e}")
                    write_result({"axis": "precision", "precision": precision,
                                  "audio": name, "error": str(e)})
            del model
            torch.cuda.empty_cache()
        except Exception as e:
            log(f"FAIL precision={precision}: {e}")

    # All remaining axes use bf16 default model.
    torch.cuda.empty_cache()
    model = load_model("bf16")
    cfg = prep_cfg(model)

    # ===== chunk_len axis (A.1) =====
    for chunk_len in [10.0, 20.0, 40.0, 60.0]:
        for name in SUBSET:
            wav, dur = wavs[name]
            axis = f"A1_chunk_len={chunk_len}"
            log(f"running {axis} on {name}")
            try:
                text, t, vram = run_buffered(
                    model, cfg, wav, dur,
                    chunk_len=chunk_len, total_buffer=chunk_len, batch_size=8,
                    precision="bf16",
                )
                write_result({
                    "axis": "chunk_len", "chunk_len": chunk_len,
                    "overlap": 0, "batch_size": 8, "precision": "bf16",
                    "decoding": "greedy", "pnc": "yes", "timestamps": False,
                    "audio": name, "duration_audio": dur,
                    "duration_infer": round(t, 2), "vram_peak_mb": round(vram, 0),
                    "chars": len(text), "text": text,
                })
            except Exception as e:
                log(f"  FAIL: {e}")
                write_result({"axis": "chunk_len", "chunk_len": chunk_len,
                              "audio": name, "error": str(e)})

    # ===== overlap axis (A.2) — chunk_len=40, total_buffer=chunk_len+overlap =====
    for overlap in [1.0, 2.0, 5.0, 10.0]:
        chunk_len = 40.0
        total_buffer = chunk_len + overlap
        for name in SUBSET:
            wav, dur = wavs[name]
            axis = f"A2_overlap={overlap}"
            log(f"running {axis} on {name}")
            try:
                text, t, vram = run_buffered(
                    model, cfg, wav, dur,
                    chunk_len=chunk_len, total_buffer=total_buffer,
                    batch_size=8, precision="bf16",
                )
                write_result({
                    "axis": "overlap", "chunk_len": chunk_len,
                    "overlap": overlap, "batch_size": 8, "precision": "bf16",
                    "decoding": "greedy", "pnc": "yes", "timestamps": False,
                    "audio": name, "duration_audio": dur,
                    "duration_infer": round(t, 2), "vram_peak_mb": round(vram, 0),
                    "chars": len(text), "text": text,
                })
            except Exception as e:
                log(f"  FAIL: {e}")
                write_result({"axis": "overlap", "overlap": overlap,
                              "audio": name, "error": str(e)})

    # ===== batch_size axis (B.1) =====
    for batch_size in [1, 2, 4, 8, 16]:
        for name in SINGLE:
            wav, dur = wavs[name]
            axis = f"B1_batch={batch_size}"
            log(f"running {axis} on {name}")
            try:
                text, t, vram = run_buffered(
                    model, cfg, wav, dur,
                    chunk_len=40.0, total_buffer=40.0, batch_size=batch_size,
                    precision="bf16",
                )
                write_result({
                    "axis": "batch_size", "batch_size": batch_size,
                    "chunk_len": 40.0, "overlap": 0, "precision": "bf16",
                    "decoding": "greedy", "pnc": "yes", "timestamps": False,
                    "audio": name, "duration_audio": dur,
                    "duration_infer": round(t, 2), "vram_peak_mb": round(vram, 0),
                    "chars": len(text), "text": text,
                })
            except Exception as e:
                log(f"  FAIL: {e}")
                write_result({"axis": "batch_size", "batch_size": batch_size,
                              "audio": name, "error": str(e)})

    # ===== pnc axis (C.2) =====
    for pnc in ["no", "yes"]:
        for name in SUBSET:
            wav, dur = wavs[name]
            axis = f"C2_pnc={pnc}"
            log(f"running {axis} on {name}")
            try:
                text, t, vram = run_buffered(
                    model, cfg, wav, dur,
                    chunk_len=40.0, total_buffer=40.0, batch_size=8,
                    precision="bf16", pnc=pnc,
                )
                write_result({
                    "axis": "pnc", "pnc": pnc,
                    "chunk_len": 40.0, "overlap": 0, "batch_size": 8,
                    "precision": "bf16", "decoding": "greedy",
                    "timestamps": False,
                    "audio": name, "duration_audio": dur,
                    "duration_infer": round(t, 2), "vram_peak_mb": round(vram, 0),
                    "chars": len(text), "text": text,
                })
            except Exception as e:
                log(f"  FAIL: {e}")
                write_result({"axis": "pnc", "pnc": pnc,
                              "audio": name, "error": str(e)})

    # ===== timestamps axis (C.3) — when on, NeMo recommends chunk_len=10 =====
    for ts_flag, chunk_len in [(False, 40.0), (True, 10.0)]:
        for name in SUBSET:
            wav, dur = wavs[name]
            axis = f"C3_timestamps={ts_flag}"
            log(f"running {axis} (chunk_len={chunk_len}) on {name}")
            try:
                text, t, vram = run_buffered(
                    model, cfg, wav, dur,
                    chunk_len=chunk_len, total_buffer=chunk_len, batch_size=8,
                    precision="bf16", timestamps=ts_flag,
                )
                write_result({
                    "axis": "timestamps", "timestamps": ts_flag,
                    "chunk_len": chunk_len, "overlap": 0, "batch_size": 8,
                    "precision": "bf16", "decoding": "greedy", "pnc": "yes",
                    "audio": name, "duration_audio": dur,
                    "duration_infer": round(t, 2), "vram_peak_mb": round(vram, 0),
                    "chars": len(text), "text": text,
                })
            except Exception as e:
                log(f"  FAIL: {e}")
                write_result({"axis": "timestamps", "timestamps": ts_flag,
                              "audio": name, "error": str(e)})

    # ===== beam search (C.1) =====
    try:
        from nemo.collections.asr.parts.submodules.multitask_decoding import MultiTaskDecodingConfig
        for strategy, beam_size in [("beam", 4), ("greedy", 1)]:
            decoding_cfg = MultiTaskDecodingConfig()
            decoding_cfg.strategy = strategy
            if hasattr(decoding_cfg, "beam"):
                decoding_cfg.beam.beam_size = beam_size
            try:
                model.change_decoding_strategy(decoding_cfg)
            except Exception as e:
                log(f"change_decoding_strategy failed for {strategy}: {e}")
                continue
            for name in SUBSET:
                wav, dur = wavs[name]
                axis = f"C1_decoding={strategy}_beam={beam_size}"
                log(f"running {axis} on {name}")
                try:
                    text, t, vram = run_buffered(
                        model, cfg, wav, dur,
                        chunk_len=40.0, total_buffer=40.0, batch_size=8,
                        precision="bf16",
                    )
                    write_result({
                        "axis": "decoding", "decoding": strategy,
                        "beam_size": beam_size,
                        "chunk_len": 40.0, "overlap": 0, "batch_size": 8,
                        "precision": "bf16", "pnc": "yes", "timestamps": False,
                        "audio": name, "duration_audio": dur,
                        "duration_infer": round(t, 2), "vram_peak_mb": round(vram, 0),
                        "chars": len(text), "text": text,
                    })
                except Exception as e:
                    log(f"  FAIL: {e}")
                    write_result({"axis": "decoding", "decoding": strategy,
                                  "audio": name, "error": str(e)})
    except Exception as e:
        log(f"beam search axis skipped: {e}")

    # ===== Path-1 (.transcribe() auto-chunking) — D.1 =====
    log("=== Path 1: model.transcribe() with auto-chunking ===")
    for name in SUBSET:
        wav, dur = wavs[name]
        axis = "D1_path1_transcribe"
        log(f"running {axis} on {name}")
        try:
            text, t, vram = run_transcribe_path1(
                model, wav, batch_size=1, pnc="yes",
                timestamps=False, precision="bf16",
            )
            write_result({
                "axis": "path1", "path": "transcribe_auto_chunking",
                "batch_size": 1, "precision": "bf16",
                "audio": name, "duration_audio": dur,
                "duration_infer": round(t, 2), "vram_peak_mb": round(vram, 0),
                "chars": len(text), "text": text,
            })
        except Exception as e:
            log(f"  FAIL: {e}")
            write_result({"axis": "path1", "audio": name, "error": str(e)})

    log("=== sweep complete ===")


if __name__ == "__main__":
    sweep()