#!/usr/bin/env python3
"""Normalize and score Echo and whisper.cpp common-track benchmark outputs."""

from __future__ import annotations

import argparse
import json
import re
import statistics
from pathlib import Path


def words(value: str) -> list[str]:
    return re.findall(r"[^\W_]+", value.casefold(), flags=re.UNICODE)


def edits(reference: list[str], hypothesis: list[str]) -> int:
    previous = list(range(len(hypothesis) + 1))
    for reference_index, reference_word in enumerate(reference):
        current = [reference_index + 1] + [0] * len(hypothesis)
        for hypothesis_index, hypothesis_word in enumerate(hypothesis):
            current[hypothesis_index + 1] = min(
                previous[hypothesis_index + 1] + 1,
                current[hypothesis_index] + 1,
                previous[hypothesis_index]
                + (reference_word != hypothesis_word),
            )
        previous = current
    return previous[-1]


def percentile(values: list[float], quantile: float) -> float:
    ordered = sorted(values)
    index = round((len(ordered) - 1) * quantile)
    return ordered[index]


def whisper_text(payload: dict) -> str:
    if isinstance(payload.get("transcription"), list):
        return " ".join(
            segment.get("text", "").strip()
            for segment in payload["transcription"]
            if segment.get("text")
        ).strip()
    if isinstance(payload.get("text"), str):
        return payload["text"].strip()
    if isinstance(payload.get("result"), dict):
        return str(payload["result"].get("text", "")).strip()
    raise ValueError("Unrecognized whisper.cpp JSON structure")


def load_rows(
    manifest: dict,
    engine: str,
    result_path: Path,
    benchmark_root: Path,
) -> tuple[list[dict], dict]:
    fixtures = {fixture["id"]: fixture for fixture in manifest["fixtures"]}
    if engine == "echo":
        payload = json.loads(result_path.read_text())
        rows = []
        for result in payload["results"]:
            fixture = fixtures[result["id"]]
            rows.append(
                {
                    "id": result["id"],
                    "split": "clean" if "_clean_" in result["id"] else "other",
                    "reference": fixture["reference"],
                    "hypothesis": result["transcript"],
                    "modelMs": result.get("modelMs"),
                    "audioDurationMs": result.get("audioDurationMs"),
                }
            )
        return rows, {
            "model": payload["model"],
            "pipeline": payload.get("pipeline", []),
        }

    rows = []
    for fixture in manifest["fixtures"]:
        audio_path = benchmark_root / fixture["audio"]
        candidates = [
            Path(str(audio_path) + ".json"),
            audio_path.with_suffix(".json"),
            result_path.parent / f"{fixture['id']}.json",
        ]
        output_path = next((path for path in candidates if path.exists()), None)
        if output_path is None:
            raise FileNotFoundError(f"No whisper result for {fixture['id']}")
        payload = json.loads(output_path.read_text())
        rows.append(
            {
                "id": fixture["id"],
                "split": "clean" if "_clean_" in fixture["id"] else "other",
                "reference": fixture["reference"],
                "hypothesis": whisper_text(payload),
                "modelMs": None,
                "audioDurationMs": None,
            }
        )
    return rows, {"model": "whisper-large-v3-turbo-q5_0", "pipeline": ["whisper.cpp"]}


def summarize(rows: list[dict], metadata: dict) -> dict:
    scored = []
    total_edits = 0
    total_reference_words = 0
    for row in rows:
        reference_words = words(row["reference"])
        hypothesis_words = words(row["hypothesis"])
        edit_count = edits(reference_words, hypothesis_words)
        total_edits += edit_count
        total_reference_words += len(reference_words)
        scored.append(
            {
                **row,
                "referenceWordCount": len(reference_words),
                "editCount": edit_count,
                "wordErrorRate": edit_count / len(reference_words),
            }
        )

    by_split = {}
    for split in ["clean", "other"]:
        subset = [row for row in scored if row["split"] == split]
        split_edits = sum(row["editCount"] for row in subset)
        split_words = sum(row["referenceWordCount"] for row in subset)
        by_split[split] = {
            "rows": len(subset),
            "meanUtteranceWER": statistics.mean(
                row["wordErrorRate"] for row in subset
            ),
            "corpusWER": split_edits / split_words,
        }

    model_times = [
        row["modelMs"] for row in scored if isinstance(row.get("modelMs"), (int, float))
    ]
    audio_times = [
        row["audioDurationMs"]
        for row in scored
        if isinstance(row.get("audioDurationMs"), (int, float))
    ]
    summary = {
        **metadata,
        "rows": len(scored),
        "meanUtteranceWER": statistics.mean(
            row["wordErrorRate"] for row in scored
        ),
        "corpusWER": total_edits / total_reference_words,
        "bySplit": by_split,
        "modelTiming": None,
    }
    if model_times:
        summary["modelTiming"] = {
            "meanMs": statistics.mean(model_times),
            "p50Ms": statistics.median(model_times),
            "p90Ms": percentile(model_times, 0.90),
            "totalMs": sum(model_times),
            "audioDurationMs": sum(audio_times),
            "realTimeMultiple": (
                sum(audio_times) / sum(model_times) if sum(model_times) else None
            ),
        }
    return {"summary": summary, "results": scored}


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--engine", choices=["echo", "whisper"], required=True)
    parser.add_argument("--result", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    arguments = parser.parse_args()

    manifest = json.loads(arguments.manifest.read_text())
    rows, metadata = load_rows(
        manifest,
        arguments.engine,
        arguments.result,
        arguments.manifest.resolve().parent,
    )
    report = summarize(rows, metadata)
    arguments.output.parent.mkdir(parents=True, exist_ok=True)
    arguments.output.write_text(
        json.dumps(report, indent=2, ensure_ascii=False) + "\n"
    )
    print(json.dumps(report["summary"], indent=2))


if __name__ == "__main__":
    main()
