state.py

from __future__ import annotations

import operator
from typing import Annotated, TypedDict

from langgraph.graph import add_messages


class ImageRef(TypedDict):
    """Lightweight reference to an image stored on disk."""
    id: str              # unique ID e.g. "page_5_crop_2"
    path: str            # disk path to PNG
    label: str           # human-readable: "Page 5 (M-101 HVAC Plan) - Gymnasium diffusers"
    page_num: int
    crop_type: str       # "full_page" | "crop" | "annotated"
    width: int
    height: int


class CropTask(TypedDict):
    """A single crop+annotate instruction produced by the planner."""
    page_num: int
    crop_instruction: str    # for code_execution: "Crop to the gymnasium area showing all ductwork"
    annotate: bool           # whether to run nano-banana annotation after cropping
    annotation_prompt: str   # for nano-banana: "Draw red numbered boxes around each diffuser"
    label: str               # human-readable label for the resulting image
    priority: int            # execution order (legends first = 0, detail = 1)


class PageMetadata(TypedDict):
    """Structured description of a single PDF page, generated by background metadata extraction."""
    page_num: int
    sheet_id: str
    sheet_title: str
    discipline: str
    page_type: str
    description: str
    key_elements: list[str]
    spatial_coverage: str


class DrawingReaderState(TypedDict):
    # -- Conversation --
    messages: Annotated[list, add_messages]
    question: str

    # -- Document (set at ingest time) --
    pdf_path: str
    page_image_dir: str           # directory with pre-rendered PNGs
    num_pages: int

    # -- Page Metadata (optional, for metadata-based planning) --
    page_metadata_json: str       # JSON string of per-page metadata; empty = not available

    # -- Page Selection (from planner) --
    legend_pages: list[int]
    target_pages: list[int]

    # -- Crop Planning --
    crop_tasks: list[CropTask]

    # -- Images (additive reducer — new refs are appended, never replaced) --
    image_refs: Annotated[list[ImageRef], operator.add]

    # -- Analysis --
    gemini_analysis: str
    gpt_analysis: str
    final_answer: str

    # -- Control --
    needs_more_investigation: bool
    investigation_round: int
    max_rounds: int
    enable_consensus: bool
    enable_annotation: bool
    status_message: str