#!/usr/bin/env python3
"""
Long-form transcription for the Hebrew GigaAM CTC model, without pyannote.

GigaAM ships `transcribe_longform`, but it imports `vad_utils`, which needs the
`longform` extra -- and that extra pins torch==2.10 / torchcodec==0.10. Installing
it would downgrade torch underneath a running training job, and it also drags in
a gated pyannote segmentation model. Neither is acceptable here.

So: energy-based segmentation, numpy and torch only, nothing new installed.
This is not a learned VAD and does not try to be. Its only job is to find
*silences deep enough to cut at*, so no word is severed and no chunk exceeds
what the model was trained on (<=20 s). For that, energy is enough.

  python longform.py --checkpoint ... --manifest ... [--device cuda]
"""

import argparse
import json
import os
import sys
from pathlib import Path
from typing import List, Sequence, Tuple

import numpy as np
import soundfile as sf
import torch

import gigaam  # noqa: E402
from gigaam.utils import normalize_raw_text  # noqa: E402

SAMPLE_RATE = 16000


def frame_energy_db(x: np.ndarray, win: int, hop: int) -> np.ndarray:
    n = 1 + max(0, (len(x) - win) // hop)
    if n <= 0:
        return np.zeros(0, dtype="float32")
    idx = np.arange(win)[None, :] + hop * np.arange(n)[:, None]
    frames = x[idx]
    rms = np.sqrt(np.maximum((frames ** 2).mean(axis=1), 1e-12))
    return 20.0 * np.log10(rms)


def speech_mask(
    x: np.ndarray,
    hop_s: float = 0.010,
    win_s: float = 0.030,
    rel_db: float = 35.0,
    min_speech_s: float = 0.20,
    min_silence_s: float = 0.30,
) -> np.ndarray:
    """Per-frame speech/silence, thresholded relative to the clip's own loud parts.

    A fixed dBFS threshold does not survive contact with real recordings -- a
    quiet WhatsApp note and a loud podcast differ by tens of dB. Thresholding
    relative to this clip's 95th-percentile energy adapts to both.
    """
    win, hop = int(win_s * SAMPLE_RATE), int(hop_s * SAMPLE_RATE)
    e = frame_energy_db(x, win, hop)
    if not len(e):
        return np.zeros(0, dtype=bool)
    loud = np.percentile(e, 95)
    mask = e > (loud - rel_db)

    # Bridge short gaps, then drop short blips. Order matters: closing before
    # opening keeps a brief inter-word pause from fragmenting one word run.
    def runs(m: np.ndarray, value: bool):
        out, start = [], None
        for i, v in enumerate(m):
            if v == value and start is None:
                start = i
            elif v != value and start is not None:
                out.append((start, i)); start = None
        if start is not None:
            out.append((start, len(m)))
        return out

    for s, e_ in runs(mask, False):
        if (e_ - s) * hop_s < min_silence_s:
            mask[s:e_] = True
    for s, e_ in runs(mask, True):
        if (e_ - s) * hop_s < min_speech_s:
            mask[s:e_] = False
    return mask


def segment_audio(
    x: np.ndarray,
    max_len_s: float = 18.0,
    min_len_s: float = 1.0,
    pad_s: float = 0.15,
    hop_s: float = 0.010,
) -> List[Tuple[float, float]]:
    """Return [(start_s, end_s)] chunks, each <= max_len_s, cut inside silence."""
    total_s = len(x) / SAMPLE_RATE
    if total_s <= max_len_s:
        return [(0.0, total_s)]

    mask = speech_mask(x, hop_s=hop_s)
    if not mask.any():
        return [(0.0, total_s)]

    # Speech regions, then split any that is still too long at its quietest point.
    regions, start = [], None
    for i, v in enumerate(mask):
        if v and start is None:
            start = i
        elif not v and start is not None:
            regions.append((start * hop_s, i * hop_s)); start = None
    if start is not None:
        regions.append((start * hop_s, len(mask) * hop_s))

    win = int(0.030 * SAMPLE_RATE)
    hop = int(hop_s * SAMPLE_RATE)
    e = frame_energy_db(x, win, hop)

    out: List[Tuple[float, float]] = []
    queue = list(regions)
    while queue:
        s, t = queue.pop(0)
        if t - s <= max_len_s:
            if t - s >= min_len_s or not out:
                out.append((s, t))
            else:
                # Too short to stand alone: glue onto the previous chunk.
                ps, pt = out[-1]
                if s - pt < 2.0 and t - ps <= max_len_s:
                    out[-1] = (ps, t)
                else:
                    out.append((s, t))
            continue
        # Cut at the quietest frame in the middle third, so we never sever a word
        # at an arbitrary fixed offset.
        a, b = int((s + (t - s) / 3) / hop_s), int((s + 2 * (t - s) / 3) / hop_s)
        a, b = max(a, 0), min(b, len(e))
        cut = (a + int(np.argmin(e[a:b]))) * hop_s if b > a else (s + t) / 2
        queue.insert(0, (cut, t))
        queue.insert(0, (s, cut))

    # Pad for a little context, clamped to the clip.
    return [(max(0.0, s - pad_s), min(total_s, t + pad_s)) for s, t in out]


@torch.inference_mode()
def transcribe_long(model, path: str, max_len_s: float = 18.0, batch_size: int = 8) -> str:
    x, sr = sf.read(path, dtype="float32")
    if x.ndim > 1:
        x = x.mean(axis=1)
    if sr != SAMPLE_RATE:
        import torchaudio.functional as taF
        x = taF.resample(torch.from_numpy(x), sr, SAMPLE_RATE).numpy()

    # Telephony scoring (GIGAAM_TELEPHONY_P=1): degrade BEFORE segmentation, so
    # the energy VAD sees phone-band audio, which is what it faces in production.
    _tel = float(os.getenv("GIGAAM_TELEPHONY_P", "0"))
    if _tel > 0:
        from gigaam.utils import telephony_degrade
        x = telephony_degrade(torch.from_numpy(np.ascontiguousarray(x)).float()).numpy()

    spans = segment_audio(x, max_len_s=max_len_s)
    chunks = [torch.from_numpy(x[int(s * SAMPLE_RATE): int(t * SAMPLE_RATE)]) for s, t in spans]
    chunks = [c for c in chunks if len(c) >= int(0.1 * SAMPLE_RATE)]
    if not chunks:
        return ""

    texts: List[str] = []
    dev = model._device
    for i in range(0, len(chunks), batch_size):
        batch = chunks[i: i + batch_size]
        lens = torch.tensor([len(c) for c in batch], dtype=torch.long)
        pad = torch.zeros(len(batch), int(lens.max()), dtype=torch.float32)
        for j, c in enumerate(batch):
            pad[j, : len(c)] = c
        enc, enc_len = model.forward(pad.to(dev).to(model._dtype), lens.to(dev))
        for text, _ in model._decode(enc, enc_len, lens.to(dev), False):
            if text.strip():
                texts.append(text.strip())
    return " ".join(texts)


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--checkpoint", required=True)
    ap.add_argument("--manifest", required=True)
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--max-len", type=float, default=18.0)
    ap.add_argument("--batch-size", type=int, default=8)
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--out", default=None)
    args = ap.parse_args()

    import editdistance
    import time

    base = Path(args.manifest).parent
    rows = []
    with open(args.manifest, encoding="utf-8") as f:
        f.readline()
        for line in f:
            p = line.rstrip("\n").split("\t")
            if len(p) >= 3:
                rows.append((base / p[0], float(p[1]), p[2]))
    if args.limit:
        rows = rows[: args.limit]

    model = gigaam.load_model(args.checkpoint, device=args.device)
    we = wt = ce = ct = 0
    results = []
    t0 = time.time()
    for path, dur, ref in rows:
        hyp = transcribe_long(model, str(path), args.max_len, args.batch_size)
        r, h = normalize_raw_text(ref), normalize_raw_text(hyp)
        we += editdistance.eval(r.split(), h.split()); wt += len(r.split())
        ce += editdistance.eval(r, h); ct += len(r)
        results.append({"path": str(path), "duration": dur, "ref": r, "hyp": h})
    wall = time.time() - t0
    audio = sum(r[1] for r in rows)

    print(f"clips : {len(rows)}  audio {audio/3600:.2f} h")
    print(f"WER   : {100*we/max(wt,1):.2f}%  ({we}/{wt})")
    print(f"CER   : {100*ce/max(ct,1):.2f}%  ({ce}/{ct})")
    print(f"wall  : {wall:.1f}s  ({audio/wall:.1f}x realtime)")
    if args.out:
        with open(args.out, "w", encoding="utf-8") as f:
            for r in results:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")


if __name__ == "__main__":
    main()