"""
Vision-pass batch runner for 2629 Erskine Park (priority sheets).

Renders pre-defined PDF pages, calls Claude vision with the lfcs-drawing-extract-vision
structured-extraction prompt, and saves per-sheet JSON to scripts/vision/output/2629/.

Usage:
    py scripts/vision/lfcs-erskine-vision.py [--only S2.00,S2.01]
"""

from __future__ import annotations

import argparse
import base64
import io
import json
import os
import re
import sys
from dataclasses import dataclass, field
from pathlib import Path

from anthropic import Anthropic
from dotenv import load_dotenv
from pdf2image import convert_from_path
from PIL import Image

REPO = Path(__file__).resolve().parents[2]
JOB_FOLDER = (
    REPO
    / "HQ-Vault"
    / "11_OpsMan_LFCS"
    / "Operations"
    / "Jobs"
    / "Upcoming"
    / "2629 - TCE - Erskine Park WWTP Slab"
    / "03 - Plans and Specifications"
    / "03a - Issued for Construction"
)
STRUCTURAL_PDF = JOB_FOLDER / "250643_Erskine Park_90% Design_Structural package.pdf"
CIVIL_PDF = JOB_FOLDER / "250643_Erskine Park_90% Design_Civil package.pdf"
OUTPUT_DIR = REPO / "scripts" / "vision" / "output" / "2629"

MODEL = "claude-opus-4-7"
DPI = 200
MAX_EDGE_PX = 1568
JPEG_QUALITY = 85
MAX_TOKENS = 4096


@dataclass
class Sheet:
    drawing_no: str
    pdf_path: Path
    page: int
    discipline: str
    sheet_size: str
    lfcs_scope: str
    title: str


SHEETS: list[Sheet] = [
    Sheet("S2.00", STRUCTURAL_PDF, 3, "structural", "A1", "primary",
          "WWTP Slab Plans"),
    Sheet("S2.01", STRUCTURAL_PDF, 4, "structural", "A1", "primary",
          "WWTP Slab Reinforcement Plans"),
    Sheet("S2.05", STRUCTURAL_PDF, 5, "structural", "A1", "primary",
          "WWTP Slab Sections"),
    Sheet("S3.00", STRUCTURAL_PDF, 6, "structural", "A1",
          "primary-walkway-concrete + reference-steel-walkway",
          "Pedestrian Walkway Plan and Details"),
    Sheet("C103", CIVIL_PDF, 8, "civil", "A1", "governs-walkway-concrete-pavement",
          "Walkway Plan"),
]


PROMPT_TEMPLATE = """This is page {page} of {pdf_filename} ({drawing_no} - "{title}", sheet_size {sheet_size}, discipline {discipline}, lfcs_scope {lfcs_scope}).

You are running the lfcs-drawing-extract-vision pass for LFCS Civil Solutions on the 2629 Erskine Park WWTP Slab bid (TCE Contracting head contractor, due 2026-05-07). LFCS scope: concrete slab + walkway + bund wall (BW1). NOT in scope: bored piles, stormwater sump, sludge platform steel, walkway steel framing, existing wall demolition.

Extract ONLY items visible on the drawing. Do NOT infer. If a value is not visible, omit the field.

Reply with a SINGLE JSON OBJECT (no markdown fences, no preamble) matching this schema:

{{
  "drawing_no": "{drawing_no}",
  "page": {page},
  "geometry_dimensions": [
    {{"dimension_value_mm": <number>, "dimension_what": "<short description>", "location_on_sheet": "<where on the drawing>"}}
  ],
  "materials_called_out": [
    {{"material": "<verbatim label>", "what_for": "<element>", "location": "<where on sheet>"}}
  ],
  "section_markers": [
    {{"section_label": "<e.g. SECTION 1, A-A>", "references_drawing": "<sheet ref or 'this sheet'>"}}
  ],
  "detail_bubbles": [
    {{"detail_label": "<e.g. DETAIL 1>", "what_in": "<what it shows>", "references_drawing": "<sheet ref>"}}
  ],
  "drawn_clashes_with_text_labels": [
    {{"what": "<element>", "where": "<location>", "observation": "<what differs vs text label>"}}
  ],
  "hand_marked_annotations": [
    {{"annotation": "<text>", "what": "<what it relates to>", "where": "<location>"}}
  ],
  "construction_notes_visible": ["<verbatim note from drawing>"],
  "lfcs_scope_relevant_tags": ["<tag:value>"],
  "lfcs_quantity_signals": [
    {{"element": "<short name>", "approx_quantity": "<e.g. ~50 m2, qty 4 pad footings>", "evidence": "<where on sheet>"}}
  ],
  "drawing_quality_assessment": "clean|smudged|partial|illegible",
  "extraction_confidence": "High|Medium|Low|NeedsReview",
  "summary_per_drawing": "<1-paragraph summary of vision-only findings, <=120 words>"
}}

Tag conventions: material:X (e.g. material:N16-200-EW), geometry:X (e.g. geometry:slab-270mm), risk:X, discrepancy:X, existing:X, scope:LFCS-in or scope:LFCS-out.

Quantity signals: only items LFCS prices (slab, walkway slab, walkway pad footings, bund wall BW1, integrated 500x1200 RC beams). Skip out-of-scope items (piles, steel framing, sump, demo).
"""


