"""VLM-based structural labeling of PDF pages.

Sends page images to Gemini Flash Lite, gets back structural signals
(form grids, table types, key-value pairs, etc.), saves to CSV.

Usage:
    agentic/.venv/bin/python agentic/label_pages.py pdfs/01-practice.pdf
    agentic/.venv/bin/python agentic/label_pages.py /path/to/pdfs/ --output labels.csv
"""

from __future__ import annotations

import argparse
import csv
import io
import sys
import time
from pathlib import Path
from typing import Literal

from dotenv import load_dotenv
from pydantic import BaseModel, Field, conint, constr, model_validator

load_dotenv(Path(__file__).parent / ".env")

# Natural PDF lives one level up
sys.path.insert(0, str(Path(__file__).parent.parent))


# ── Pydantic output model ───────────────────────────────────────────────────

class PageStructure(BaseModel):
    """Structural signals for a single PDF page image."""

    # Quality
    has_handwriting: bool = Field(default=False,
        description="Any handwritten text visible (notes, annotations, signatures)")
    has_hand_filled_form: bool = Field(default=False,
        description="A printed form with handwritten entries in the fields. "
        "If true, has_handwriting must also be true.")
    has_rotated_90: bool = Field(default=False,
        description="Page content is rotated 90 degrees")
    has_low_quality_scan: bool = Field(default=False,
        description="Faded, noisy, poor contrast, visible scan artifacts")

    # Box structures
    has_boxed_fields: bool = Field(default=False,
        description="Individual data fields each in their own box/rectangle — "
        "e.g. separate boxes for 'Name', 'Address', 'DOB', or unevenly split "
        "cells in a form grid. Each box contains ONE field or value. "
        "NOT large section containers.")
    has_organizing_boxes: bool = Field(default=False,
        description="Larger rectangles that group sections of content together — "
        "e.g. a 'Witness Information' box containing several fields inside, "
        "or a 'Comments' area boxed off from the rest. These are containers "
        "for multiple items, NOT individual field boxes.")
    has_checkboxes: bool = Field(default=False,
        description="Checkbox squares or radio buttons visible")

    # Table structures — H and V are independent axes
    table_count: int = Field(default=0,
        description="Number of distinct data tables on this page (0 if none)")
    table_row_count: Literal["0", "1-5", "6-20", "21-50", "50+"] = Field(
        default="0",
        description="Approximate total data rows across all tables. "
        "Use '0' if no tables.")
    has_column_headers: bool = Field(default=False,
        description="A distinct header row immediately above repeated data rows. "
        "Only true if you can identify 2+ header words above at least 2 columns "
        "of repeated row entries. A page title or section heading is NOT a column header.")
    table_has_h_ruling_lines: bool = Field(default=False,
        description="Tables have horizontal ruling lines between rows")
    table_has_h_zebra_stripes: bool = Field(default=False,
        description="Tables have alternating row background colors/shading")
    table_has_h_whitespace: bool = Field(default=False,
        description="Table rows separated by whitespace gaps ONLY — no ruling lines, "
        "no stripes. If ruling lines separate rows, do NOT also flag this.")
    table_has_v_ruling_lines: bool = Field(default=False,
        description="Tables have vertical ruling lines between columns")
    table_has_v_whitespace: bool = Field(default=False,
        description="Table columns separated by whitespace gaps or text alignment ONLY — "
        "no vertical ruling lines. If vertical lines separate columns, do NOT also flag this.")

    # Layout structures
    has_key_value_pairs: bool = Field(default=False,
        description="Label: value pairs laid out spatially "
        "(label on left/value on right, or label above value)")
    has_section_headers: bool = Field(default=False,
        description="Distinct section headings (bold, larger, or underlined) "
        "dividing the page into named sections")
    has_narrative: bool = Field(default=False,
        description="Flowing paragraph text — sentences and paragraphs, "
        "not just labels or data")
    has_multi_column_text: bool = Field(default=False,
        description="Narrative text in 2+ columns side by side (newspaper-style). "
        "NOT a data table — flowing prose with wrapped lines in each column.")
    has_empty_form: bool = Field(default=False,
        description="A blank/unfilled form template — fields are empty, no data "
        "has been entered. Has labels and structure but no actual values. "
        "Contrast with has_hand_filled_form (form WITH entries).")
    has_fill_in_blanks: bool = Field(default=False,
        description="Underline blanks or dot leaders for entries "
        "(Name: _______ or Total........$)")
    has_visual_media: bool = Field(default=False,
        description="Photographs, charts, graphs, maps, or other non-text visual content")
    has_multiple_records: bool = Field(default=False,
        description="Page contains multiple similar records or entries")

    # Free text — layout only, not content
    page_description: str = Field(
        description="1-3 sentences describing PAGE LAYOUT and FORMATTING only. "
        "Spatial arrangement, structure types, visual patterns. "
        "Do NOT describe content/subject matter. "
        "Example: 'Dense form grid in upper 2/3 with labeled cells, "
        "horizontal rule, then narrative text at bottom. Small logo top-left.'")

    @model_validator(mode="after")
    def validate_consistency(self) -> "PageStructure":
        if self.has_empty_form and self.has_hand_filled_form:
            raise ValueError("A form cannot be both empty and hand-filled.")
        if self.has_hand_filled_form and not self.has_handwriting:
            raise ValueError("has_hand_filled_form requires has_handwriting to be true.")
        if self.table_count == 0:
            if self.table_row_count != "0":
                raise ValueError("table_row_count must be '0' when table_count is 0.")
            if any([self.table_has_h_ruling_lines, self.table_has_h_zebra_stripes,
                    self.table_has_h_whitespace, self.table_has_v_ruling_lines,
                    self.table_has_v_whitespace, self.has_column_headers]):
                raise ValueError("All table_has_* and has_column_headers must be false when table_count is 0.")
        return self


