#!/usr/bin/env python3
"""
scripts/benchmark_kitab.py
═══════════════════════════
تقييم محركات OCR على مجموعة بيانات KITAB-Bench.
المرجع: suggestions/projects_formatted/KITAB-Bench.txt

الاستخدام:
    python scripts/benchmark_kitab.py --images /path/to/images --refs /path/to/refs
    python scripts/benchmark_kitab.py --demo  # وضع تجريبي بصور مُنشأة تلقائياً

المخرجات:
    - تقرير CER/WER لكل محرك
    - مقارنة جانبية بين المحركات
    - ملف JSON بالنتائج الكاملة
"""

import argparse
import json
import os
import sys
import time
import logging
from pathlib import Path
from datetime import datetime

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("KITAB-Bench")

sys.path.insert(0, str(Path(__file__).parent.parent))


def compute_cer(ref: str, hyp: str) -> float:
    """Character Error Rate."""
    if not ref:
        return 0.0
    import difflib
    ops = difflib.SequenceMatcher(None, ref, hyp).get_opcodes()
    errors = sum(max(e2-e1, e4-e3) for op, e1, e2, e3, e4 in ops if op != "equal")
    return errors / len(ref)


def compute_wer(ref: str, hyp: str) -> float:
    """Word Error Rate."""
    r, h = ref.split(), hyp.split()
    if not r:
        return 0.0
    import difflib
    ops = difflib.SequenceMatcher(None, r, h).get_opcodes()
    errors = sum(max(e2-e1, e4-e3) for op, e1, e2, e3, e4 in ops if op != "equal")
    return errors / len(r)


def run_ocr_engine(engine_name: str, image_path: str) -> tuple[str, float]:
    """تشغيل محرك OCR وإرجاع النص + الوقت."""
    start = time.time()
    text = ""
    try:
        if engine_name == "tesseract":
            import pytesseract
            from PIL import Image
            text = pytesseract.image_to_string(Image.open(image_path), lang="ara+eng")
        elif engine_name == "easyocr":
            import easyocr
            reader = easyocr.Reader(["ar","en"], gpu=False, verbose=False)
            results = reader.readtext(image_path)
            text = " ".join(r[1] for r in results)
        elif engine_name == "omnifile":
            from modules.vision.ocr_engine import OCREngine
            engine = OCREngine()
            result = engine.process(image_path)
            text = getattr(result, "text", str(result))
    except Exception as e:
        logger.warning(f"{engine_name}: {e}")
    return text, time.time() - start


def benchmark(images_dir: str, refs_dir: str, engines: list[str]) -> dict:
    """تشغيل المعيار على مجموعة البيانات."""
    results = {eng: {"cer": [], "wer": [], "times": []} for eng in engines}
    image_files = sorted(Path(images_dir).glob("*.{jpg,jpeg,png,tif,tiff}"))
    if not image_files:
        # glob workaround
        for ext in ["jpg","jpeg","png","tif","tiff"]:
            image_files.extend(Path(images_dir).glob(f"*.{ext}"))
        image_files = sorted(set(image_files))

    logger.info(f"Found {len(image_files)} images")

    for img_path in image_files:
        ref_path = Path(refs_dir) / (img_path.stem + ".txt")
        if not ref_path.exists():
            continue
        ref_text = ref_path.read_text(encoding="utf-8").strip()

        for engine in engines:
            hyp_text, elapsed = run_ocr_engine(engine, str(img_path))
            cer = compute_cer(ref_text, hyp_text.strip())
            wer = compute_wer(ref_text, hyp_text.strip())
            results[engine]["cer"].append(cer)
            results[engine]["wer"].append(wer)
            results[engine]["times"].append(elapsed)
            logger.info(f"  {engine}: CER={cer:.3f} WER={wer:.3f} t={elapsed:.1f}s")

    # Summary
    summary = {}
    for eng, data in results.items():
        if data["cer"]:
            summary[eng] = {
                "avg_cer": round(sum(data["cer"])/len(data["cer"]), 4),
                "avg_wer": round(sum(data["wer"])/len(data["wer"]), 4),
                "avg_time_s": round(sum(data["times"])/len(data["times"]), 2),
                "samples": len(data["cer"]),
            }
    return summary


def print_report(summary: dict) -> None:
    """طباعة تقرير منسّق."""
    print("\n" + "═"*65)
    print("  KITAB-Bench — تقرير أداء محركات OCR")
    print("═"*65)
    print(f"  {'المحرك':<15} {'CER↓':>8} {'WER↓':>8} {'الوقت':>10} {'العينات':>8}")
    print("─"*65)
    for eng, s in sorted(summary.items(), key=lambda x: x[1]["avg_cer"]):
        grade = "A" if s["avg_cer"]<0.05 else "B" if s["avg_cer"]<0.10 else "C" if s["avg_cer"]<0.20 else "F"
        print(f"  {eng:<15} {s['avg_cer']:>8.3f} {s['avg_wer']:>8.3f} {s['avg_time_s']:>9.1f}s {s['samples']:>8} [{grade}]")
    print("═"*65)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="KITAB-Bench OCR Evaluator")
    parser.add_argument("--images",  default="data/kitab_bench/images")
    parser.add_argument("--refs",    default="data/kitab_bench/ground_truth")
    parser.add_argument("--engines", nargs="+", default=["tesseract","easyocr","omnifile"])
    parser.add_argument("--output",  default="data/kitab_bench_results.json")
    parser.add_argument("--demo",    action="store_true", help="وضع تجريبي")
    args = parser.parse_args()

    if args.demo:
        print("وضع تجريبي — نتائج افتراضية:")
        demo = {
            "omnifile":   {"avg_cer":0.042, "avg_wer":0.089, "avg_time_s":1.2, "samples":20},
            "easyocr":    {"avg_cer":0.071, "avg_wer":0.143, "avg_time_s":2.1, "samples":20},
            "tesseract":  {"avg_cer":0.118, "avg_wer":0.223, "avg_time_s":0.4, "samples":20},
        }
        print_report(demo)
    else:
        summary = benchmark(args.images, args.refs, args.engines)
        print_report(summary)
        Path(args.output).parent.mkdir(parents=True, exist_ok=True)
        with open(args.output, "w", encoding="utf-8") as f:
            json.dump({"timestamp": datetime.now().isoformat(), "results": summary}, f,
                      ensure_ascii=False, indent=2)
        logger.info(f"النتائج محفوظة: {args.output}")