"""Run accuracy benchmarks on evaluation images.

Usage:
    python benchmarks/run_accuracy.py \
        --model-path AutomatosX/AX-Unlimited-OCR-3B-MoE-MLX-MXFP8 \
        --eval-dir ./benchmarks/eval_data/ \
        --output benchmarks/results/accuracy_results.json
"""

from __future__ import annotations

import argparse
from collections import Counter
import json
import os
import sys
import time
from pathlib import Path

# Add the project and src roots for both package and benchmark imports when the
# file is executed directly (``python benchmarks/run_accuracy.py``).
PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(PROJECT_ROOT / "src"))

from benchmarks.evaluate_cer import (  # noqa: E402
    compute_cer, compute_wer, compute_digit_cer,
    compute_cjk_cer, detect_repetition, extract_cjk_characters,
)
from benchmarks.normalize_output import normalize_ocr_output  # noqa: E402
from benchmarks.evaluate_tables import evaluate_tables  # noqa: E402
from unlimited_ocr.engine import SlidingWindowNoRepeatNGramProcessor  # noqa: E402
from unlimited_ocr.profiles import get_profile  # noqa: E402


def load_category_map(eval_dir: Path) -> dict[str, str]:
    """Map image filenames to categories from ``manifest.json`` when present."""
    manifest_path = Path(eval_dir) / "manifest.json"
    if not manifest_path.is_file():
        return {}
    try:
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    except (OSError, UnicodeError, json.JSONDecodeError):
        return {}
    samples = manifest.get("samples")
    if not isinstance(samples, list):
        return {}
    mapping: dict[str, str] = {}
    for sample in samples:
        if not isinstance(sample, dict):
            continue
        category = sample.get("category")
        image = sample.get("image")
        if not isinstance(category, str) or not category or not isinstance(image, str):
            continue
        mapping[Path(image).name] = category
    return mapping


def aggregate_by_category(per_file: list[dict]) -> dict[str, dict]:
    """Aggregate CER-family metrics by category label on each per-file row."""
    buckets: dict[str, list[dict]] = {}
    for row in per_file:
        if not isinstance(row, dict):
            continue
        category = row.get("category")
        if not isinstance(category, str) or not category:
            category = "uncategorized"
        buckets.setdefault(category, []).append(row)

    summary: dict[str, dict] = {}
    for category, rows in sorted(buckets.items()):
        digit_scores = [r["digit_cer"] for r in rows if r.get("ref_digit_count", 0) > 0]
        cjk_scores = [r["cjk_cer"] for r in rows if r.get("ref_cjk_count", 0) > 0]
        table_scores = [r["table_score"] for r in rows if r.get("table_score") is not None]
        summary[category] = {
            "num_samples": len(rows),
            "mean_cer": (
                sum(r["cer"] for r in rows) / len(rows) if rows else None
            ),
            "mean_digit_cer": (
                sum(digit_scores) / len(digit_scores) if digit_scores else None
            ),
            "mean_cjk_cer": (
                sum(cjk_scores) / len(cjk_scores) if cjk_scores else None
            ),
            "mean_table_score": (
                sum(table_scores) / len(table_scores) if table_scores else None
            ),
            "num_digit_samples": len(digit_scores),
            "num_cjk_samples": len(cjk_scores),
            "num_table_samples": len(table_scores),
        }
    return summary


