"""
ocr_confidence.py
-----------------
Reads an existing olmOCR .md output file and scores each page for transcription
confidence using the DeepInfra API. Does NOT re-run OCR.

Requirements:
    pip install requests

Usage:
    python ocr_confidence.py --api_key "YOUR_KEY"
    python ocr_confidence.py --api_key "YOUR_KEY" --input "path/to/file.md" --output "path/to/out.json"
"""

import argparse
import json
import sys
import time
import warnings
from pathlib import Path

import requests

# ---------------------------------------------------------------------------
# Load .env from project root (if present) — no external dependencies needed
# ---------------------------------------------------------------------------

def _load_dotenv() -> None:
    env_path = Path(__file__).resolve().parent.parent / ".env"
    if not env_path.exists():
        return
    import os
    with open(env_path, encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            key, _, value = line.partition("=")
            os.environ.setdefault(key.strip(), value.strip())

_load_dotenv()

# ---------------------------------------------------------------------------
# Configuration — mirrors ocr_wardiaries.py
# ---------------------------------------------------------------------------

DEEPINFRA_API_URL = "https://api.deepinfra.com/v1/openai/chat/completions"
# olmOCR is a vision model — use a chat LLM for text-based confidence review
MODEL             = "google/gemma-3-27b-it"
MAX_RETRIES       = 3
RETRY_DELAY       = 5   # seconds between retries

_PROJECT_ROOT = Path(__file__).resolve().parent.parent
DEFAULT_INPUT  = _PROJECT_ROOT / "Inputs" / "ocr-output" / "Calgary-Highlanders_War-Diary_Sep44_olmocr.md"
DEFAULT_OUTPUT = _PROJECT_ROOT / "Inputs" / "ocr-output" / "Calgary-Highlanders_War-Diary_Sep44_confidence.json"

CONFIDENCE_PROMPT = (
    "You are reviewing an OCR transcription of a WWII war diary. Your job is to identify words that may have been misread by the OCR scanner — not to correct the soldier's original spelling or interpret abbreviations. Rate OCR accuracy only. Do not suggest what words "should" be. Respond in JSON only:
     {"score": 7, "uncertain_words": ["Loon", "Fme"], "notes": "Possible OCR misread in line 3"}
)

ERROR_RESULT = {"score": 0, "uncertain_words": [], "notes": "API error"}


# ---------------------------------------------------------------------------
# Parse olmOCR markdown into {page_num: text} dict
# Same logic as parseOCRByPage() in p44-ocr-viewer.html
# ---------------------------------------------------------------------------

def parse_ocr_by_page(raw_text: str) -> dict[int, str]:
    """Split olmOCR markdown on '## Page N' headings into a page-keyed dict."""
    pages: dict[int, str] = {}
    lines = raw_text.split("\n")
    page_num: int | None = None
    buf: list[str] = []

    for line in lines:
        m_head = _PAGE_HEADING.match(line)
        if m_head:
            if page_num is not None:
                pages[page_num] = "\n".join(buf)
            page_num = int(m_head.group(1))
            buf = [line]          # keep the heading at the top
        elif page_num is not None:
            buf.append(line)

    if page_num is not None:
        pages[page_num] = "\n".join(buf)

    return pages


import re as _re
_PAGE_HEADING = _re.compile(r"^## Page (\d+)\s*$")


# ---------------------------------------------------------------------------
# API call
# ---------------------------------------------------------------------------

def score_page(page_text: str, page_num: int, api_key: str) -> dict:
    """
    Send one page's OCR text to DeepInfra for confidence scoring.
    Returns a dict with keys: score, uncertain_words, notes.
    On failure returns ERROR_RESULT.
    """
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
    }
    payload = {
        "model": MODEL,
        "messages": [
            {
                "role": "user",
                "content": (
                    f"{CONFIDENCE_PROMPT}\n\n"
                    f"--- BEGIN OCR TEXT (page {page_num}) ---\n"
                    f"{page_text}\n"
                    f"--- END OCR TEXT ---"
                ),
            }
        ],
        "max_tokens": 256,
        "temperature": 0.0,
    }

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

            # Strip markdown code fences if the model wrapped the JSON
            if raw_content.startswith("```"):
                raw_content = _re.sub(r"^```[a-z]*\n?", "", raw_content)
                raw_content = _re.sub(r"\n?```$", "", raw_content)

            result = json.loads(raw_content)

            # Validate expected keys; fill missing ones with defaults
            score = int(result.get("score", 0))
            uncertain_words = result.get("uncertain_words", [])
            notes = result.get("notes", "")

            if not isinstance(uncertain_words, list):
                uncertain_words = []

            return {"score": score, "uncertain_words": uncertain_words, "notes": notes}

        except requests.exceptions.HTTPError as exc:
            warnings.warn(f"Page {page_num}: HTTP error on attempt {attempt}/{MAX_RETRIES}: {exc}")
            if attempt < MAX_RETRIES:
                time.sleep(RETRY_DELAY)

        except json.JSONDecodeError as exc:
            warnings.warn(
                f"Page {page_num}: Malformed JSON from API on attempt {attempt}/{MAX_RETRIES}: {exc}"
            )
            if attempt < MAX_RETRIES:
                time.sleep(RETRY_DELAY)

        except Exception as exc:  # noqa: BLE001
            warnings.warn(f"Page {page_num}: Unexpected error on attempt {attempt}/{MAX_RETRIES}: {exc}")
            if attempt < MAX_RETRIES:
                time.sleep(RETRY_DELAY)

    warnings.warn(f"Page {page_num}: All {MAX_RETRIES} attempts failed — storing error result.")
    return dict(ERROR_RESULT)  # return a fresh copy


