"""
ocr_correction_test.py
----------------------
Tests Gemini Pro as a text-only correction pass over olmOCR output.
No images sent — text in, structured corrections out.

Reads:
  - olmOCR output .md file (already generated)
  - CP Stacey Victory Campaign PDF (extracts relevant text)

Sends to Gemini Pro via DeepInfra:
  - Stacey context (unit + date relevant passages)
  - Standard abbreviations/terminology
  - olmOCR page text
  - Correction prompt

Outputs:
  - JSON file with suggested corrections per page
  - Summary markdown showing before/after for each suggestion

Usage:
    python scripts/ocr_correction_test.py

Requirements:
    pip install requests python-dotenv pypdf pdfplumber
"""

import json
import os
import re
import sys
from pathlib import Path

import pdfplumber
import requests
from dotenv import load_dotenv

load_dotenv()

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

PROJECT_ROOT = Path("G:/IdeaProjects/AI-Prototype")

OLMOCR_OUTPUT = PROJECT_ROOT / "outputs_step1_olmocr" / "Calgary-Highlanders_War-Diary_Sep44_control-test_comparison_olmocr-2-7b.md"
STACEY_PDF = PROJECT_ROOT / "Inputs" / "CP-Stacey_Official-History_The-Victory-Campaign.pdf"
OUTPUT_DIR    = PROJECT_ROOT / "outputs_step1_olmocr"

DEEPINFRA_API_KEY = os.getenv("DEEPINFRA_API_KEY")
API_URL           = "https://api.deepinfra.com/v1/openai/chat/completions"
MODEL             = "google/gemini-2.5-pro"

# Pages to test — the ones with known olmOCR errors from comparison
TEST_PAGES = [4, 5, 6, 7, 30, 32]

# ── Stacey extraction ─────────────────────────────────────────────────────────

# Keywords that indicate relevance to Calgary Highlanders Sep 44 operations
STACEY_KEYWORDS = [
    "Calgary Highlanders",
    "Calg High",
    "5th Canadian Infantry Brigade",
    "5 CIB",
    "2nd Canadian Infantry Division",
    "Loon Plage",
    "Bourbourgville",
    "Brecht",
    "Scheldt",
    "Crerar",
    "MacLauchlan",
    "Dieppe",
    "September 1944",
    "Sep 44",
]

MAX_STACEY_CHARS = 12_000  # ~3,000 tokens — enough context without blowing cost


def extract_stacey_context(pdf_path: Path, keywords: list[str], max_chars: int) -> str:
    """
    Extract relevant passages from Stacey by keyword matching.
    Returns a condensed string of the most relevant paragraphs.
    """
    if not pdf_path.exists():
        print(f"WARNING: Stacey PDF not found at {pdf_path}")
        return "[Stacey context not available — PDF not found]"

    print(f"Extracting Stacey context from {pdf_path.name}...")
    relevant_chunks = []

    try:
        with pdfplumber.open(pdf_path) as pdf:
            print(f"  Stacey: {len(pdf.pages)} pages")
            for i, page in enumerate(pdf.pages):
                text = page.extract_text()
                if not text:
                    continue
                # Check if any keyword appears on this page
                text_lower = text.lower()
                if any(kw.lower() in text_lower for kw in keywords):
                    relevant_chunks.append(f"[Stacey p.{i+1}]\n{text.strip()}")

    except Exception as e:
        print(f"  WARNING: Error reading Stacey PDF: {e}")
        return "[Stacey context extraction failed]"

    if not relevant_chunks:
        print("  WARNING: No relevant Stacey passages found")
        return "[No relevant Stacey passages found for this unit/period]"

    combined = "\n\n".join(relevant_chunks)

    # Truncate if needed
    if len(combined) > max_chars:
        combined = combined[:max_chars] + "\n\n[...Stacey context truncated for length...]"

    print(f"  Found {len(relevant_chunks)} relevant Stacey pages, {len(combined)} chars")
    return combined


# ── olmOCR page extraction ────────────────────────────────────────────────────

