#!/usr/bin/env python3
"""OCR Benchmark: Current system (PyMuPDF + Google Vision) vs Mistral OCR 4.0

Usage:
    python scripts/ocr_benchmark_mistral.py [--docs N] [--output PATH]

Downloads PDFs from MinIO, calls Mistral OCR API, compares against
already-extracted text stored in the DB, and writes a Markdown report.
"""

from __future__ import annotations

import argparse
import asyncio
import base64
import json
import re
import subprocess
import sys
import tempfile
import time
from pathlib import Path
from typing import TypedDict

import httpx

# ── Config ───────────────────────────────────────────────────────────────────

MISTRAL_API_KEY = "UsZjLCX30ev6pox0KgvXyuFP3ktsPYpN"
MISTRAL_OCR_MODEL = "mistral-ocr-latest"
MISTRAL_OCR_URL = "https://api.mistral.ai/v1/ocr"
MINIO_ALIAS = "legalminio"
MINIO_BUCKET = "legal-documents"

# Documents to benchmark — selected for diversity:
#  - main appeal  (40 pp, 4.4 MB, large digital doc)
#  - permit       (9 pp,  1.7 MB, likely partially scanned)
#  - protocol     (5 pp,  178 KB, administrative typed)
#  - response     (12 pp, ~390 KB, digital legal brief)
DOCS_TO_BENCHMARK = [
    {
        "title": "כתב ערר",
        "minio_key": "cases/1027-04-26/documents/originals/כתב ערר - מושב נחם נ׳ ועדה מקומית בית שמש.pdf",
        "doc_type": "appeal",
        "pages": 40,
    },
    {
        "title": "נספח 1 — היתר הבנייה",
        "minio_key": "cases/1027-04-26/documents/originals/נספח 1 - היתר הבנייה (פורסם 17.03.26).pdf",
        "doc_type": "permit",
        "pages": 9,
    },
    {
        "title": "נספח 14 — פרוטוקול ועדה מחוזית",
        "minio_key": "cases/1027-04-26/documents/originals/נספח 14 - פרוטוקול דיון ועדה מחוזית - 06.06.23.pdf",
        "doc_type": "protocol",
        "pages": 5,
    },
    {
        "title": "תגובת המשיבה 3",
        "minio_key": "cases/1027-04-26/documents/originals/תגובת המשיבה 3 לערר ולבקשה להתליית היתר בנייה.pdf",
        "doc_type": "response",
        "pages": 12,
    },
]

# ── DB access (read-only: fetch already-extracted text) ──────────────────────

def _fetch_extracted_texts(case_number: str) -> dict[str, str]:
    """Pull extracted_text from DB for all docs in the case."""
    import psycopg2  # type: ignore
    conn = psycopg2.connect(
        host="localhost", port=5433, dbname="legal_ai",
        user="legal_ai", password="od0ASJZFYibOlWK59krLvvETmgqwlXe8",
    )
    cur = conn.cursor()
    cur.execute(
        """
        SELECT d.title, d.extracted_text, d.file_path, d.page_count
        FROM documents d
        JOIN cases c ON c.id = d.case_id
        WHERE c.case_number = %s AND d.extracted_text IS NOT NULL
        """,
        (case_number,),
    )
    rows = cur.fetchall()
    cur.close()
    conn.close()
    return {row[0]: {"text": row[1], "file_path": row[2], "pages": row[3]} for row in rows}


# ── MinIO download ────────────────────────────────────────────────────────────