# ---------------------------------------------------------------------------
# Summary helpers
# ---------------------------------------------------------------------------

def print_summary(results: dict) -> None:
    high = sum(1 for v in results.values() if v["score"] >= 8)
    med  = sum(1 for v in results.values() if 5 <= v["score"] <= 7)
    low  = sum(1 for v in results.values() if 1 <= v["score"] <= 4)
    err  = sum(1 for v in results.values() if v["score"] == 0)

    total = len(results)
    print("\n" + "=" * 50)
    print(f"CONFIDENCE SCORING COMPLETE — {total} pages scored")
    print("=" * 50)
    print(f"  High  (8-10) : {high:>4}  ({high/total*100:.1f}%)" if total else "  High  (8-10) :    0")
    print(f"  Medium (5-7) : {med:>4}  ({med/total*100:.1f}%)"  if total else "  Medium (5-7) :    0")
    print(f"  Low   (1-4)  : {low:>4}  ({low/total*100:.1f}%)"  if total else "  Low   (1-4)  :    0")
    if err:
        print(f"  API errors   : {err:>4}")
    print("=" * 50)


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

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Score OCR confidence for each page of an olmOCR .md output file."
    )
    parser.add_argument(
        "--api_key", default=None,
        help="DeepInfra API key (defaults to DEEPINFRA_API_KEY env var)."
    )
    parser.add_argument(
        "--input", default=str(DEFAULT_INPUT),
        help=f"Path to olmOCR .md file. Default: {DEFAULT_INPUT}"
    )
    parser.add_argument(
        "--output", default=str(DEFAULT_OUTPUT),
        help=f"Path to write confidence JSON. Default: {DEFAULT_OUTPUT}"
    )
    parser.add_argument(
        "--delay", type=float, default=0.5,
        help="Seconds to pause between API calls (default: 0.5)."
    )
    args = parser.parse_args()

    import os
    api_key = args.api_key or os.environ.get("DEEPINFRA_API_KEY", "")
    if not api_key:
        print("ERROR: No API key provided. Use --api_key or set DEEPINFRA_API_KEY in .env", file=sys.stderr)
        sys.exit(1)

    input_path  = Path(args.input)
    output_path = Path(args.output)

    # ── Read source file ──────────────────────────────────────────────────────
    if not input_path.exists():
        print(f"ERROR: Input file not found: {input_path}", file=sys.stderr)
        sys.exit(1)

    print(f"Reading OCR source: {input_path}")
    raw_text = input_path.read_text(encoding="utf-8")

    pages = parse_ocr_by_page(raw_text)
    if not pages:
        print("ERROR: No '## Page N' headings found in the input file.", file=sys.stderr)
        sys.exit(1)

    total_pages   = len(pages)
    sorted_pages  = sorted(pages.keys())
    print(f"Found {total_pages} pages (Page {sorted_pages[0]} – {sorted_pages[-1]}).")

    # ── Load existing results (resume support) ────────────────────────────────
    results: dict[str, dict] = {}
    if output_path.exists():
        try:
            results = json.loads(output_path.read_text(encoding="utf-8"))
            already_done = len(results)
            print(f"Resuming — {already_done} page(s) already scored, skipping them.")
        except (json.JSONDecodeError, OSError) as exc:
            warnings.warn(f"Could not read existing output ({exc}); starting fresh.")
            results = {}

    output_path.parent.mkdir(parents=True, exist_ok=True)

    # ── Score each page ───────────────────────────────────────────────────────
    scored_this_run = 0
    for page_num in sorted_pages:
        page_key = str(page_num)

        if page_key in results:
            continue    # already scored — skip

        page_text = pages[page_num]
        result    = score_page(page_text, page_num, api_key)
        results[page_key] = result
        scored_this_run += 1

        # Progress line
        score_str = str(result["score"]) if result["score"] > 0 else "ERR"
        print(f"  Page {page_num}/{sorted_pages[-1]} — score: {score_str}")

        # Write after every page so a crash loses minimal work
        output_path.write_text(
            json.dumps(results, indent=2, ensure_ascii=False),
            encoding="utf-8"
        )

        if scored_this_run > 0 and page_num != sorted_pages[-1]:
            time.sleep(args.delay)

    # ── Final save & summary ──────────────────────────────────────────────────
    output_path.write_text(
        json.dumps(results, indent=2, ensure_ascii=False),
        encoding="utf-8"
    )
    print(f"\nResults written to: {output_path}")
    print_summary(results)


if __name__ == "__main__":
    main()