# ── VLM prompt ───────────────────────────────────────────────────────────────

SYSTEM_PROMPT = """\
You are analyzing PDF page images for structural/layout features.
Focus ONLY on visual structure — how content is arranged — not what it says.
You may read short text snippets to identify headers; do not summarize content.

CRITICAL DISTINCTIONS:

BOXED FIELDS vs ORGANIZING BOXES:
- "has_boxed_fields" = individual data fields each in their own small box.
  Each box holds ONE piece of data (Name, DOB, SSN). Think form cells.
  These may form a grid or be scattered — the key is each box = one field.
- "has_organizing_boxes" = larger rectangles grouping sections of content.
  A "Witness Information" box containing several fields inside, or a
  "Comments" section boxed off. These are containers, not individual fields.
- A page can have BOTH (field boxes inside organizing boxes).

TABLE H/V STRUCTURE (independent axes):
- Horizontal (row separation) and vertical (column separation) are INDEPENDENT.
- Example: zebra-striped rows + whitespace column gaps =
  table_has_h_zebra_stripes=true, table_has_v_whitespace=true.
- RULING LINES vs WHITESPACE are MUTUALLY EXCLUSIVE per axis.
  If ruling lines separate rows, do NOT also flag h_whitespace.
  If ruling lines separate columns, do NOT also flag v_whitespace.
- All table_has_* fields must be false when table_count=0.

COLUMN HEADERS vs PAGE TITLES:
- "has_column_headers" = header row immediately above repeated data rows.
  Must have 2+ header words above at least 2 columns of repeated row entries.
- A page title, document heading, or section heading is NOT a column header.

MULTI-COLUMN TEXT vs TABLE:
- "has_multi_column_text" = narrative prose in 2+ columns with wrapped paragraph lines.
- Tables have short cell-like entries aligned by rows. Columns of prose are NOT tables.

FORM FILL STATUS:
- "has_empty_form" = a blank form template with labels/structure but no data entered.
- "has_hand_filled_form" = a printed form with handwritten entries filled in.
  If true, has_handwriting must also be true.
- A page can be has_empty_form with has_fill_in_blanks (blank form with underlines).
- has_empty_form and has_hand_filled_form should NOT both be true for the same page.

ROW COUNT:
- table_row_count: count data rows only (not header rows). Round to nearest bucket.

PAGE DESCRIPTION:
- Describe LAYOUT AND FORMATTING ONLY. 1-3 sentences, max 300 characters.
- Say "dense grid of labeled cells" NOT "arrest report for John Smith".
- Describe spatial relationships: "upper half", "left column", "below the rule".\
"""

USER_PROMPT = "Analyze this page image for structural features."


# ── CSV columns ──────────────────────────────────────────────────────────────

CSV_COLUMNS = [
    "pdf_path", "page_number",
    # Quality
    "has_handwriting", "has_hand_filled_form", "has_rotated_90", "has_low_quality_scan",
    # Box structures
    "has_boxed_fields", "has_organizing_boxes", "has_checkboxes",
    # Table
    "table_count", "table_row_count", "has_column_headers",
    "table_has_h_ruling_lines", "table_has_h_zebra_stripes", "table_has_h_whitespace",
    "table_has_v_ruling_lines", "table_has_v_whitespace",
    # Layout
    "has_key_value_pairs", "has_section_headers", "has_narrative", "has_empty_form",
    "has_multi_column_text", "has_fill_in_blanks", "has_visual_media",
    "has_multiple_records",
    # Description
    "page_description",
]


