"""
step1_ocr_wardiaries_claude-4-6.py
------------------------------------
Converts war diary PDFs to structured text using Claude Sonnet 4.6 via the Anthropic API.
Drop-in replacement for step1_ocr_wardiaries_olmocr.py — same output format, same CLI args.

Output format: one .md file per PDF with ## Page N markers.
- Diary pages with Place/Date/Hour/Summary tables are output as HTML tables,
  preserving column structure for downstream extraction.
- All other pages (admin, appendices, forms, etc.) are output as plain text.

This structured output is the source of truth for the pipeline:
  - llm_extract.py reads the HTML tables for position and narrative extraction
  - step2_json_to_viewer_md.py flattens the HTML tables to readable prose for the OCR viewer

Requirements:
    pip install requests pillow python-dotenv

Poppler must be on PATH (pdftoppm command must work).

API key is read from ANTHROPIC_API_KEY in your .env file.

Usage:
    # Test run — pages 7 to 27 only:
    python scripts/step1_ocr_wardiaries_claude-4-6.py --single "G:/path/to/diary.pdf" --output_dir "G:/path/to/outputs" --first_page 7 --last_page 27

    # Full diary:
    python scripts/step1_ocr_wardiaries_claude-4-6.py --single "G:/path/to/diary.pdf" --output_dir "G:/path/to/outputs"

    # All PDFs in a folder:
    python scripts/step1_ocr_wardiaries_claude-4-6.py --input_dir "G:/path/to/pdfs" --output_dir "G:/path/to/outputs"
"""

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

from dotenv import load_dotenv
from PIL import Image
import requests

load_dotenv()

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

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

# ---------------------------------------------------------------------------
# OCR prompt — identical to step1_ocr_wardiaries_olmocr.py
# ---------------------------------------------------------------------------

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, MR 0675 Sheet 861).
Keep military abbreviations exactly as written (e.g. Bn, Bde, HQ, OR, C.O., RHC, R de Mais).
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 columns Place / Date / Hour / Summary / Remarks:
  Output the table as HTML using <table>, <tr>, <th>, <td> tags.
  - Place column: all lines exactly as written, separated by <br> tags.
    Include country names, place names, grid references, and sheet references.
  - Date column: exact date as written, <br> if split across lines.
  - Hour column: exact time as written, or empty if absent.
  - Summary column: full verbatim text, paragraphs separated by <br>. Every word. Nothing omitted.
  - Remarks column: exact text, or empty if absent.
  If the page has multiple rows (different dates or continuation entries), each row is a separate <tr>.
  Any text outside the table (headers, weather notes, stamps) is output as plain text
  before or after the <table> tag in the order it appears on the page.

If the page does NOT contain a Place/Date/Hour/Summary table:
  Output all text as plain text in the order it appears on the page.

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]]:
    """
    Render pages of a PDF to PNG images using pdftoppm (poppler).
    Returns a list of (pdf_page_number, image_path) tuples.
    """
    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 for {pdf_path.name}:\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_jpeg(image_path: Path, quality: int = 85) -> str:
    """
    Load PNG, compress to JPEG at given quality, return base64.
    Anthropic has a 5 MB image limit — war diary PNGs at 150 DPI are ~6-7 MB.
    JPEG at 85% brings them to ~0.8-1.2 MB with no meaningful loss of OCR detail.
    """
    img = Image.open(image_path).convert("RGB")
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=quality)
    return base64.b64encode(buf.getvalue()).decode("utf-8")


