#!/usr/bin/env python3
from __future__ import annotations

import argparse
import io
import json
import time
from pathlib import Path

import numpy as np
import pyarrow.parquet as pq
import soundfile as sf
import torch
from transformers import WhisperForConditionalGeneration, WhisperProcessor

from rescore_gold69_v2_fair import fair_text


def read_audio(cell):
    wav, sr = sf.read(io.BytesIO(cell["bytes"]), dtype="float32")
    if wav.ndim > 1:
        wav = wav.mean(axis=-1)
    return wav, sr


def load_gold(parquet_path: Path):
    table = pq.read_table(parquet_path)
    rows = []
    for row in table.to_pylist():
        wav, sr = read_audio(row["audio"])
        rows.append({"id": str(row["programmeitem_id"]), "wav": wav, "sr": sr, "ref": row["gold_transcript"] or ""})
    return rows


@torch.no_grad()
def run_model(model_id: str, rows, device: str):
    processor = WhisperProcessor.from_pretrained(model_id)
    model = WhisperForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.float16 if device.startswith("cuda") else torch.float32)
    model.to(device).eval()
    params = sum(p.numel() for p in model.parameters())
    preds = []
    times = []
    for i, row in enumerate(rows, 1):
        inputs = processor(row["wav"], sampling_rate=row["sr"], return_tensors="pt")
        feats = inputs.input_features.to(device)
        if device.startswith("cuda"):
            feats = feats.half()
            torch.cuda.synchronize()
        t0 = time.time()
        try:
            ids = model.generate(feats, language="fa", task="transcribe")
        except TypeError:
            ids = model.generate(feats)
        if device.startswith("cuda"):
            torch.cuda.synchronize()
        times.append(time.time() - t0)
        hyp = processor.batch_decode(ids, skip_special_tokens=True)[0]
        preds.append(hyp)
        print(f"[pred] {model_id} {i}/{len(rows)}", flush=True)
    del model
    if device.startswith("cuda"):
        torch.cuda.empty_cache()
    return preds, times, params


def score(rows, preds):
    import jiwer

    refs_words = [fair_text(r["ref"], keep_space=True) for r in rows]
    hyps_words = [fair_text(h, keep_space=True) for h in preds]
    refs_chars = [fair_text(r["ref"], keep_space=False) for r in rows]
    hyps_chars = [fair_text(h, keep_space=False) for h in preds]
    return jiwer.wer(refs_words, hyps_words), jiwer.cer(refs_chars, hyps_chars)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--gold-parquet", required=True)
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--device", default="cuda:0")
    ap.add_argument("models", nargs="+")
    args = ap.parse_args()

    rows = load_gold(Path(args.gold_parquet))
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    results = []
    for model_id in args.models:
        safe = model_id.replace("/", "__")
        pred_path = out_dir / f"{safe}__gold69_v2_predictions.jsonl"
        preds, times, params = run_model(model_id, rows, args.device)
        with pred_path.open("w") as f:
            for row, hyp in zip(rows, preds):
                f.write(json.dumps({"id": row["id"], "ref": row["ref"], "hyp": hyp}, ensure_ascii=False) + "\n")
        wer, cer = score(rows, preds)
        rec = {
            "model": model_id,
            "n": len(rows),
            "gold69_v2_fair_wer": wer,
            "gold69_v2_fair_cer": cer,
            "mean_decode_ms": float(np.mean(times) * 1000.0),
            "params_b": params / 1e9,
            "predictions": str(pred_path),
            "status": "complete",
        }
        results.append(rec)
        print(json.dumps(rec, ensure_ascii=False), flush=True)
    (out_dir / "whisper_gold69_v2_results.json").write_text(json.dumps(results, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()