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

import argparse
import json
import re
import unicodedata
from pathlib import Path

import jiwer
import pyarrow.parquet as pq

PERSIAN_DIGITS = "۰۱۲۳۴۵۶۷۸۹٠١٢٣٤٥٦٧٨٩"
ASCII_DIGITS = "01234567890123456789"
DIGIT_MAP = str.maketrans(PERSIAN_DIGITS, ASCII_DIGITS)
DIACRITICS_RE = re.compile(r"[\u064b-\u065f\u0670\u06d6-\u06ed]")
PUNCT_SYMBOL_RE = re.compile(r"[\W_]+", re.UNICODE)


def fair_text(text: str, keep_space: bool) -> str:
    text = unicodedata.normalize("NFKC", text or "")
    text = text.replace("ي", "ی").replace("ك", "ک").replace("ۀ", "ه").replace("ة", "ه")
    text = text.replace("أ", "ا").replace("إ", "ا").replace("ٱ", "ا").replace("آ", "ا")
    text = text.replace("ؤ", "و").replace("ئ", "ی").replace("ى", "ی")
    text = text.translate(DIGIT_MAP)
    text = DIACRITICS_RE.sub("", text)
    text = text.replace("\u200c", " ").replace("\u200d", "").replace("\u200b", "")
    text = PUNCT_SYMBOL_RE.sub(" ", text)
    text = re.sub(r"\s+", " ", text).strip()
    if not keep_space:
        text = text.replace(" ", "")
    return text


def score_pair(ref_list: list[str], hyp_list: list[str]) -> dict:
    refs_words = [fair_text(r, keep_space=True) for r in ref_list]
    hyps_words = [fair_text(h, keep_space=True) for h in hyp_list]
    refs_chars = [fair_text(r, keep_space=False) for r in ref_list]
    hyps_chars = [fair_text(h, keep_space=False) for h in hyp_list]
    return {
        "gold69_v2_fair_wer": jiwer.wer(refs_words, hyps_words),
        "gold69_v2_fair_cer": jiwer.cer(refs_chars, hyps_chars),
        "gold69_v2_space_insensitive_wer": jiwer.wer(refs_chars, hyps_chars),
        "gold69_v2_space_insensitive_cer": jiwer.cer(refs_chars, hyps_chars),
    }


def load_refs(parquet_path: Path) -> dict[str, str]:
    table = pq.read_table(parquet_path, columns=["programmeitem_id", "gold_transcript"])
    ids = table.column("programmeitem_id").to_pylist()
    refs = table.column("gold_transcript").to_pylist()
    return {str(i): r or "" for i, r in zip(ids, refs)}


def load_preds(path: Path) -> dict[str, str]:
    out = {}
    with path.open() as f:
        for line in f:
            row = json.loads(line)
            out[str(row["id"])] = row.get("hyp") or ""
    return out


def score(refs: dict[str, str], preds: dict[str, str]) -> dict:
    ids = [i for i in refs if i in preds]
    missing = [i for i in refs if i not in preds]
    refs_words = [fair_text(refs[i], keep_space=True) for i in ids]
    hyps_words = [fair_text(preds[i], keep_space=True) for i in ids]
    refs_chars = [fair_text(refs[i], keep_space=False) for i in ids]
    hyps_chars = [fair_text(preds[i], keep_space=False) for i in ids]
    return {
        "n": len(ids),
        "missing_ids": missing,
        "gold69_v2_fair_wer": jiwer.wer(refs_words, hyps_words),
        "gold69_v2_fair_cer": jiwer.cer(refs_chars, hyps_chars),
        "gold69_v2_space_insensitive_wer": jiwer.wer(refs_chars, hyps_chars),
        "gold69_v2_space_insensitive_cer": jiwer.cer(refs_chars, hyps_chars),
    }


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--gold-parquet", required=True)
    ap.add_argument("--predictions-root", default=".")
    ap.add_argument("--out", required=True)
    args = ap.parse_args()

    root = Path(args.predictions_root)
    refs = load_refs(Path(args.gold_parquet))
    pred_files = {
        "VisualEars v2 final greedy": root / "visualears_text_policy/eval/step1568_spoken_greedy/predictions/gold69_predictions.jsonl",
        "VisualEars v2 final + 4gram LM": root / "visualears_text_policy/eval/step1568_spoken_lm_a0p5_b1_beam50/predictions/gold69_predictions.jsonl",
        "VisualEars v2 checkpoint-1000": root / "visualears_eval/results/checkpoint1000_ddp_bestknown/predictions/gold69_predictions.jsonl",
        "VisualEars step7000": root / "remote_artifacts/visualears-xlsr-300m-target269h-step7000/double_eval/results/gold69_predictions.jsonl",
        "student_mms1b": root / "stallion_predictions/mms1b_trained_double_eval_results/gold69_predictions.jsonl",
        "Peacockery/omni-ctc-300m-v2-fleurs-fa-ir-thomcles-continue": root / "persian_asr_leaderboard/Peacockery__omni-ctc-300m-v2-fleurs-fa-ir-thomcles-continue__gold69.jsonl",
    }
    rows = []
    for model, path in pred_files.items():
        if not path.exists():
            rows.append({"model": model, "status": "missing_predictions", "predictions": str(path)})
            continue
        rec = score(refs, load_preds(path))
        rec.update({"model": model, "status": "complete", "predictions": str(path)})
        rows.append(rec)

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(rows, ensure_ascii=False, indent=2))
    print(json.dumps(rows, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()