"""
step3_extract-posn_claude-4-6.py
----------------------------------
Extracts structured data from verified OCR text produced by
step1_ocr_wardiaries_claude-4-6.py (or the olmOCR equivalent).

Handles two page types:

  DIARY PAGES (HTML tables — Place / Date / Hour / Summary / Remarks):
    Extracts dated narrative entries and all position records.
    Output: {stem}_claude-4-6_narratives.json
            {stem}_claude-4-6_positions.json

  APPENDIX PAGES (plain text — message forms, movement orders, arty tables,
    field returns, ISUMs, patrol reports, traces, etc.):
    Extracts document metadata and all grid references / positions mentioned.
    Output: {stem}_claude-4-6_appx_records.json
            {stem}_claude-4-6_appx_positions.json

Requirements:
    pip install requests python-dotenv

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

Usage:
    python scripts/step3_extract-posn_claude-4-6.py --input "G:/path/to/step1_Calgary-Highlanders_War-Diary_Sep44_claude-4-6.md" --output_dir "G:/path/to/outputs_step3_llm"
    python scripts/step3_extract-posn_claude-4-6.py --input "..." --output_dir "..." --first_page 7 --last_page 27
"""

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

from dotenv import load_dotenv
import requests

load_dotenv()

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

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

# ---------------------------------------------------------------------------
# Prompts
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = """\
You are a military historian extracting structured data from a WWII Canadian Army
war diary file (Calgary Highlanders, September–October 1944).

You will receive the raw text of one page. Determine the page type and extract
accordingly. Output valid JSON only. No prose before or after the JSON.

════════════════════════════════════════
PAGE TYPE DETECTION
════════════════════════════════════════

page_type must be one of:
  "diary"           — HTML table with Place / Date / Hour / Summary / Remarks columns
  "message_form"    — Army Form C2136 or similar signal/message form
  "movement_order"  — Numbered movement or operation order
  "arty_table"      — Artillery DF / SOS / fire task table
  "field_return"    — Strength return (officers or other ranks)
  "isum"            — Intelligence summary
  "patrol_report"   — Patrol programme or patrol report
  "trace"           — Map, sketch map, disposition diagram, or trace
  "admin"           — Part I/II orders, nominal rolls, boilerplate instructions
  "other"           — Anything that does not fit the above

════════════════════════════════════════
OUTPUT SCHEMA — ALL PAGES
════════════════════════════════════════

{
  "page_type": "diary | message_form | movement_order | ...",
  "date_warning": null or "explanation",

  // ── DIARY PAGES ONLY ──────────────────────────────────────────────────────
  "narratives": [
    {
      "date": "1 Sep 44",
      "date_inferred": false,
      "date_span": null,
      "place": "Ste. Foy",
      "summary": "Full verbatim narrative text...",
      "weather": "Fair and warm",
      "no_change": false,
      "source_page": 7
    }
  ],

  // ── APPENDIX PAGES ONLY ───────────────────────────────────────────────────
  "appx_record": {
    "page_type": "message_form",
    "document_id": "GO-7",
    "date": "27 Sep 44",
    "time": "2205A",
    "from_unit": "HQ RCA 2 Cdn Inf Div",
    "to_units": ["5 CIB", "RHC", "CALG HIGHRS"],
    "subject": "DF Task Table No 35 — amendment",
    "summary": "One sentence description of what this document contains.",
    "units_mentioned": ["Calgary Highlanders", "RHC", "R de Mais", "5 Cdn Fd Regt"],
    "source_page": 12
  },

  // ── ALL PAGES — positions found anywhere on the page ─────────────────────
  "positions": [
    {
      "date": "1 Sep 44",
      "date_inferred": false,
      "hour": null,
      "grid": "255330",
      "grid_inferred": false,
      "place_name": "Ste. Foy",
      "sheet_ref": "Sheet 861",
      "category": "END_OF_DAY",
      "subunit": null,
      "friendly_unit": null,
      "no_movement": false,
      "confidence": "high",
      "context": "arrived in the little village of Ste. Foy east of Longueville",
      "source_page": 7
    }
  ]
}

For diary pages:    populate "narratives" and "positions". Set "appx_record" to null.
For appendix pages: populate "appx_record" and "positions". Set "narratives" to [].
For blank pages:    set all arrays empty, appx_record null.

════════════════════════════════════════
DIARY PAGE RULES (page_type = "diary")
════════════════════════════════════════

Each <tr> in the HTML table is one diary entry. Extract:

narratives:
  date        — from Date <td>. Join multiline ("1 Sep\n44" → "1 Sep 44").
  date_span   — if entry covers multiple dates ("6-8 Sep"), set to "6-8 Sep 44",
                date to the first date.
  date_inferred — true if date carried forward from a previous row.
  place       — primary place name from Place <td>, excluding grid/sheet refs.
  summary     — full verbatim text from Summary <td>. Nothing omitted.
  weather     — weather note if present on the page. null if absent.
  no_change   — true if entry states no change from previous day.

positions (cast a wide net — extract from BOTH Place <td> AND Summary <td>):
  PLACE COLUMN → category END_OF_DAY (last entry per date) or UNIT_MOVEMENT.
  SUMMARY TEXT → category SUBUNIT, FRIENDLY, ENEMY, PATROL, or MISC.
  Only ONE position per date may have category END_OF_DAY.

════════════════════════════════════════
APPENDIX PAGE RULES (all other page_types)
════════════════════════════════════════

appx_record:
  document_id   — order/form number if visible (e.g. "Mov Order No 5",
                  "ISUM No 45", "GO-3", "DF Task Table No 35"). null if absent.
  date          — date of the document. null if not determinable.
  time          — time or date-time group if present (e.g. "272205A"). null if absent.
  from_unit     — originating unit/HQ. null if absent.
  to_units      — list of addressees. [] if absent.
  subject       — subject line or a one-sentence description of purpose.
  summary       — one to three sentence plain-English summary of what this
                  document records or orders. No verbatim transcription.
  units_mentioned — all unit names appearing anywhere on the page.

positions (extract ALL grid references and named locations):
  Extract every grid reference (4-digit, 6-digit, 8-digit) and every named
  location from the entire page, regardless of context.
  Use the same position schema as diary pages.
  category:
    UNIT_MOVEMENT — if associated with the Calgary Highlanders' own movement
    FRIENDLY      — if associated with another Allied unit
    ENEMY         — if associated with enemy forces
    PATROL        — if from a patrol programme or patrol report
    DF_TASK       — if from an DF/SOS artillery task table
    MISC          — everything else (objectives, named features, route points)
  date            — take from document date if not stated per-position. null if unknown.
  confidence:
    "high"   — grid explicitly written next to a place name
    "medium" — grid present but context unclear
    "low"    — place name only, no grid; or grid with no place name context

════════════════════════════════════════
GRID REFERENCE RULES (all page types)
════════════════════════════════════════

grid: digits only, no prefix, no punctuation.
  "MR 2468" → "2468"    "GR 442891" → "442891"    "MR 24.68" → "2468"
If the diary gives a 4-digit grid, expand: "2553" → "255535". Set grid_inferred = true.
Do not guess or expand 6-digit grids. Set grid_inferred = false.
8-digit grids: keep first 3 + last 3 digits → "44289100" → "442891".
5-digit grids: set grid = null, confidence = "low", note in context.
If no grid is available, set grid to null and record place_name instead.

sheet_ref: carry forward the most recent sheet reference seen on the page.
hour: HHMM format if stated ("0930"). null if absent.
  Ignore Hour <td> values that look like years ("44") or grid refs.

Output ONLY the JSON object. No preamble. No explanation. No trailing text.\
"""


