"""
models.py — Pydantic request/response schemas for Space 2.

Changes vs previous version:
  - ChunkSchema: added `outline_path` and `source_block_id` fields so the
    enriched payload produced by the new serialiser.py round-trips cleanly.
    Both are Optional with defaults so older payloads (fields absent) still
    validate without error.

  - IngestPayload: added all enriched fields that serialiser.py now emits:
      block_to_chunks      { block_id: [chunk_id, ...] }
      outline              nested section tree
      chunk_outline_path   { chunk_id: [h1_title, h2_title, ...] }
      content_type_index   { "paragraph": [...], "table": [...], "figure": [...] }
      page_index           { "1": [block_id, ...], "2": [...] }
    All are Optional[...] = None so old single-field payloads keep working.
    `graph` is also Optional — serialiser.py supports include_graph=False for
    lightweight payloads.
"""
from __future__ import annotations

from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field


# ── Ingest (Space 1 → Space 2) ────────────────────────────────────────────────

class ChunkSchema(BaseModel):
    chunk_id:        str
    doc_name:        str
    page_num:        int
    region_type:     str
    text:            str
    section_title:   str
    bbox:            List[float]
    confidence:      float
    char_count:      int
    table_html:      Optional[str]        = None
    figure_path:     Optional[str]        = None
    source_block_id: Optional[str]        = None
    outline_path:    Optional[List[str]]  = None   # [h1_title, h2_title, ...]
    # Space 1 FIX #1: carry-forward context is now a separate field, NOT in text.
    # chunk.text is always clean; context is the previous sentence for
    # display / re-ranking only. Never embed this field.
    context:         Optional[str]        = None


class IngestPayload(BaseModel):
    doc_name:            str
    total_pages:         int
    reading_order:       List[str]
    chunks:              List[ChunkSchema]

    # Node-link graph — Optional because serialiser supports include_graph=False
    graph:               Optional[Dict[str, Any]]  = None

    # Enriched indexes added by the new serialiser.py (Optional for back-compat)
    block_to_chunks:     Optional[Dict[str, List[str]]]       = None
    outline:             Optional[List[Dict[str, Any]]]        = None
    chunk_outline_path:  Optional[Dict[str, List[str]]]        = None
    content_type_index:  Optional[Dict[str, List[str]]]        = None
    page_index:          Optional[Dict[str, List[str]]]        = None


class IngestResponse(BaseModel):
    doc_name:       str
    chunks_indexed: int
    message:        str


# ── Query (caller → Space 2) ──────────────────────────────────────────────────

class QueryRequest(BaseModel):
    question:   str             = Field(..., min_length=3)
    doc_name:   Optional[str]   = Field(
        None,
        description="Restrict retrieval to a single ingested document. "
                    "If omitted, all documents are searched.",
    )
    top_k:      Optional[int]   = Field(
        None,
        description="Override the default RETRIEVAL_TOP_K for this request.",
        ge=1, le=20,
    )


class RetrievedChunk(BaseModel):
    chunk_id:       str
    doc_name:       str
    page_num:       int
    region_type:    str
    section_title:  str
    text:           str
    score:          float   # final score after graph re-ranking
    table_html:     Optional[str]       = None
    outline_path:   Optional[List[str]] = None  # full breadcrumb for generator context headers


class RAGResponse(BaseModel):
    question:           str
    answer:             str
    retrieved_chunks:   List[RetrievedChunk]
    doc_names_searched: List[str]


# ── Health ────────────────────────────────────────────────────────────────────

class HealthResponse(BaseModel):
    status:         str
    docs_ingested:  int
    embedder_ready: bool
    llm_ready:      bool