def download_from_minio(minio_key: str, dest: Path) -> None:
    """Download a file from MinIO using mcli."""
    src = f"{MINIO_ALIAS}/{MINIO_BUCKET}/{minio_key}"
    result = subprocess.run(
        ["mcli", "cp", src, str(dest)],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise RuntimeError(f"mcli cp failed: {result.stderr}")


# ── Text quality metrics ─────────────────────────────────────────────────────

_HEBREW_RE = re.compile(r'[֐-׿]')
_WORD_RE = re.compile(r'\S+')


def compute_metrics(text: str) -> dict:
    """Compute quality metrics for extracted text."""
    if not text:
        return {"chars": 0, "words": 0, "avg_word_len": 0,
                "single_char_pct": 0, "words_per_line": 0,
                "hebrew_pct": 0, "quality_ok": False}

    words = _WORD_RE.findall(text)
    n_words = len(words)
    if n_words == 0:
        return {"chars": len(text), "words": 0, "avg_word_len": 0,
                "single_char_pct": 0, "words_per_line": 0,
                "hebrew_pct": 0, "quality_ok": False}

    avg_len = sum(len(w) for w in words) / n_words
    single_char_pct = sum(1 for w in words if len(w) == 1) / n_words

    lines = [ln for ln in text.split("\n") if ln.strip()]
    words_per_line = n_words / len(lines) if lines else 0

    letters = re.findall(r'[a-zA-Z֐-׿]', text)
    hebrew_pct = (
        sum(1 for c in letters if _HEBREW_RE.match(c)) / len(letters)
        if letters else 0
    )

    quality_ok = (
        n_words >= 10
        and avg_len >= 2.5
        and single_char_pct <= 0.4
        and words_per_line >= 3.0
        and hebrew_pct >= 0.5
    )

    return {
        "chars": len(text),
        "words": n_words,
        "avg_word_len": round(avg_len, 2),
        "single_char_pct": round(single_char_pct * 100, 1),
        "words_per_line": round(words_per_line, 1),
        "hebrew_pct": round(hebrew_pct * 100, 1),
        "quality_ok": quality_ok,
    }


def count_known_abbrev_errors(text: str) -> int:
    """Count common Hebrew abbreviation OCR errors (pre-fix indicators)."""
    patterns = ['עוהייד', 'עוייד', 'הנייל', 'ביהמייש', 'עייי', 'בייכ', 'תמייא']
    return sum(text.count(p) for p in patterns)


def count_correct_abbrevs(text: str) -> int:
    """Count correctly rendered Hebrew abbreviations."""
    correct = ['עו"ד', 'הנ"ל', 'ביהמ"ש', 'ע"י', 'ב"כ', 'תמ"א', 'ס"ק']
    return sum(text.count(p) for p in correct)


# ── Mistral OCR ───────────────────────────────────────────────────────────────

def call_mistral_ocr(pdf_path: Path) -> tuple[str, float]:
    """Call Mistral OCR API on a PDF. Returns (extracted_text, elapsed_seconds)."""
    with open(pdf_path, "rb") as f:
        pdf_b64 = base64.b64encode(f.read()).decode()

    payload = {
        "model": MISTRAL_OCR_MODEL,
        "document": {
            "type": "document_url",
            "document_url": f"data:application/pdf;base64,{pdf_b64}",
        },
        "include_image_base64": False,
    }

    t0 = time.time()
    with httpx.Client(timeout=300.0) as client:
        resp = client.post(
            MISTRAL_OCR_URL,
            headers={
                "Authorization": f"Bearer {MISTRAL_API_KEY}",
                "Content-Type": "application/json",
            },
            json=payload,
        )
    elapsed = time.time() - t0

    if resp.status_code != 200:
        raise RuntimeError(f"Mistral OCR error {resp.status_code}: {resp.text[:500]}")

    data = resp.json()
    # Response: {"pages": [{"index": 0, "markdown": "..."}, ...]}
    pages = data.get("pages", [])
    text = "\n\n".join(p.get("markdown", "") for p in pages)
    return text, elapsed


# ── Report ────────────────────────────────────────────────────────────────────

def render_metric_table(current: dict, mistral: dict) -> str:
    rows = [
        ("תווים", f"{current['chars']:,}", f"{mistral['chars']:,}"),
        ("מילים", f"{current['words']:,}", f"{mistral['words']:,}"),
        ("אורך מילה ממוצע", str(current['avg_word_len']), str(mistral['avg_word_len'])),
        ("% מילים חד-תוויות", f"{current['single_char_pct']}%", f"{mistral['single_char_pct']}%"),
        ("מילים לשורה", str(current['words_per_line']), str(mistral['words_per_line'])),
        ("% תווים עבריים", f"{current['hebrew_pct']}%", f"{mistral['hebrew_pct']}%"),
        ("איכות כוללת", "✅" if current['quality_ok'] else "❌", "✅" if mistral['quality_ok'] else "❌"),
    ]
    lines = ["| מדד | OCR נוכחי | Mistral OCR |", "|-----|-----------|-------------|"]
    for label, cur_val, mis_val in rows:
        lines.append(f"| {label} | {cur_val} | {mis_val} |")
    return "\n".join(lines)


def build_report(results: list[dict], output_path: Path) -> None:
    lines = [
        "# השוואת OCR: מערכת נוכחית מול Mistral OCR",
        f"\n**תיק:** 1027-04-26 — בל\"מ מושב נחם מפעל בטון בית שמש  ",
        f"**תאריך:** {time.strftime('%Y-%m-%d %H:%M')}  ",
        f"**מודל Mistral:** `{MISTRAL_OCR_MODEL}`  ",
        f"**מערכת נוכחית:** PyMuPDF (born-digital) + Google Cloud Vision (scanned)  ",
        "\n---\n",
        "## סיכום מנהלים\n",
    ]

    # Summary table
    sum_lines = ["| מסמך | עמודים | נוכחי תווים | Mistral תווים | Mistral זמן | עדיפות |",
                 "|------|--------|-------------|---------------|-------------|--------|"]
    for r in results:
        if "error" in r:
            sum_lines.append(f"| {r['title']} | {r['pages']} | — | שגיאה | — | — |")
            continue
        cur_chars = r["current_metrics"]["chars"]
        mis_chars = r["mistral_metrics"]["chars"]
        winner = "🔵 נוכחי" if cur_chars > mis_chars * 1.05 else (
            "🟢 Mistral" if mis_chars > cur_chars * 1.05 else "⚖️ שקול")
        sum_lines.append(
            f"| {r['title']} | {r['pages']} | {cur_chars:,} | {mis_chars:,} | "
            f"{r['mistral_elapsed']:.1f}s | {winner} |"
        )
    lines.extend(sum_lines)
    lines.append("\n---\n")

    # Per-document detail
    for r in results:
        lines.append(f"## {r['title']} ({r['pages']} עמודים)\n")
        if "error" in r:
            lines.append(f"**שגיאה ב-Mistral OCR:** `{r['error']}`\n")
            continue

        lines.append(render_metric_table(r["current_metrics"], r["mistral_metrics"]))
        lines.append("")

        cur_abbr_err = count_known_abbrev_errors(r["current_text"])
        mis_abbr_err = count_known_abbrev_errors(r["mistral_text"])
        cur_abbr_ok = count_correct_abbrevs(r["current_text"])
        mis_abbr_ok = count_correct_abbrevs(r["mistral_text"])

        lines.append(f"\n**קיצורים עבריים:**")
        lines.append(f"- נוכחי: {cur_abbr_ok} נכונים, {cur_abbr_err} שגויים")
        lines.append(f"- Mistral: {mis_abbr_ok} נכונים, {mis_abbr_err} שגויים")
        lines.append(f"\n**זמן Mistral:** {r['mistral_elapsed']:.1f} שניות\n")

        # Side-by-side first 600 chars
        lines.append("### דוגמת טקסט — 600 תווים ראשונים\n")
        lines.append("**מערכת נוכחית:**")
        lines.append("```")
        lines.append((r["current_text"] or "")[:600].replace("```", "'''"))
        lines.append("```\n")
        lines.append("**Mistral OCR:**")
        lines.append("```")
        lines.append((r["mistral_text"] or "")[:600].replace("```", "'''"))
        lines.append("```\n")
        lines.append("---\n")

    output_path.write_text("\n".join(lines), encoding="utf-8")
    print(f"\n✅ דוח נשמר: {output_path}")


# ── Main ──────────────────────────────────────────────────────────────────────

def main() -> None:
    parser = argparse.ArgumentParser(description="OCR benchmark: current vs Mistral")
    parser.add_argument("--docs", type=int, default=4,
                        help="כמה מסמכים לבדוק (ברירת מחדל: 4)")
    parser.add_argument("--output", type=str,
                        default="/home/chaim/legal-ai/data/audit/ocr-benchmark-mistral.md",
                        help="נתיב לדוח הפלט")
    args = parser.parse_args()

    docs = DOCS_TO_BENCHMARK[: args.docs]
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    print(f"🔍 שולף טקסטים קיימים מה-DB עבור תיק 1027-04-26...")
    db_texts = _fetch_extracted_texts("1027-04-26")
    print(f"   נמצאו {len(db_texts)} מסמכים עם טקסט מחולץ")

    results = []

    with tempfile.TemporaryDirectory() as tmp_dir:
        for doc in docs:
            print(f"\n📄 מעבד: {doc['title']} ({doc['pages']} עמודים)")

            # --- Current OCR text from DB ---
            current_text = ""
            for title, info in db_texts.items():
                if doc["title"].split("—")[0].strip() in title or doc["minio_key"].split("/")[-1] in (info.get("file_path") or ""):
                    current_text = info["text"] or ""
                    break
            if not current_text:
                # Try by file_path suffix match
                key_name = doc["minio_key"].split("/")[-1]
                for title, info in db_texts.items():
                    if key_name in (info.get("file_path") or ""):
                        current_text = info["text"] or ""
                        break

            if not current_text:
                print(f"   ⚠️  לא נמצא טקסט ב-DB, מחפש לפי שם מסמך...")
                # fallback: match by doc_type + rough title
                for title, info in db_texts.items():
                    if doc["doc_type"] in title.lower() or doc["title"][:8] in title:
                        current_text = info["text"] or ""
                        break

            print(f"   נוכחי: {len(current_text):,} תווים")

            # --- Download PDF from MinIO ---
            pdf_name = doc["minio_key"].split("/")[-1]
            pdf_path = Path(tmp_dir) / pdf_name
            print(f"   מוריד מ-MinIO...")
            try:
                download_from_minio(doc["minio_key"], pdf_path)
                print(f"   הורד: {pdf_path.stat().st_size:,} bytes")
            except Exception as e:
                print(f"   ❌ שגיאה בהורדה: {e}")
                results.append({"title": doc["title"], "pages": doc["pages"], "error": str(e)})
                continue

            # --- Mistral OCR ---
            print(f"   קורא Mistral OCR API...")
            try:
                mistral_text, elapsed = call_mistral_ocr(pdf_path)
                print(f"   Mistral: {len(mistral_text):,} תווים ({elapsed:.1f}s)")
            except Exception as e:
                print(f"   ❌ שגיאת Mistral API: {e}")
                results.append({
                    "title": doc["title"], "pages": doc["pages"],
                    "current_text": current_text,
                    "current_metrics": compute_metrics(current_text),
                    "mistral_text": "", "mistral_metrics": compute_metrics(""),
                    "mistral_elapsed": 0, "error": str(e),
                })
                continue

            results.append({
                "title": doc["title"],
                "pages": doc["pages"],
                "current_text": current_text,
                "current_metrics": compute_metrics(current_text),
                "mistral_text": mistral_text,
                "mistral_metrics": compute_metrics(mistral_text),
                "mistral_elapsed": elapsed,
            })

    print("\n📊 בונה דוח...")
    build_report(results, output_path)

    # Also save raw texts for manual inspection
    raw_dir = output_path.parent / "ocr-benchmark-raw"
    raw_dir.mkdir(exist_ok=True)
    for r in results:
        safe = r["title"].replace("/", "-").replace(" ", "_")[:40]
        if "current_text" in r:
            (raw_dir / f"{safe}_current.txt").write_text(r["current_text"], encoding="utf-8")
        if "mistral_text" in r:
            (raw_dir / f"{safe}_mistral.txt").write_text(r["mistral_text"], encoding="utf-8")
    print(f"💾 טקסטים גולמיים נשמרו: {raw_dir}")


if __name__ == "__main__":
    main()
