"""
ocr_model_comparison.py
-----------------------
Runs multiple vision LLMs on the same war diary pages and outputs
labelled plain text for side-by-side accuracy comparison.

One output file per model, named:
  {diary_stem}_comparison_{model_label}.md

Each file uses the same ## Page N format as the main OCR pipeline
so you can compare directly against the olmOCR output.

Requirements:
    pip install requests pillow python-dotenv

Poppler must be on PATH.
API key read from DEEPINFRA_API_KEY in .env file.

Usage:
    python scripts/ocr_model_comparison.py ^
      --single "G:/path/to/diary.pdf" ^
      --output_dir "G:/path/to/outputs_step1_olmocr" ^
      --first_page 7 ^
      --last_page 16
"""

import argparse
import base64
import os
import re
import subprocess
import sys
import tempfile
import time
from pathlib import Path

from dotenv import load_dotenv
import requests

load_dotenv()

# ---------------------------------------------------------------------------
# Models to test
# ---------------------------------------------------------------------------
# Add or remove models here. label is used in the output filename.
# All use the same DeepInfra OpenAI-compatible endpoint.

MODELS = [
    {
        "label": "gemini-2.5-pro",
        "model": "google/gemini-2.5-pro",
    },
    {
        "label": "qwen3-vl-30b",
        "model": "Qwen/Qwen3-VL-30B-A3B-Instruct",
    },
    {
        "label": "olmocr-2-7b",
        "model": "allenai/olmOCR-2-7B-1025",
    },
]



# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

DEEPINFRA_API_URL = "https://api.deepinfra.com/v1/openai/chat/completions"
DPI               = 150
MAX_RETRIES       = 3
RETRY_DELAY       = 15

# ---------------------------------------------------------------------------
# Prompt — same for all models so comparison is fair
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = """\
You are transcribing a scanned WWII Canadian Army war diary page.
Transcribe every word exactly as written. Do not correct spelling,
grammar, or abbreviations. Do not add commentary or explanations.
Keep grid references exactly as written (e.g. MR 2553, GR 442891).
Keep military abbreviations exactly as written (e.g. Bn, Bde, HQ, OR, C.O.).
Transcribe all text visible on the page — headers, titles, form numbers,
stamps, handwritten notes, signatures, and margin notations. Nothing omitted.
If the page contains a table with Place / Date / Hour / Summary columns,
output it as an HTML table using <table>, <tr>, <th>, <td>, <br> tags,
preserving column structure exactly.
If the page is genuinely blank, output only: [BLANK]\
"""

USER_PROMPT = (
    "Transcribe every word on this page exactly as written. "
    "Use an HTML table for any Place/Date/Hour/Summary diary table. "
    "Nothing omitted, nothing added."
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def pdf_to_images(pdf_path: Path, output_dir: Path, dpi: int = DPI,
                  first_page: int = None, last_page: int = None) -> list[tuple[int, Path]]:
    prefix = output_dir / pdf_path.stem
    cmd = ["pdftoppm", "-r", str(dpi), "-png"]
    if first_page is not None:
        cmd += ["-f", str(first_page)]
    if last_page is not None:
        cmd += ["-l", str(last_page)]
    cmd += [str(pdf_path), str(prefix)]

    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"pdftoppm failed:\n{result.stderr}")

    images = sorted(output_dir.glob(f"{pdf_path.stem}-*.png"))
    result_pairs = []
    for img_path in images:
        m = re.search(r'-(\d+)\.png$', img_path.name)
        page_num = int(m.group(1)) if m else len(result_pairs) + (first_page or 1)
        result_pairs.append((page_num, img_path))
    return result_pairs