def run_accuracy_benchmark(
    model_path: str,
    eval_dir: Path,
    prompt: str = "<image>document parsing.",
    max_tokens: int | None = None,
    profile: str = "accurate",
    backend: str = "mlx",
    vllm_base_url: str = "http://127.0.0.1:8000/v1",
    vllm_api_key: str | None = None,
    served_revision: str | None = None,
    save_predictions: Path | None = None,
) -> dict:
    """Run accuracy benchmark on evaluation dataset.

    Expects eval_dir structure:
        eval_dir/
        ├── images/          # Input images
        └── ground_truth/    # Reference .txt files (same stem)
    """
    images_dir = eval_dir / "images"
    gt_dir = eval_dir / "ground_truth"

    if not images_dir.is_dir():
        raise FileNotFoundError(f"Images directory not found: {images_dir}")
    if not gt_dir.is_dir():
        raise FileNotFoundError(f"Ground truth directory not found: {gt_dir}")

    if not isinstance(prompt, str) or prompt.count("<image>") != 1:
        raise ValueError("prompt must contain exactly one literal '<image>' token")
    generation_profile = get_profile(profile)
    token_limit = generation_profile.max_tokens if max_tokens is None else max_tokens
    if not isinstance(token_limit, int) or isinstance(token_limit, bool) or token_limit < 1:
        raise ValueError("max_tokens must be a positive integer")
    if backend not in {"mlx", "vllm"}:
        raise ValueError("backend must be one of: mlx, vllm")
    if served_revision is not None and (
        not isinstance(served_revision, str) or not served_revision.strip()
    ):
        raise ValueError("served_revision must be a non-empty string or None")

    # Collect image files
    image_files = sorted(
        f for f in images_dir.iterdir()
        if f.is_file() and f.suffix.lower() in {
            ".jpg", ".jpeg", ".png", ".tiff", ".tif", ".webp", ".bmp"
        }
    )

    print(f"Found {len(image_files)} images")
    if not image_files:
        raise ValueError(f"No supported images found in: {images_dir}")

    duplicate_stems = sorted(
        stem for stem, count in Counter(path.stem for path in image_files).items()
        if count > 1
    )
    if duplicate_stems:
        raise ValueError(
            "Multiple input images share the same stem: "
            + ", ".join(duplicate_stems)
        )

    missing_ground_truth = [
        image_file.name
        for image_file in image_files
        if not (gt_dir / f"{image_file.stem}.txt").is_file()
    ]
    if missing_ground_truth:
        raise FileNotFoundError(
            "Missing ground-truth text for: " + ", ".join(missing_ground_truth)
        )

    evaluation_files = [
        (image_file, gt_dir / f"{image_file.stem}.txt")
        for image_file in image_files
    ]
    category_map = load_category_map(eval_dir)
    if save_predictions is not None:
        save_predictions = Path(save_predictions)
        save_predictions.mkdir(parents=True, exist_ok=True)

    print(f"Loading model: {model_path}")
    if backend == "mlx":
        # Import only after validating the dataset and generation options.
        from mlx_vlm import generate as mlx_generate, load

        model, processor = load(model_path, revision=served_revision)
        return _evaluate_accuracy_files(
            model_path=model_path,
            evaluation_files=evaluation_files,
            prompt=prompt,
            token_limit=token_limit,
            profile=profile,
            backend=backend,
            served_revision=served_revision,
            generation_profile=generation_profile,
            model=model,
            processor=processor,
            generate=mlx_generate,
            engine=None,
            category_map=category_map,
            save_predictions=save_predictions,
        )

    from unlimited_ocr.vllm_backend import VLLMEngine, VLLMEngineConfig

    engine = VLLMEngine(
        VLLMEngineConfig(
            model_path=model_path,
            base_url=vllm_base_url,
            api_key=vllm_api_key,
            max_tokens=token_limit,
            profile=profile,
        )
    )
    try:
        engine.ensure_ready()
        return _evaluate_accuracy_files(
            model_path=model_path,
            evaluation_files=evaluation_files,
            prompt=prompt,
            token_limit=token_limit,
            profile=profile,
            backend=backend,
            served_revision=served_revision,
            generation_profile=generation_profile,
            model=None,
            processor=None,
            generate=None,
            engine=engine,
            category_map=category_map,
            save_predictions=save_predictions,
        )
    finally:
        engine.unload()