def extract_olmocr_pages(md_path: Path, page_numbers: list[int]) -> dict[int, str]:
    """
    Parse the olmOCR markdown output and extract specific page texts.
    Returns dict of {page_number: page_text}.
    """
    if not md_path.exists():
        print(f"ERROR: olmOCR output not found at {md_path}")
        sys.exit(1)

    content = md_path.read_text(encoding="utf-8")
    pages = {}

    for page_num in page_numbers:
        # Match ## Page N — olmocr-2-7b header
        pattern = rf"## Page {page_num} — olmocr-2-7b\n(.*?)(?=\n## Page |\Z)"
        match = re.search(pattern, content, re.DOTALL)
        if match:
            pages[page_num] = match.group(1).strip()
        else:
            print(f"  WARNING: Page {page_num} not found in olmOCR output")

    return pages


# ── Reference context ─────────────────────────────────────────────────────────

ABBREVIATIONS_CONTEXT = """
STANDARD WWII CANADIAN ARMY ABBREVIATIONS AND TERMINOLOGY:

Unit abbreviations:
- Bn = Battalion
- Coy = Company
- Bde = Brigade
- Div = Division
- RHC = Royal Hamilton Light Infantry (Rileys)
- R de Mais / R de M = Regiment de Maisonneuve
- Calg Highrs / Calgary H = Calgary Highlanders
- 5 CIB = 5th Canadian Infantry Brigade
- 2 CID / 2 CDN INF DIV = 2nd Canadian Infantry Division
- RCA = Royal Canadian Artillery
- RCASC = Royal Canadian Army Service Corps
- RCEME = Royal Canadian Electrical and Mechanical Engineers

Rank abbreviations:
- Lt.-Col. / Lieut.-Col. = Lieutenant Colonel
- Maj. / Major = Major
- Capt. = Captain
- Lieut. / Lt. = Lieutenant
- Cpl. = Corporal
- Pte. = Private
- C.O. = Commanding Officer
- I.O. = Intelligence Officer
- Adjt. / Adj. = Adjutant
- B.M. = Brigade Major
- O.C. = Officer Commanding

Military terms:
- Bn HQ = Battalion Headquarters
- "O" Group = Orders Group (briefing)
- recce = reconnaissance
- embussed = loaded into vehicles
- debussed = unloaded from vehicles
- MG = Machine Gun
- pdrs = pounders (artillery shell weight, e.g. 25 pdrs = 25-pounder guns)
- MR = Map Reference
- Grid refs are 6-digit numbers referencing military map sheets
- NTR = Nothing To Report
- DF = Defensive Fire (artillery task)
- SOS = Special On-call (artillery task, highest priority)

Key personnel Sep 44:
- Lt.-Col. D.G. MacLauchlan = CO, Calgary Highlanders
- General H.D.G. Crerar = Commander, First Canadian Army (NOT Montgomery)
- Brigadier W.J. Megill = Commander, 5th Canadian Infantry Brigade
- Major R.G. Slater = Brigade Major, 5 CIB
- R.L. Ellis = Major, Calgary Highlanders
- Major Kearns = OC Able Company, Calgary Highlanders

Key locations Sep 44 (Calgary Highlanders axis):
- Dieppe, France (2-3 Sep 44) — liberation, memorial service
- Loon Plage, France (7-8 Sep 44) — NOT "Loon Place", NOT "Caen Place"
- Bourbourgville, France (7 Sep 44) — canal crossing area
- Borsbeek, Belgium (16-18 Sep 44) — 5 CIB HQ
- Antwerp, Belgium (Sep 44)
- Brecht, Belgium (30 Sep 44) — Able Coy street fighting

Note: The march past at Dieppe on 3 Sep 44 was for General Crerar
(First Canadian Army), NOT General Montgomery.
"""


# ── Correction prompt ─────────────────────────────────────────────────────────