def flatten(ps: PageStructure, pdf_path: str, page_num: int) -> dict:
    """Convert PageStructure to a flat dict for CSV."""
    d = ps.model_dump()
    d["pdf_path"] = pdf_path
    d["page_number"] = page_num
    return d


def render_page_to_bytes(page, resolution: int) -> bytes:
    """Render a page to PNG bytes."""
    img = page.render(resolution=resolution)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()


def collect_pdf_paths(inputs: list[str]) -> list[Path]:
    """Expand input paths: files stay as-is, directories are globbed for *.pdf."""
    paths = []
    for inp in inputs:
        p = Path(inp)
        if p.is_dir():
            paths.extend(sorted(p.glob("**/*.pdf")))
        elif p.is_file() and p.suffix.lower() == ".pdf":
            paths.append(p)
        else:
            print(f"  SKIP: {inp} (not a PDF or directory)")
    return paths


def main():
    parser = argparse.ArgumentParser(description="Label PDF pages with structural signals via VLM")
    parser.add_argument("inputs", nargs="+", help="PDF files or directories")
    parser.add_argument("--output", default="agentic/output/page_labels.csv", help="Output CSV path")
    parser.add_argument("--resolution", type=int, default=72, help="Render DPI (default: 72)")
    parser.add_argument("--model", default="google-gla:gemini-3.1-flash-lite-preview", help="VLM model")
    parser.add_argument("--append", action="store_true", help="Append to existing CSV")
    parser.add_argument("--limit", type=int, default=None, help="Max pages to process")
    args = parser.parse_args()

    # Collect PDFs
    pdf_paths = collect_pdf_paths(args.inputs)
    if not pdf_paths:
        print("No PDFs found.")
        return

    print(f"Found {len(pdf_paths)} PDFs")

    # Create agent
    from pydantic_ai import Agent
    from pydantic_ai.messages import BinaryContent
    from pydantic_ai.settings import ModelSettings

    agent = Agent(
        args.model,
        output_type=PageStructure,
        instructions=SYSTEM_PROMPT,
        retries=3,
    )

    # Open CSV
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    mode = "a" if args.append else "w"
    csvfile = open(output_path, mode, newline="", encoding="utf-8")
    writer = csv.DictWriter(csvfile, fieldnames=CSV_COLUMNS, extrasaction="ignore")
    if not args.append or not output_path.exists() or output_path.stat().st_size == 0:
        writer.writeheader()

    # Process
    from natural_pdf import PDF

    total_pages = 0
    total_errors = 0
    total_tokens = 0
    start_time = time.monotonic()

    for pdf_path in pdf_paths:
        try:
            pdf = PDF(str(pdf_path))
        except Exception as e:
            print(f"  FAIL open: {pdf_path} — {e}")
            continue

        for page_idx, page in enumerate(pdf.pages):
            if args.limit and total_pages >= args.limit:
                pdf.close()
                break

            page_start = time.monotonic()
            try:
                img_bytes = render_page_to_bytes(page, args.resolution)
                result = agent.run_sync(
                    [USER_PROMPT, BinaryContent(data=img_bytes, media_type="image/png")],
                    model_settings=ModelSettings(temperature=0),
                )
                row = flatten(result.output, str(pdf_path), page_idx)
                writer.writerow(row)
                csvfile.flush()

                tokens = result.usage().total_tokens
                total_tokens += tokens
                elapsed = time.monotonic() - page_start

                # Progress
                signals = []
                if row.get("has_dense_form_grid"):
                    signals.append("form_grid")
                if row.get("table_count", 0) > 0:
                    signals.append(f"{row['table_count']}tbl")
                if row.get("has_key_value_pairs"):
                    signals.append("kv")
                if row.get("has_narrative"):
                    signals.append("narrative")
                if row.get("has_checkboxes"):
                    signals.append("checkboxes")
                sig_str = ", ".join(signals) if signals else "—"
                print(f"  OK: {pdf_path.name} p{page_idx} — [{sig_str}] ({elapsed:.1f}s, {tokens}tok)")

            except Exception as e:
                # Write error row
                error_row = {"pdf_path": str(pdf_path), "page_number": page_idx,
                             "page_description": f"ERROR: {e}"}
                writer.writerow(error_row)
                csvfile.flush()
                print(f"  ERR: {pdf_path.name} p{page_idx} — {e}")
                total_errors += 1

            total_pages += 1

        pdf.close()

        if args.limit and total_pages >= args.limit:
            break

    csvfile.close()

    # Summary
    elapsed = time.monotonic() - start_time
    print(f"\nDone: {total_pages} pages, {total_errors} errors, "
          f"{total_tokens} tokens, {elapsed:.0f}s")
    print(f"Output: {output_path}")


if __name__ == "__main__":
    main()