def image_to_base64(image_path: Path) -> str:
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def ocr_page(image_path: Path, api_key: str, model: str, page_num: int) -> str:
    b64 = image_to_base64(image_path)

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
    }

    payload = {
        "model": model,
        "messages": [
            {
                "role": "system",
                "content": SYSTEM_PROMPT,
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/png;base64,{b64}"
                        },
                    },
                    {
                        "type": "text",
                        "text": USER_PROMPT,
                    },
                ],
            }
        ],
        "max_tokens": 4096,
        "temperature": 0.0,
    }

    for attempt in range(1, MAX_RETRIES + 1):
        try:
            response = requests.post(
                DEEPINFRA_API_URL,
                headers=headers,
                json=payload,
                timeout=180,
            )
            response.raise_for_status()
            data = response.json()
            return data["choices"][0]["message"]["content"].strip()

        except requests.exceptions.HTTPError as e:
            print(f"      HTTP error attempt {attempt}/{MAX_RETRIES}: {e}")
            if attempt < MAX_RETRIES:
                time.sleep(RETRY_DELAY)
            else:
                return f"[OCR FAILED — page {page_num} — {model} — HTTP error: {e}]"
        except Exception as e:
            print(f"      Error attempt {attempt}/{MAX_RETRIES}: {e}")
            if attempt < MAX_RETRIES:
                time.sleep(RETRY_DELAY)
            else:
                return f"[OCR FAILED — page {page_num} — {model} — Error: {e}]"


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def run_comparison(pdf_path: Path, output_dir: Path, api_key: str,
                   first_page: int = None, last_page: int = None):

    print(f"\n{'='*60}")
    print(f"OCR Model Comparison: {pdf_path.name}")
    if first_page or last_page:
        print(f"Pages: {first_page or 1} → {last_page or 'end'}")
    print(f"Models: {', '.join(m['label'] for m in MODELS)}")
    print(f"{'='*60}")

    with tempfile.TemporaryDirectory() as tmp_dir:
        tmp_path = Path(tmp_dir)

        print(f"\nRendering PDF pages at {DPI} DPI...")
        page_images = pdf_to_images(pdf_path, tmp_path,
                                    first_page=first_page, last_page=last_page)
        total = len(page_images)
        print(f"Found {total} pages.")

        # Store results per model: {label: [(page_num, text), ...]}
        results = {m["label"]: [] for m in MODELS}

        for i, (pdf_page_num, image_path) in enumerate(page_images, start=1):
            print(f"\nPage {pdf_page_num} ({i}/{total})")

            for model_cfg in MODELS:
                label = model_cfg["label"]
                model = model_cfg["model"]
                print(f"  [{label}]...", end=" ", flush=True)

                page_text = ocr_page(image_path, api_key, model, pdf_page_num)
                results[label].append((pdf_page_num, page_text))
                print("done")

                # Delay between models to avoid rate limits
                time.sleep(1)

            # Slightly longer delay between pages
            if i < total:
                time.sleep(1)

    # Write one output file per model
    output_dir.mkdir(parents=True, exist_ok=True)
    stem = pdf_path.stem

    for model_cfg in MODELS:
        label = model_cfg["label"]
        model_id = model_cfg["model"]
        output_path = output_dir / f"{stem}_comparison_{label}.md"

        lines = [
            f"# {stem}",
            f"Model: {model_id}",
            f"Label: {label}",
            f"Pages: {first_page or 1}–{last_page or total}",
            f"DPI: {DPI}",
            "",
        ]

        for page_num, page_text in results[label]:
            lines.append(f"## Page {page_num} — {label}")
            lines.append("")
            lines.append(page_text)
            lines.append("")

        output_path.write_text("\n".join(lines), encoding="utf-8")
        print(f"\nSaved: {output_path.name}")

    print(f"\n{'='*60}")
    print(f"Complete. {len(MODELS)} comparison files written to:")
    print(f"  {output_dir}")
    print("\nFiles:")
    for model_cfg in MODELS:
        print(f"  {stem}_comparison_{model_cfg['label']}.md")
    print(f"\nCompare these against:")
    print(f"  {stem}_olmocr.md  (existing olmOCR output)")


def main():
    parser = argparse.ArgumentParser(
        description="Compare multiple vision LLMs on the same war diary pages."
    )
    parser.add_argument(
        "--single",
        required=True,
        help="Path to a single PDF file to process.",
    )
    parser.add_argument(
        "--output_dir",
        required=True,
        help="Folder to write comparison output files.",
    )
    parser.add_argument(
        "--first_page",
        type=int,
        default=None,
        help="First page to process (1-based).",
    )
    parser.add_argument(
        "--last_page",
        type=int,
        default=None,
        help="Last page to process (1-based).",
    )
    args = parser.parse_args()

    api_key = os.getenv("DEEPINFRA_API_KEY")
    if not api_key:
        print("Error: DEEPINFRA_API_KEY not set in .env file")
        sys.exit(1)

    pdf_path = Path(args.single)
    if not pdf_path.exists():
        print(f"Error: PDF not found: {pdf_path}")
        sys.exit(1)

    run_comparison(
        pdf_path,
        Path(args.output_dir),
        api_key,
        first_page=args.first_page,
        last_page=args.last_page,
    )


if __name__ == "__main__":
    main()
