"""Pydantic models for PageIndex structured LLM outputs."""
from __future__ import annotations

from typing import Any, Optional

from pydantic import BaseModel, Field


# --- TOC Detection & Validation ---

class TocDetectionResult(BaseModel):
    """Result of checking whether a page contains a table of contents."""

    thinking: str = Field(description="Reasoning about TOC presence")
    toc_detected: str = Field(description="'yes' or 'no'")


class TocCompletionCheck(BaseModel):
    """Result of checking whether a TOC extraction is complete."""

    thinking: str = Field(description="Reasoning about completeness")
    completed: str = Field(description="'yes' or 'no'")


class PageIndexDetection(BaseModel):
    """Result of detecting page numbers within a table of contents."""

    thinking: str = Field(description="Reasoning about page number presence")
    page_index_given_in_toc: str = Field(description="'yes' or 'no'")


# --- TOC Structure ---

class TocItem(BaseModel):
    """A single entry in a parsed table of contents."""

    structure: Optional[str] = Field(
        None,
        description="Hierarchical index like '1', '1.1', '1.2.3' or None"
    )
    title: str = Field(description="Title of the section")
    page: Optional[int] = Field(
        None,
        description="Page number (when extracted from TOC with page numbers)"
    )
    physical_index: Optional[str] = Field(
        None,
        description="Physical index tag like '<physical_index_X>' or None"
    )
    start: Optional[str] = Field(
        None,
        description="'yes' or 'no' indicating if section starts in partial document"
    )


class TocJson(BaseModel):
    """Full table of contents in structured JSON form."""

    table_of_contents: list[TocItem] = Field(
        description="List of TOC entries"
    )


# --- Title Verification ---

class TitleAppearanceCheck(BaseModel):
    """Result of checking if a section title appears on a specific page."""

    thinking: str = Field(description="Reasoning about title appearance")
    answer: str = Field(description="'yes' or 'no'")


class TitleStartCheck(BaseModel):
    """Result of checking if a section starts at the beginning of a page."""

    thinking: str = Field(description="Reasoning about section start position")
    start_begin: str = Field(description="'yes' or 'no'")


# --- Physical Index Fixing ---

class PhysicalIndexFix(BaseModel):
    """Result of finding the correct physical index for a section."""

    thinking: str = Field(
        description="Reasoning about which page contains the section start"
    )
    physical_index: str = Field(
        description="Physical index tag like '<physical_index_X>'"
    )


# --- TOC Generation (no-TOC mode) ---

class GeneratedTocItem(BaseModel):
    """A TOC entry generated by LLM from document content (no-TOC mode)."""

    structure: str = Field(description="Hierarchical index like '1', '1.1'")
    title: str = Field(description="Title of the section (original from text)")
    physical_index: str = Field(
        description="Physical index tag '<physical_index_X>'"
    )


# --- Tree Search ---

class TreeSearchResult(BaseModel):
    """Result of an LLM tree search over a PageIndex tree."""

    thinking: str = Field(description="Reasoning about which nodes are relevant")
    node_list: list[str] = Field(
        description="List of node_id values that likely contain the answer"
    )


# --- PageIndex Tree Node (for serialization) ---

class PageIndexNode(BaseModel):
    """A node in the PageIndex tree structure."""

    title: str
    node_id: Optional[str] = None
    start_index: Optional[int] = None
    end_index: Optional[int] = None
    summary: Optional[str] = None
    prefix_summary: Optional[str] = None
    text: Optional[str] = None
    line_num: Optional[int] = None
    nodes: Optional[list[PageIndexNode]] = None

    model_config = {"extra": "allow"}


class PageIndexTree(BaseModel):
    """Top-level PageIndex document representation."""

    doc_name: str
    doc_description: Optional[str] = None
    structure: list[PageIndexNode] = Field(default_factory=list)


# --- Document Description ---

class DocDescription(BaseModel):
    """A one-sentence document description."""

    description: str = Field(
        description="One-sentence description distinguishing this document"
    )