def build_correction_prompt(page_num: int, ocr_text: str, stacey_context: str) -> str:
    return f"""You are a WWII military history expert helping to correct OCR errors in Canadian Army war diary transcriptions.

You have been given:
1. A reference block of standard abbreviations and known facts about the Calgary Highlanders, September 1944
2. Relevant passages from C.P. Stacey's official history "The Victory Campaign" (Vol III of the Official History of the Canadian Army in the Second World War)
3. Raw OCR output from page {page_num} of the Calgary Highlanders War Diary, September 1944

Your task: identify OCR errors and suggest corrections. Focus on:
- Wrong names (people, places, units)
- Garbled military terms or abbreviations
- Repeated text loops (where the OCR got stuck repeating a phrase)
- Missing words or broken sentences
- Factual errors that contradict the reference context

Do NOT correct:
- Original spelling errors made by the diarist (these are historically important)
- Archaic spellings or period-appropriate language
- Abbreviations that are standard for the period

Return ONLY a JSON array. Each element must have exactly these fields:
  "page": {page_num},
  "original": "the exact text as it appears in the OCR output",
  "correction": "the corrected text",
  "confidence": "high" | "medium" | "low",
  "reason": "brief explanation of why this is an error"

If you find no errors, return an empty array: []
Do not return any text outside the JSON array.

═══════════════════════════════════════════════
REFERENCE: ABBREVIATIONS AND KNOWN FACTS
═══════════════════════════════════════════════
{ABBREVIATIONS_CONTEXT}

═══════════════════════════════════════════════
REFERENCE: C.P. STACEY — THE VICTORY CAMPAIGN
═══════════════════════════════════════════════
{stacey_context}

═══════════════════════════════════════════════
OCR TEXT TO CORRECT — PAGE {page_num}
═══════════════════════════════════════════════
{ocr_text}
"""


# ── API call ──────────────────────────────────────────────────────────────────

def call_gemini_correction(prompt: str, page_num: int) -> list[dict]:
    """
    Send correction request to Gemini Pro via DeepInfra.
    Returns list of correction dicts.
    """
    if not DEEPINFRA_API_KEY:
        print("ERROR: DEEPINFRA_API_KEY not set in .env")
        sys.exit(1)

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

    payload = {
        "model": MODEL,
        "max_tokens": 2000,
        "temperature": 0.1,  # Low temp — we want deterministic corrections
        "messages": [
            {
                "role": "user",
                "content": prompt,
            }
        ],
    }

    print(f"  Calling Gemini Pro for page {page_num}...")
    print(f"  Prompt length: {len(prompt):,} chars (~{len(prompt)//4:,} tokens)")

    try:
        response = requests.post(
            API_URL,
            headers=headers,
            json=payload,
            timeout=120,
        )
        response.raise_for_status()
    except requests.exceptions.Timeout:
        print(f"  ERROR: Timeout on page {page_num}")
        return []
    except requests.exceptions.RequestException as e:
        print(f"  ERROR: API call failed: {e}")
        return []

    data = response.json()

    # Extract content
    try:
        raw = data["choices"][0]["message"]["content"].strip()
    except (KeyError, IndexError) as e:
        print(f"  ERROR: Unexpected response structure: {e}")
        print(f"  Response: {data}")
        return []

    # Strip think blocks if present
    raw = re.sub(r"<think>.*?</think>", "", raw, flags=re.DOTALL).strip()

    # Strip markdown code fences if present
    raw = re.sub(r"^```(?:json)?\s*", "", raw)
    raw = re.sub(r"\s*```$", "", raw)
    raw = raw.strip()

    # Parse JSON
    try:
        corrections = json.loads(raw)
        if not isinstance(corrections, list):
            print(f"  WARNING: Expected JSON array, got {type(corrections)}")
            return []
        print(f"  Found {len(corrections)} correction(s)")
        return corrections
    except json.JSONDecodeError as e:
        print(f"  ERROR: JSON parse failed: {e}")
        print(f"  Raw response (first 500 chars): {raw[:500]}")
        return []


# ── Output formatting ─────────────────────────────────────────────────────────