def make_user_prompt(page_num: int, page_text: str) -> str:
    return (
        f"This is page {page_num} of the war diary file. "
        f"Determine the page type, then extract accordingly.\n\n"
        f"PAGE TEXT:\n{page_text}"
    )


# ---------------------------------------------------------------------------
# MD page parser
# ---------------------------------------------------------------------------

def parse_md_pages(md_text: str) -> dict[int, str]:
    pages = {}
    page_re = re.compile(r'^## Page (\d+)\s*$', re.MULTILINE)
    matches = list(page_re.finditer(md_text))
    for i, m in enumerate(matches):
        page_num = int(m.group(1))
        start = m.end()
        end = matches[i + 1].start() if i + 1 < len(matches) else len(md_text)
        pages[page_num] = md_text[start:end].strip()
    return pages


# ---------------------------------------------------------------------------
# JSON response parser
# ---------------------------------------------------------------------------

def parse_json_response(raw: str, page_num: int) -> dict | None:
    cleaned = raw.strip()
    cleaned = re.sub(r'^```json\s*', '', cleaned, flags=re.IGNORECASE)
    cleaned = re.sub(r'^```\s*', '', cleaned)
    cleaned = re.sub(r'\s*```$', '', cleaned)
    cleaned = cleaned.strip()
    try:
        return json.loads(cleaned)
    except json.JSONDecodeError as e:
        print(f"    WARNING: JSON parse failed for page {page_num}: {e}")
        print(f"    Raw response (first 300 chars): {raw[:300]}")
        return None