def render_page(pdf_path: Path, page: int) -> tuple[bytes, str]:
    poppler_path = os.getenv("POPPLER_PATH") or None
    images = convert_from_path(
        str(pdf_path),
        dpi=DPI,
        first_page=page,
        last_page=page,
        fmt="png",
        poppler_path=poppler_path,
    )
    if not images:
        raise RuntimeError(f"pdf2image returned no images for {pdf_path} page {page}")
    img = images[0]
    longest = max(img.size)
    if longest > MAX_EDGE_PX:
        scale = MAX_EDGE_PX / longest
        img = img.resize((int(img.size[0] * scale), int(img.size[1] * scale)), Image.LANCZOS)
    if img.mode != "RGB":
        img = img.convert("RGB")
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=JPEG_QUALITY, optimize=True)
    return buf.getvalue(), "image/jpeg"


def call_vision(client: Anthropic, jpeg_bytes: bytes, prompt: str) -> tuple[str, dict]:
    encoded = base64.standard_b64encode(jpeg_bytes).decode("ascii")
    msg = client.messages.create(
        model=MODEL,
        max_tokens=MAX_TOKENS,
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "source": {
                            "type": "base64",
                            "media_type": "image/jpeg",
                            "data": encoded,
                        },
                    },
                    {"type": "text", "text": prompt},
                ],
            }
        ],
    )
    text = "".join(b.text for b in msg.content if b.type == "text")
    usage = {
        "input_tokens": msg.usage.input_tokens,
        "output_tokens": msg.usage.output_tokens,
    }
    return text, usage


def parse_json_strict(text: str) -> dict:
    text = text.strip()
    fence = re.match(r"^```(?:json)?\s*(.*?)\s*```$", text, re.DOTALL)
    if fence:
        text = fence.group(1).strip()
    return json.loads(text)


def main() -> int:
    load_dotenv()
    if not os.getenv("ANTHROPIC_API_KEY"):
        print("ERROR: ANTHROPIC_API_KEY not set", file=sys.stderr)
        return 2
    if not os.getenv("POPPLER_PATH"):
        print("WARN: POPPLER_PATH not set; relying on system PATH", file=sys.stderr)

    parser = argparse.ArgumentParser()
    parser.add_argument("--only", help="Comma-separated drawing numbers to run (e.g. S2.00,S2.01)")
    args = parser.parse_args()

    selected = SHEETS
    if args.only:
        wanted = {s.strip() for s in args.only.split(",")}
        selected = [s for s in SHEETS if s.drawing_no in wanted]

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    client = Anthropic()
    totals = {"input": 0, "output": 0}

    for sheet in selected:
        print(f"\n=== {sheet.drawing_no} ({sheet.pdf_path.name} p{sheet.page}) ===", file=sys.stderr)
        if not sheet.pdf_path.exists():
            print(f"  ERROR: missing PDF {sheet.pdf_path}", file=sys.stderr)
            continue

        jpeg, _ = render_page(sheet.pdf_path, sheet.page)
        print(f"  rendered: {len(jpeg) / 1024:.1f} KB", file=sys.stderr)

        prompt = PROMPT_TEMPLATE.format(
            page=sheet.page,
            pdf_filename=sheet.pdf_path.name,
            drawing_no=sheet.drawing_no,
            title=sheet.title,
            sheet_size=sheet.sheet_size,
            discipline=sheet.discipline,
            lfcs_scope=sheet.lfcs_scope,
        )
        text, usage = call_vision(client, jpeg, prompt)
        totals["input"] += usage["input_tokens"]
        totals["output"] += usage["output_tokens"]
        print(f"  tokens: in={usage['input_tokens']} out={usage['output_tokens']}", file=sys.stderr)

        out_path = OUTPUT_DIR / f"{sheet.drawing_no}.json"
        try:
            parsed = parse_json_strict(text)
        except json.JSONDecodeError as e:
            print(f"  WARN: JSON parse failed ({e}); writing raw .txt", file=sys.stderr)
            out_path.with_suffix(".txt").write_text(text, encoding="utf-8")
            continue

        out_path.write_text(json.dumps(parsed, indent=2), encoding="utf-8")
        print(f"  wrote: {out_path.relative_to(REPO)}", file=sys.stderr)

    print(f"\nTOTAL TOKENS: input={totals['input']} output={totals['output']}", file=sys.stderr)
    cost_in = totals["input"] / 1_000_000 * 15
    cost_out = totals["output"] / 1_000_000 * 75
    print(f"APPROX COST (Opus 4.7 list): ${cost_in + cost_out:.2f}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())