def _evaluate_accuracy_files(
    *,
    model_path: str,
    evaluation_files: list[tuple[Path, Path]],
    prompt: str,
    token_limit: int,
    profile: str,
    backend: str,
    served_revision: str | None,
    generation_profile,
    model,
    processor,
    generate,
    engine,
    category_map: dict[str, str] | None = None,
    save_predictions: Path | None = None,
) -> dict:
    """Evaluate already-loaded MLX or vLLM inference resources."""

    results = []
    total_time = 0.0

    for i, (img_file, gt_file) in enumerate(evaluation_files):
        reference = normalize_ocr_output(gt_file.read_text(encoding="utf-8"))

        # Run inference
        start = time.perf_counter()
        if backend == "mlx":
            response = generate(
                model, processor,
                prompt=prompt,
                image=[str(img_file)],
                max_tokens=token_limit,
                temperature=generation_profile.temperature,
                top_p=generation_profile.top_p,
                repetition_penalty=generation_profile.repetition_penalty,
                logits_processors=[SlidingWindowNoRepeatNGramProcessor(
                    generation_profile.no_repeat_ngram_size,
                    128,
                )],
                verbose=False,
            )
        else:
            response = engine.infer(
                str(img_file),
                custom_prompt=prompt,
                max_tokens=token_limit,
                profile=profile,
            )
        elapsed = time.perf_counter() - start
        total_time += elapsed

        # Extract text
        if isinstance(response, str):
            hypothesis = response
        elif hasattr(response, "text"):
            hypothesis = response.text
        else:
            raise TypeError(
                "Inference backend returned an unsupported accuracy response "
                f"({type(response).__name__})"
            )
        hypothesis = normalize_ocr_output(hypothesis)

        # Compute metrics
        cer = compute_cer(reference, hypothesis)
        wer = compute_wer(reference, hypothesis)
        digit_cer = compute_digit_cer(reference, hypothesis)
        cjk_cer = compute_cjk_cer(reference, hypothesis)
        repetition = detect_repetition(hypothesis)
        table_result = evaluate_tables(reference, hypothesis)
        table_score = (
            table_result["mean_score"]
            if table_result["num_ref_tables"] > 0
            else None
        )

        category = (category_map or {}).get(img_file.name)
        result = {
            "file": img_file.name,
            "category": category,
            "cer": cer,
            "wer": wer,
            "digit_cer": digit_cer,
            "cjk_cer": cjk_cer,
            "repetition_rate": repetition,
            "table_score": table_score,
            "elapsed_seconds": elapsed,
            "ref_length": len(reference),
            "hyp_length": len(hypothesis),
            "ref_digit_count": sum(
                char.isascii() and char.isdigit() for char in reference
            ),
            "ref_cjk_count": len(extract_cjk_characters(reference)),
        }
        results.append(result)

        if save_predictions is not None:
            (save_predictions / f"{img_file.stem}.txt").write_text(
                hypothesis + ("\n" if hypothesis and not hypothesis.endswith("\n") else ""),
                encoding="utf-8",
            )

        print(f"  [{i+1}/{len(evaluation_files)}] {img_file.name}: CER={cer:.4f} "
              f"DigitCER={digit_cer:.4f} ({elapsed:.1f}s)")

    # Aggregate
    import numpy as np
    digit_scores = [r["digit_cer"] for r in results if r["ref_digit_count"] > 0]
    cjk_scores = [r["cjk_cer"] for r in results if r["ref_cjk_count"] > 0]
    table_scores = [r["table_score"] for r in results if r["table_score"] is not None]
    generation_settings = {
        "temperature": generation_profile.temperature,
        "top_p": generation_profile.top_p,
        "repetition_penalty": generation_profile.repetition_penalty,
        "no_repeat_ngram_size": generation_profile.no_repeat_ngram_size,
        "ngram_window": 128,
    }
    if backend == "vllm":
        # The official runtime recipe always uses 35, even when an MLX
        # generation profile has an experimental alternate value.
        generation_settings.update({
            "no_repeat_ngram_size": 35,
            "skip_special_tokens": False,
        })
    by_category = aggregate_by_category(results)
    summary = {
        "model_path": Path(model_path).name if Path(model_path).is_dir() else model_path,
        "served_revision": served_revision,
        "backend": backend,
        "prompt": prompt,
        "max_tokens": token_limit,
        "profile": profile,
        "generation_settings": generation_settings,
        "num_images": len(evaluation_files),
        "num_samples": len(results),
        "total_time_seconds": total_time,
        "mean_cer": float(np.mean([r["cer"] for r in results])) if results else 0,
        "mean_wer": float(np.mean([r["wer"] for r in results])) if results else 0,
        "num_digit_samples": len(digit_scores),
        "num_cjk_samples": len(cjk_scores),
        "mean_digit_cer": float(np.mean(digit_scores)) if digit_scores else None,
        "mean_cjk_cer": float(np.mean(cjk_scores)) if cjk_scores else None,
        "num_table_samples": len(table_scores),
        "mean_table_score": float(np.mean(table_scores)) if table_scores else None,
        "mean_repetition_rate": float(np.mean([r["repetition_rate"] for r in results])) if results else 0,
        "by_category": by_category,
        "per_file": results,
    }

    return summary