# ---------------------------------------------------------------------------
# API call — Anthropic
# ---------------------------------------------------------------------------

def extract_page(page_num: int, page_text: str, api_key: str) -> dict | None:
    """Send one page to Claude via Anthropic API and return parsed extraction result."""

    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": make_user_prompt(page_num, page_text),
            },
        ],
    }

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

            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=" ")

            raw_content = data["content"][0]["text"]
            return parse_json_response(raw_content, page_num)

        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 None
        except Exception as e:
            print(f"    Error on page {page_num, attempt}/{MAX_RETRIES}: {e}")
            if attempt < MAX_RETRIES:
                time.sleep(RETRY_DELAY)
            else:
                return None


# ---------------------------------------------------------------------------
# End-of-day assignment (diary positions only)
# ---------------------------------------------------------------------------

def assign_end_of_day(positions: list[dict]) -> list[dict]:
    from collections import defaultdict
    date_groups: dict[str, list[int]] = defaultdict(list)
    for i, p in enumerate(positions):
        if p.get('date'):
            date_groups[p['date']].append(i)

    for date, indices in date_groups.items():
        eod_indices = [i for i in indices if positions[i].get('category') == 'END_OF_DAY']
        if len(eod_indices) == 1:
            continue
        if len(eod_indices) > 1:
            for i in eod_indices[:-1]:
                positions[i]['category'] = 'UNIT_MOVEMENT'
            continue
        with_grid = [i for i in indices if positions[i].get('grid')]
        unit_movement = [i for i in with_grid if positions[i].get('category') == 'UNIT_MOVEMENT']
        if unit_movement:
            positions[unit_movement[-1]]['category'] = 'END_OF_DAY'
        elif with_grid:
            positions[with_grid[-1]]['category'] = 'END_OF_DAY'
    return positions


def post_process_positions(positions: list[dict]) -> list[dict]:
    for p in positions:
        raw = re.sub(r'[^0-9]', '', p.get('grid') or '')
        if len(raw) == 4:
            p['grid'] = raw[0:2] + '5' + raw[2:4] + '5'
            p['grid_inferred'] = True
        elif len(raw) == 6:
            p['grid'] = raw
            p['grid_inferred'] = False
        elif len(raw) == 8:
            p['grid'] = raw[0:3] + raw[4:7]
            p['grid_inferred'] = False
        elif len(raw) == 5:
            p['grid'] = None
            p['grid_inferred'] = False
            p['confidence'] = 'low'
            p['context'] = (p.get('context') or '') + ' [5-figure grid — needs human review]'
        else:
            p['grid'] = None
            p['grid_inferred'] = False
            p['confidence'] = 'low'
    return positions


def dedup_positions(positions: list[dict]) -> list[dict]:
    seen = set()
    result = []
    for p in positions:
        key = (p.get('date'), p.get('grid'), (p.get('context') or '')[:60])
        if key not in seen:
            seen.add(key)
            result.append(p)
    return result