def ocr_page(image_path: Path, api_key: str, page_num: int) -> str:
    """Send one page image to Claude via Anthropic API and return the text."""
    b64 = image_to_base64_jpeg(image_path)

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

    payload = {
        "model": MODEL,
        "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 and estimated cost
            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}", end=" ")

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

        except requests.exceptions.HTTPError as e:
            print(f"    HTTP error on page {page_num}, 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} — HTTP error: {e}]"
        except Exception as e:
            print(f"    Error on page {page_num}, attempt {attempt}/{MAX_RETRIES}: {e}")
            if attempt < MAX_RETRIES:
                time.sleep(RETRY_DELAY)
            else:
                return f"[OCR FAILED — page {page_num} — Error: {e}]"


def has_table(page_text: str) -> bool:
    """Return True if the page output contains an HTML table."""
    return bool(re.search(r'<table', page_text, re.IGNORECASE))


def ocr_pdf(pdf_path: Path, output_dir: Path, api_key: str,
            first_page: int = None, last_page: int = None) -> Path:
    """
    OCR a single PDF and write structured output to a .md file.
    Diary pages are output as HTML tables; all other pages as plain text.
    """
    print(f"\n{'='*60}")
    print(f"Processing: {pdf_path.name}")
    if first_page or last_page:
        print(f"Pages: {first_page or 1} → {last_page or 'end'}")
    print(f"{'='*60}")

    output_md = output_dir / f"step1_{pdf_path.stem}_claude-4-6.md"

    if output_md.exists():
        print(f"  Already processed — skipping.")
        print(f"  Delete {output_md.name} to reprocess.")
        return output_md

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

        print(f"  Rendering PDF pages to images at {DPI} DPI...")
        try:
            page_images = pdf_to_images(pdf_path, tmp_path,
                                        first_page=first_page, last_page=last_page)
        except RuntimeError as e:
            print(f"  ERROR: {e}")
            return output_md

        total = len(page_images)
        print(f"  Found {total} pages to process.")

        lines = [
            f"# {pdf_path.stem}",
            f"OCR by Claude Sonnet 4.6 ({MODEL}) via Anthropic",
            f"Source: {pdf_path.name} — pages {first_page or 1}–{last_page or total}",
            "",
        ]

        diary_count = 0
        other_count = 0

        for i, (pdf_page_num, image_path) in enumerate(page_images, start=1):
            print(f"  Page {pdf_page_num} ({i}/{total})...", end=" ", flush=True)

            page_text = ocr_page(image_path, api_key, pdf_page_num)
            is_diary = has_table(page_text)

            lines.append(f"## Page {pdf_page_num}")
            lines.append("")
            lines.append(page_text)
            lines.append("")

            if is_diary:
                diary_count += 1
                print("done — diary (table)")
            else:
                other_count += 1
                print("done — plain text")

            if i < total:
                time.sleep(0.5)

    output_dir.mkdir(parents=True, exist_ok=True)
    output_md.write_text("\n".join(lines), encoding="utf-8")

    print(f"\n  Saved: {output_md}")
    print(f"  Summary: {diary_count} diary pages (HTML tables), {other_count} other pages (plain text)")

    return output_md


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

def main():
    parser = argparse.ArgumentParser(
        description="OCR war diary PDFs to structured text using Claude Sonnet 4.6 via Anthropic API."
    )
    parser.add_argument(
        "--input_dir",
        default=None,
        help="Folder containing PDF files to process.",
    )
    parser.add_argument(
        "--output_dir",
        required=True,
        help="Folder to write output .md files.",
    )
    parser.add_argument(
        "--single",
        default=None,
        help="Process a single PDF file instead of a whole folder.",
    )
    parser.add_argument(
        "--first_page",
        type=int,
        default=None,
        help="First PDF page to process (1-based).",
    )
    parser.add_argument(
        "--last_page",
        type=int,
        default=None,
        help="Last PDF page to process (1-based).",
    )
    args = parser.parse_args()

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

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    if args.single:
        pdfs = [Path(args.single)]
    elif args.input_dir:
        pdfs = sorted(Path(args.input_dir).glob("*.pdf"))
    else:
        print("Error: provide --input_dir or --single")
        sys.exit(1)

    if not pdfs:
        print("No PDF files found.")
        sys.exit(1)

    print(f"Found {len(pdfs)} PDF(s) to process.")
    print(f"Output directory: {output_dir}")
    print(f"Model: {MODEL}")
    print(f"DPI: {DPI}")

    results = []
    for pdf_path in pdfs:
        output_md = ocr_pdf(
            pdf_path, output_dir, api_key,
            first_page=args.first_page,
            last_page=args.last_page,
        )
        results.append(output_md)

    print(f"\n{'='*60}")
    print(f"Complete. {len(results)} file(s) processed.")
    print("Output files:")
    for r in results:
        print(f"  {r}")


if __name__ == "__main__":
    main()