def main():
    parser = argparse.ArgumentParser(description="Run OCR accuracy benchmark")
    parser.add_argument("--model-path", required=True)
    parser.add_argument("--eval-dir", required=True, type=Path)
    parser.add_argument("--output", type=Path, default=Path("benchmarks/results/accuracy_results.json"))
    parser.add_argument("--prompt", default="<image>document parsing.")
    parser.add_argument("--max-tokens", type=int, default=None,
                        help="Override the selected profile's token limit")
    parser.add_argument("--profile", default="accurate",
                        choices=["accurate", "fast", "long-document", "plain-text", "markdown"])
    parser.add_argument("--backend", default="mlx", choices=["mlx", "vllm"])
    parser.add_argument("--vllm-base-url", default="http://127.0.0.1:8000/v1")
    parser.add_argument("--vllm-api-key", default=os.environ.get("VLLM_API_KEY"))
    parser.add_argument(
        "--served-revision",
        default=None,
        help="Revision pinned when this model was started (required for BF16 release evidence)",
    )
    parser.add_argument(
        "--save-predictions",
        type=Path,
        default=None,
        help="Optional directory to write normalized hypothesis text per image stem",
    )
    args = parser.parse_args()

    print("=" * 60)
    print("Unlimited-OCR Accuracy Benchmark")
    print("=" * 60)

    results = run_accuracy_benchmark(
        model_path=args.model_path,
        eval_dir=args.eval_dir,
        prompt=args.prompt,
        max_tokens=args.max_tokens,
        profile=args.profile,
        backend=args.backend,
        vllm_base_url=args.vllm_base_url,
        vllm_api_key=args.vllm_api_key,
        served_revision=args.served_revision,
        save_predictions=args.save_predictions,
    )

    # Save results
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(results, indent=2, ensure_ascii=False, allow_nan=False),
        encoding="utf-8",
    )

    def metric(value):
        return "N/A" if value is None else f"{value:.4f}"

    print(f"\n{'='*60}")
    print(f"Results: CER={results['mean_cer']:.4f} | "
          f"DigitCER={metric(results['mean_digit_cer'])} | "
          f"CJK={metric(results['mean_cjk_cer'])}")
    if results.get("by_category"):
        print("By category:")
        for category, stats in results["by_category"].items():
            print(
                f"  {category}: n={stats['num_samples']} "
                f"CER={metric(stats['mean_cer'])} "
                f"Digit={metric(stats['mean_digit_cer'])} "
                f"CJK={metric(stats['mean_cjk_cer'])}"
            )
    print(f"Saved to: {args.output}")


if __name__ == "__main__":
    main()