"""
ocr_model_comparison_anthropic.py
----------------------------------
Runs Claude Sonnet 4.6 via the Anthropic API on war diary pages.
Produces the same ## Page N — label output format as ocr_model_comparison.py
so results can be directly compared against Gemini Pro, olmOCR, etc.

Output file:
  {diary_stem}_comparison_claude-sonnet-4-6.md

Requirements:
    pip install requests pillow python-dotenv

Poppler must be on PATH.
ANTHROPIC_API_KEY read from .env file.

Usage:
    python scripts/ocr_model_comparison_anthropic.py ^
      --single "G:/IdeaProjects/AI-Prototype/Inputs/Calgary-Highlanders_War-Diary_Sep44.pdf" ^
      --output_dir "G:/IdeaProjects/AI-Prototype/outputs_step1_olmocr" ^
      --first_page 7 ^
      --last_page 32
"""

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()

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

ANTHROPIC_API_URL = "https://api.anthropic.com/v1/messages"
MODEL_ID          = "claude-sonnet-4-6"
MODEL_LABEL       = "claude-sonnet-4-6"
DPI               = 150
MAX_RETRIES       = 3
RETRY_DELAY       = 15

# ── Prompt — identical to ocr_model_comparison.py for fair comparison ─────────

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:
    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:
    """
    Load PNG, compress to JPEG at 85% quality, return base64.
    Anthropic has a 5MB image limit — war diary PNGs at 150 DPI
    are ~6-7MB. JPEG at 85% brings them to ~0.8-1.2MB with no
    meaningful loss of OCR-relevant detail.
    """
    from PIL import Image
    import io
    img = Image.open(image_path).convert("RGB")
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=85)
    return base64.b64encode(buf.getvalue()).decode("utf-8")


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

    headers = {
        "x-api-key": api_key,
        "anthropic-version": "2023-06-01",
        "content-type": "application/json",
    }

    payload = {
        "model": MODEL_ID,
        "max_tokens": 8192,
        "system": SYSTEM_PROMPT,
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "source": {
                            "type": "base64",
                            "media_type": "image/jpeg",
                            "data": b64,
                        },
                    },
                    {
                        "type": "text",
                        "text": USER_PROMPT,
                    },
                ],
            }
        ],
    }

    for attempt in range(1, MAX_RETRIES + 1):
        try:
            response = requests.post(
                ANTHROPIC_API_URL,
                headers=headers,
                json=payload,
                timeout=180,
            )
            response.raise_for_status()
            data = response.json()

            # Print token usage
            usage = data.get("usage", {})
            input_t  = usage.get("input_tokens", 0)
            output_t = usage.get("output_tokens", 0)
            cost = (input_t / 1_000_000 * 3.00) + (output_t / 1_000_000 * 15.00)
            print(f"      tokens: {input_t} in / {output_t} out — ${cost:.4f}")

            return data["content"][0]["text"].strip()

        except requests.exceptions.HTTPError as e:
            print(f"      HTTP error attempt {attempt}/{MAX_RETRIES}: {e}")
            if hasattr(e, 'response') and e.response is not None:
                print(f"      Response: {e.response.text[:300]}")
            if attempt < MAX_RETRIES:
                time.sleep(RETRY_DELAY)
            else:
                return f"[OCR FAILED — page {page_num} — {MODEL_ID} — 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_ID} — 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_LABEL}")
    print(f"File: {pdf_path.name}")
    print(f"Pages: {first_page or 1} → {last_page or 'end'}")
    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.")

        results = []

        for i, (pdf_page_num, image_path) in enumerate(page_images, start=1):
            print(f"\nPage {pdf_page_num} ({i}/{total})...", end=" ", flush=True)
            page_text = ocr_page(image_path, api_key, pdf_page_num)
            results.append((pdf_page_num, page_text))
            print("done")

            if i < total:
                time.sleep(1)

    # Write output — same format as ocr_model_comparison.py
    output_dir.mkdir(parents=True, exist_ok=True)
    stem = pdf_path.stem
    output_path = output_dir / f"{stem}_comparison_{MODEL_LABEL}.md"

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

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

    output_path.write_text("\n".join(lines), encoding="utf-8")

    print(f"\n{'='*60}")
    print(f"Done. Output saved to:")
    print(f"  {output_path.name}")
    print(f"{'='*60}")


def main():
    parser = argparse.ArgumentParser(
        description="OCR war diary pages using Claude Sonnet 4.6 via Anthropic API."
    )
    parser.add_argument("--single",      required=True,  help="Path to PDF.")
    parser.add_argument("--output_dir",  required=True,  help="Output folder.")
    parser.add_argument("--first_page",  type=int, default=None)
    parser.add_argument("--last_page",   type=int, default=None)
    args = parser.parse_args()

    api_key = os.getenv("ANTHROPIC_API_KEY")
    if not api_key:
        print("Error: ANTHROPIC_API_KEY not set in .env")
        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()