# ---------------------------------------------------------------------------
# Main extraction loop
# ---------------------------------------------------------------------------

def extract_diary(md_path: Path, output_dir: Path, api_key: str,
                  first_page: int = None, last_page: int = None):

    print(f"\n{'='*60}")
    print(f"Extracting: {md_path.name}")
    if first_page or last_page:
        print(f"Pages: {first_page or 'start'} → {last_page or 'end'}")
    print(f"{'='*60}")

    stem = md_path.stem
    for suffix in ('_olmocr', '_claude-4-6', 'step1_'):
        stem = stem.replace(suffix, '')
    stem = stem.strip('_')

    # Diary outputs
    narratives_out   = output_dir / f"{stem}_claude-4-6_narratives.json"
    positions_out    = output_dir / f"{stem}_claude-4-6_positions.json"
    # Appendix outputs
    appx_records_out = output_dir / f"{stem}_claude-4-6_appx_records.json"
    appx_posns_out   = output_dir / f"{stem}_claude-4-6_appx_positions.json"
    # Checkpoint file — deleted on clean completion
    checkpoint_path  = output_dir / f"{stem}_claude-4-6_checkpoint.json"

    md_text = md_path.read_text(encoding='utf-8')
    all_pages = parse_md_pages(md_text)

    if not all_pages:
        print("  ERROR: No pages found in MD file. Check file format.")
        return

    page_nums = sorted(all_pages.keys())
    if first_page:
        page_nums = [p for p in page_nums if p >= first_page]
    if last_page:
        page_nums = [p for p in page_nums if p <= last_page]

    # ── Resume from checkpoint if one exists ─────────────────────────────────
    all_narratives   = []
    all_positions    = []
    all_appx_records = []
    all_appx_posns   = []
    skipped_pages    = []
    error_pages      = []
    resume_from      = None

    if checkpoint_path.exists():
        try:
            ckpt = json.loads(checkpoint_path.read_text(encoding='utf-8'))
            resume_from      = ckpt.get('last_completed_page')
            all_narratives   = ckpt.get('narratives',    [])
            all_positions    = ckpt.get('positions',     [])
            all_appx_records = ckpt.get('appx_records',  [])
            all_appx_posns   = ckpt.get('appx_positions',[])
            skipped_pages    = ckpt.get('skipped_pages', [])
            error_pages      = ckpt.get('error_pages',   [])
            print(f"  RESUMING from checkpoint — last completed page: {resume_from}")
            page_nums = [p for p in page_nums if p > resume_from]
            print(f"  Remaining pages: {len(page_nums)}")
        except Exception as e:
            print(f"  WARNING: Could not read checkpoint ({e}) — starting from scratch.")

    print(f"  Found {len(all_pages)} total pages, processing {len(page_nums)}.")
    print(f"  Model: {MODEL}")
    print(f"  Diary outputs   : {narratives_out.name} / {positions_out.name}")
    print(f"  Appendix outputs: {appx_records_out.name} / {appx_posns_out.name}")

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

    def save_checkpoint(last_page_done: int):
        checkpoint_path.write_text(
            json.dumps({
                'last_completed_page': last_page_done,
                'narratives':    all_narratives,
                'positions':     all_positions,
                'appx_records':  all_appx_records,
                'appx_positions':all_appx_posns,
                'skipped_pages': skipped_pages,
                'error_pages':   error_pages,
            }, ensure_ascii=False),
            encoding='utf-8'
        )

    for i, page_num in enumerate(page_nums, start=1):
        page_text = all_pages[page_num]

        if not page_text.strip() or page_text.strip() == '[BLANK]':
            print(f"  Page {page_num} ({i}/{len(page_nums)})... skipped (blank)")
            skipped_pages.append(page_num)
            save_checkpoint(page_num)
            continue

        print(f"  Page {page_num} ({i}/{len(page_nums)})...", end=" ", flush=True)

        result = extract_page(page_num, page_text, api_key)

        if result is None:
            print("ERROR — extraction failed")
            error_pages.append(page_num)
            save_checkpoint(page_num)
            continue

        page_type    = result.get('page_type', 'other')
        date_warning = result.get('date_warning')
        narratives   = result.get('narratives', []) or []
        positions    = result.get('positions', []) or []
        appx_record  = result.get('appx_record')

        for n in narratives:
            n['source_page'] = page_num
        for p in positions:
            p['source_page'] = page_num
        if appx_record:
            appx_record['source_page'] = page_num

        if page_type == 'diary':
            all_narratives.extend(narratives)
            all_positions.extend(positions)
            parts = [f"diary — {len(narratives)} narrative(s), {len(positions)} position(s)"]
        else:
            if appx_record:
                all_appx_records.append(appx_record)
            all_appx_posns.extend(positions)
            doc_id = (appx_record or {}).get('document_id') or ''
            parts = [f"{page_type}{' — ' + doc_id if doc_id else ''} — {len(positions)} position(s)"]

        if date_warning:
            parts.append(f"DATE WARNING: {date_warning}")
        print(", ".join(parts))

        # Save checkpoint after every successfully processed page
        save_checkpoint(page_num)

        if i < len(page_nums):
            time.sleep(0.5)

    # Post-process diary positions
    print(f"\n  Running end-of-day assignment (diary)...")
    all_positions = assign_end_of_day(all_positions)
    all_positions = post_process_positions(all_positions)
    all_positions = dedup_positions(all_positions)

    # Post-process appendix positions
    all_appx_posns = post_process_positions(all_appx_posns)
    all_appx_posns = dedup_positions(all_appx_posns)

    # Write final outputs
    narratives_out.write_text(
        json.dumps(all_narratives, indent=2, ensure_ascii=False), encoding='utf-8')
    positions_out.write_text(
        json.dumps(all_positions, indent=2, ensure_ascii=False), encoding='utf-8')
    appx_records_out.write_text(
        json.dumps(all_appx_records, indent=2, ensure_ascii=False), encoding='utf-8')
    appx_posns_out.write_text(
        json.dumps(all_appx_posns, indent=2, ensure_ascii=False), encoding='utf-8')

    # Clean up checkpoint — run completed successfully
    if checkpoint_path.exists():
        checkpoint_path.unlink()
        print(f"  Checkpoint removed.")

    print(f"\n{'='*60}")
    print(f"Complete.")
    print(f"  Diary narratives : {len(all_narratives)} entries    → {narratives_out.name}")
    print(f"  Diary positions  : {len(all_positions)} records   → {positions_out.name}")
    print(f"  Appendix records : {len(all_appx_records)} documents → {appx_records_out.name}")
    print(f"  Appendix positions: {len(all_appx_posns)} records  → {appx_posns_out.name}")
    if skipped_pages:
        print(f"  Skipped    : pages {skipped_pages} (blank)")
    if error_pages:
        print(f"  ERRORS     : pages {error_pages} — re-run with --first_page/--last_page to retry")

    warnings = [(n.get('source_page'), n.get('date'), n.get('date_warning'))
                for n in all_narratives if n.get('date_warning')]
    if warnings:
        print(f"\n  DATE WARNINGS ({len(warnings)}):")
        for page, date, warning in warnings:
            print(f"    Page {page} | date={date} | {warning}")


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

def main():
    parser = argparse.ArgumentParser(
        description="Extract narratives, positions, and appendix records from OCR .md files using Claude Sonnet 4.6."
    )
    parser.add_argument("--input",      required=True,
        help="Path to the step1_*_claude-4-6.md (or _olmocr.md) file.")
    parser.add_argument("--output_dir", required=True,
        help="Folder to write output JSON 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("ANTHROPIC_API_KEY")
    if not api_key:
        print("Error: ANTHROPIC_API_KEY not set in environment or .env file")
        sys.exit(1)

    md_path = Path(args.input)
    if not md_path.exists():
        print(f"Error: input file not found: {md_path}")
        sys.exit(1)

    extract_diary(
        md_path, Path(args.output_dir), api_key,
        first_page=args.first_page,
        last_page=args.last_page,
    )


if __name__ == "__main__":
    main()