def format_corrections_report(all_corrections: dict[int, list[dict]]) -> str:
    """
    Format all corrections into a readable markdown report.
    """
    lines = [
        "# OCR Correction Test Report",
        "## Gemini 2.5 Pro text correction with Stacey context",
        "### Calgary Highlanders War Diary Sep 44 — olmOCR output",
        "",
        f"Pages tested: {sorted(all_corrections.keys())}",
        "",
        "---",
        "",
    ]

    total_corrections = sum(len(v) for v in all_corrections.values())
    lines.append(f"**Total corrections suggested: {total_corrections}**")
    lines.append("")

    confidence_counts = {"high": 0, "medium": 0, "low": 0}

    for page_num in sorted(all_corrections.keys()):
        corrections = all_corrections[page_num]
        lines.append(f"## Page {page_num}")

        if not corrections:
            lines.append("*No corrections suggested.*")
            lines.append("")
            continue

        lines.append(f"*{len(corrections)} correction(s) suggested*")
        lines.append("")

        for i, c in enumerate(corrections, 1):
            conf = c.get("confidence", "unknown")
            confidence_counts[conf] = confidence_counts.get(conf, 0) + 1

            conf_emoji = {"high": "🔴", "medium": "🟡", "low": "🟢"}.get(conf, "⚪")

            lines.append(f"### Correction {i} — {conf_emoji} {conf.upper()} confidence")
            lines.append(f"**Original:** `{c.get('original', '[missing]')}`")
            lines.append(f"**Suggested:** `{c.get('correction', '[missing]')}`")
            lines.append(f"**Reason:** {c.get('reason', '[no reason given]')}")
            lines.append("")

    lines.append("---")
    lines.append("## Confidence Summary")
    lines.append(f"- High confidence: {confidence_counts.get('high', 0)}")
    lines.append(f"- Medium confidence: {confidence_counts.get('medium', 0)}")
    lines.append(f"- Low confidence: {confidence_counts.get('low', 0)}")

    return "\n".join(lines)


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

def main():
    print("=" * 60)
    print("OCR Correction Test — Gemini Pro + Stacey context")
    print("=" * 60)

    # 1. Extract Stacey context (done once, reused for all pages)
    stacey_context = extract_stacey_context(STACEY_PDF, STACEY_KEYWORDS, MAX_STACEY_CHARS)

    # 2. Extract olmOCR pages
    print(f"\nExtracting olmOCR pages {TEST_PAGES} from {OLMOCR_OUTPUT.name}...")
    ocr_pages = extract_olmocr_pages(OLMOCR_OUTPUT, TEST_PAGES)
    print(f"  Extracted {len(ocr_pages)} pages")

    # 3. Run correction pass for each page
    all_corrections = {}

    for page_num in TEST_PAGES:
        print(f"\n--- Page {page_num} ---")

        if page_num not in ocr_pages:
            print(f"  Skipping — page not found in olmOCR output")
            all_corrections[page_num] = []
            continue

        ocr_text = ocr_pages[page_num]

        if "[OCR FAILED" in ocr_text:
            print(f"  Skipping — olmOCR failed on this page")
            all_corrections[page_num] = []
            continue

        prompt = build_correction_prompt(page_num, ocr_text, stacey_context)
        corrections = call_gemini_correction(prompt, page_num)
        all_corrections[page_num] = corrections

    # 4. Save JSON output
    json_path = OUTPUT_DIR / "ocr_correction_test_results.json"
    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(all_corrections, f, indent=2, ensure_ascii=False)
    print(f"\nJSON results saved to {json_path}")

    # 5. Save markdown report
    report = format_corrections_report(all_corrections)
    md_path = OUTPUT_DIR / "ocr_correction_test_report.md"
    with open(md_path, "w", encoding="utf-8") as f:
        f.write(report)
    print(f"Markdown report saved to {md_path}")

    # 6. Print summary to console
    print("\n" + "=" * 60)
    print("RESULTS SUMMARY")
    print("=" * 60)
    for page_num, corrections in sorted(all_corrections.items()):
        if corrections:
            print(f"\nPage {page_num}: {len(corrections)} correction(s)")
            for c in corrections:
                conf = c.get("confidence", "?").upper()
                print(f"  [{conf}] '{c.get('original', '')}' → '{c.get('correction', '')}'")
                print(f"         {c.get('reason', '')}")
        else:
            print(f"\nPage {page_num}: No corrections suggested")


if __name__ == "__main__":
    main()