src/document_intelligence/chunks/models.py

"""
Core Data Models for Document Intelligence

Comprehensive Pydantic models for:
- Bounding boxes and spatial data
- Document chunks (text, table, chart, form fields)
- Evidence references for grounding
- Parse results and document metadata

Design principles:
- Vision-first: treat documents as visual objects
- Grounding: every extraction has evidence pointers
- Stable IDs: reproducible, hash-based chunk identifiers
- Schema-compatible: JSON export/import, Pydantic validation
"""

from __future__ import annotations

import hashlib
import json
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

from pydantic import BaseModel, Field, field_validator, model_validator


# =============================================================================
# Bounding Box Models
# =============================================================================

class BoundingBox(BaseModel):
    """
    Bounding box in XYXY format (x_min, y_min, x_max, y_max).

    Supports both pixel coordinates and normalized (0-1) coordinates.
    All spatial grounding uses this as the standard format.
    """
    x_min: float = Field(..., description="Left edge (x1)")
    y_min: float = Field(..., description="Top edge (y1)")
    x_max: float = Field(..., description="Right edge (x2)")
    y_max: float = Field(..., description="Bottom edge (y2)")

    # Coordinate system metadata
    normalized: bool = Field(default=False, description="True if 0-1 normalized")
    page_width: Optional[int] = Field(default=None, description="Page width in pixels")
    page_height: Optional[int] = Field(default=None, description="Page height in pixels")

    @field_validator('x_max')
    @classmethod
    def validate_x_max(cls, v, info):
        if 'x_min' in info.data and v < info.data['x_min']:
            raise ValueError('x_max must be >= x_min')
        return v

    @field_validator('y_max')
    @classmethod
    def validate_y_max(cls, v, info):
        if 'y_min' in info.data and v < info.data['y_min']:
            raise ValueError('y_max must be >= y_min')
        return v

    @property
    def width(self) -> float:
        return self.x_max - self.x_min

    @property
    def height(self) -> float:
        return self.y_max - self.y_min

    @property
    def area(self) -> float:
        return self.width * self.height

    @property
    def center(self) -> Tuple[float, float]:
        return ((self.x_min + self.x_max) / 2, (self.y_min + self.y_max) / 2)

    @property
    def xyxy(self) -> Tuple[float, float, float, float]:
        """Return as (x_min, y_min, x_max, y_max)."""
        return (self.x_min, self.y_min, self.x_max, self.y_max)

    @property
    def xywh(self) -> Tuple[float, float, float, float]:
        """Return as (x, y, width, height)."""
        return (self.x_min, self.y_min, self.width, self.height)

    def to_pixel(self, width: int, height: int) -> BoundingBox:
        """Convert to pixel coordinates."""
        if not self.normalized:
            return self
        return BoundingBox(
            x_min=int(self.x_min * width),
            y_min=int(self.y_min * height),
            x_max=int(self.x_max * width),
            y_max=int(self.y_max * height),
            normalized=False,
            page_width=width,
            page_height=height,
        )

    def to_normalized(self, width: int, height: int) -> BoundingBox:
        """Convert to normalized (0-1) coordinates."""
        if self.normalized:
            return self
        return BoundingBox(
            x_min=self.x_min / width,
            y_min=self.y_min / height,
            x_max=self.x_max / width,
            y_max=self.y_max / height,
            normalized=True,
            page_width=width,
            page_height=height,
        )

    def iou(self, other: BoundingBox) -> float:
        """Calculate Intersection over Union."""
        x1 = max(self.x_min, other.x_min)
        y1 = max(self.y_min, other.y_min)
        x2 = min(self.x_max, other.x_max)
        y2 = min(self.y_max, other.y_max)

        if x2 < x1 or y2 < y1:
            return 0.0

        intersection = (x2 - x1) * (y2 - y1)
        union = self.area + other.area - intersection
        return intersection / union if union > 0 else 0.0

    def contains(self, other: BoundingBox) -> bool:
        """Check if this bbox fully contains another."""
        return (
            self.x_min <= other.x_min and
            self.y_min <= other.y_min and
            self.x_max >= other.x_max and
            self.y_max >= other.y_max
        )

    def expand(self, margin: float) -> BoundingBox:
        """Expand bbox by margin pixels."""
        return BoundingBox(
            x_min=max(0, self.x_min - margin),
            y_min=max(0, self.y_min - margin),
            x_max=self.x_max + margin,
            y_max=self.y_max + margin,
            normalized=self.normalized,
            page_width=self.page_width,
            page_height=self.page_height,
        )

    def clip(self, max_width: float, max_height: float) -> BoundingBox:
        """Clip bbox to image boundaries."""
        return BoundingBox(
            x_min=max(0, self.x_min),
            y_min=max(0, self.y_min),
            x_max=min(max_width, self.x_max),
            y_max=min(max_height, self.y_max),
            normalized=self.normalized,
            page_width=self.page_width,
            page_height=self.page_height,
        )

    @classmethod
    def from_xyxy(cls, xyxy: Tuple[float, float, float, float], **kwargs) -> BoundingBox:
        """Create from (x_min, y_min, x_max, y_max) tuple."""
        return cls(x_min=xyxy[0], y_min=xyxy[1], x_max=xyxy[2], y_max=xyxy[3], **kwargs)

    @classmethod
    def from_xywh(cls, xywh: Tuple[float, float, float, float], **kwargs) -> BoundingBox:
        """Create from (x, y, width, height) tuple."""
        x, y, w, h = xywh
        return cls(x_min=x, y_min=y, x_max=x + w, y_max=y + h, **kwargs)

    def __hash__(self):
        return hash((self.x_min, self.y_min, self.x_max, self.y_max))


# =============================================================================
# Chunk Type Enumerations
# =============================================================================

class ChunkType(str, Enum):
    """
    Semantic chunk types for document segmentation.

    Covers text, tables, figures, charts, forms, and structural elements.
    Used for routing chunks to specialized extraction logic.
    """
    # Text types
    TEXT = "text"
    TITLE = "title"
    HEADING = "heading"
    PARAGRAPH = "paragraph"
    LIST = "list"
    LIST_ITEM = "list_item"

    # Structured types
    TABLE = "table"
    TABLE_CELL = "table_cell"
    FIGURE = "figure"
    CHART = "chart"
    FORMULA = "formula"
    CODE = "code"

    # Form types
    FORM_FIELD = "form_field"
    CHECKBOX = "checkbox"
    SIGNATURE = "signature"
    STAMP = "stamp"
    HANDWRITING = "handwriting"

    # Document structure
    HEADER = "header"
    FOOTER = "footer"
    PAGE_NUMBER = "page_number"
    CAPTION = "caption"
    FOOTNOTE = "footnote"
    WATERMARK = "watermark"
    LOGO = "logo"

    # Metadata
    METADATA = "metadata"
    UNKNOWN = "unknown"


class ConfidenceLevel(str, Enum):
    """Confidence level classification."""
    HIGH = "high"       # >= 0.9
    MEDIUM = "medium"   # >= 0.7
    LOW = "low"         # >= 0.5
    VERY_LOW = "very_low"  # < 0.5

    @classmethod
    def from_score(cls, score: float) -> ConfidenceLevel:
        if score >= 0.9:
            return cls.HIGH
        elif score >= 0.7:
            return cls.MEDIUM
        elif score >= 0.5:
            return cls.LOW
        else:
            return cls.VERY_LOW


# =============================================================================
# Core Document Chunk
# =============================================================================

class DocumentChunk(BaseModel):
    """
    Base document chunk with text and grounding evidence.

    This is the fundamental unit for retrieval and extraction.
    Every chunk has:
    - Stable, reproducible chunk_id (hash-based)
    - Precise spatial grounding (page, bbox)
    - Confidence score for quality assessment
    """
    # Identity
    chunk_id: str = Field(..., description="Unique, stable chunk identifier")
    doc_id: str = Field(..., description="Parent document identifier")

    # Content
    chunk_type: ChunkType = Field(..., description="Semantic type")
    text: str = Field(..., description="Text content")

    # Spatial grounding
    page: int = Field(..., ge=0, description="Zero-indexed page number")
    bbox: BoundingBox = Field(..., description="Bounding box on page")

    # Quality metrics
    confidence: float = Field(default=1.0, ge=0.0, le=1.0, description="Extraction confidence")

    # Reading order
    sequence_index: int = Field(default=0, ge=0, description="Position in reading order")

    # Source tracking
    source_path: Optional[str] = Field(default=None, description="Original file path")

    # Relationships
    parent_id: Optional[str] = Field(default=None, description="Parent chunk ID")
    children_ids: List[str] = Field(default_factory=list, description="Child chunk IDs")

    # Associated content
    caption: Optional[str] = Field(default=None, description="Caption if applicable")

    # Warnings and quality issues
    warnings: List[str] = Field(default_factory=list, description="Quality warnings")

    # Additional metadata
    extra: Dict[str, Any] = Field(default_factory=dict, description="Type-specific metadata")

    # Optional embedding (populated during indexing)
    embedding: Optional[List[float]] = Field(default=None, exclude=True)

    @property
    def confidence_level(self) -> ConfidenceLevel:
        return ConfidenceLevel.from_score(self.confidence)

    @property
    def needs_review(self) -> bool:
        """Check if chunk needs human review."""
        return self.confidence < 0.7 or len(self.warnings) > 0

    def content_hash(self) -> str:
        """Generate hash of chunk content for deduplication."""
        content = f"{self.doc_id}:{self.page}:{self.chunk_type.value}:{self.text[:200]}"
        return hashlib.sha256(content.encode()).hexdigest()[:16]

    @staticmethod
    def generate_chunk_id(
        doc_id: str,
        page: int,
        bbox: BoundingBox,
        chunk_type: ChunkType,
    ) -> str:
        """
        Generate a stable, reproducible chunk ID.

        Uses hash of (doc_id, page, bbox, type) for reproducibility.
        """
        bbox_str = f"{bbox.x_min:.2f},{bbox.y_min:.2f},{bbox.x_max:.2f},{bbox.y_max:.2f}"
        content = f"{doc_id}:p{page}:{bbox_str}:{chunk_type.value}"
        return hashlib.sha256(content.encode()).hexdigest()[:16]

    def to_retrieval_metadata(self) -> Dict[str, Any]:
        """Convert to metadata dict for vector store."""
        return {
            "chunk_id": self.chunk_id,
            "doc_id": self.doc_id,
            "chunk_type": self.chunk_type.value,
            "page": self.page,
            "bbox_xyxy": list(self.bbox.xyxy),
            "confidence": self.confidence,
            "sequence_index": self.sequence_index,
            "source_path": self.source_path,
        }

    def __hash__(self):
        return hash(self.chunk_id)


# =============================================================================
# Specialized Chunk Types
# =============================================================================

class TableCell(BaseModel):
    """A single cell in a table."""
    cell_id: str = Field(..., description="Unique cell identifier")
    row: int = Field(..., ge=0, description="Row index (0-based)")
    col: int = Field(..., ge=0, description="Column index (0-based)")
    text: str = Field(default="", description="Cell text content")
    bbox: Optional[BoundingBox] = Field(default=None, description="Cell bounding box")

    # Spanning
    rowspan: int = Field(default=1, ge=1, description="Number of rows spanned")
    colspan: int = Field(default=1, ge=1, description="Number of columns spanned")

    # Cell type
    is_header: bool = Field(default=False, description="Is header cell")

    confidence: float = Field(default=1.0, ge=0.0, le=1.0)


class TableChunk(DocumentChunk):
    """
    Specialized chunk for tables with structured cell data.

    Preserves row/column structure and supports merged cells.
    """
    chunk_type: ChunkType = Field(default=ChunkType.TABLE)

    # Table structure
    cells: List[TableCell] = Field(default_factory=list, description="All table cells")
    num_rows: int = Field(default=0, ge=0, description="Number of rows")
    num_cols: int = Field(default=0, ge=0, description="Number of columns")

    # Headers
    header_rows: List[int] = Field(default_factory=list, description="Header row indices")
    header_cols: List[int] = Field(default_factory=list, description="Header column indices")

    # Table metadata
    has_merged_cells: bool = Field(default=False)
    table_title: Optional[str] = Field(default=None)

    def get_cell(self, row: int, col: int) -> Optional[TableCell]:
        """Get cell at specific position."""
        for cell in self.cells:
            if cell.row == row and cell.col == col:
                return cell
            # Check merged cells
            if (cell.row <= row < cell.row + cell.rowspan and
                cell.col <= col < cell.col + cell.colspan):
                return cell
        return None

    def get_row(self, row: int) -> List[TableCell]:
        """Get all cells in a row."""
        return [c for c in self.cells if c.row == row]

    def get_column(self, col: int) -> List[TableCell]:
        """Get all cells in a column."""
        return [c for c in self.cells if c.col == col]

    def to_csv(self) -> str:
        """Export table to CSV format."""
        import io
        import csv

        output = io.StringIO()
        writer = csv.writer(output)

        for row_idx in range(self.num_rows):
            row_data = []
            for col_idx in range(self.num_cols):
                cell = self.get_cell(row_idx, col_idx)
                row_data.append(cell.text if cell else "")
            writer.writerow(row_data)

        return output.getvalue()

    def to_markdown(self) -> str:
        """Export table to Markdown format."""
        lines = []

        for row_idx in range(self.num_rows):
            row_cells = []
            for col_idx in range(self.num_cols):
                cell = self.get_cell(row_idx, col_idx)
                row_cells.append(cell.text if cell else "")
            lines.append("| " + " | ".join(row_cells) + " |")

            # Add separator after header
            if row_idx == 0 or row_idx in self.header_rows:
                lines.append("| " + " | ".join(["---"] * self.num_cols) + " |")

        return "\n".join(lines)

    def to_structured_json(self) -> Dict[str, Any]:
        """Export table to structured JSON with headers."""
        # Determine headers
        headers = []
        if self.header_rows:
            for col_idx in range(self.num_cols):
                cell = self.get_cell(self.header_rows[0], col_idx)
                headers.append(cell.text if cell else f"col_{col_idx}")
        else:
            headers = [f"col_{i}" for i in range(self.num_cols)]

        # Extract data rows
        data_start = max(self.header_rows) + 1 if self.header_rows else 0
        rows = []

        for row_idx in range(data_start, self.num_rows):
            row_dict = {}
            for col_idx, header in enumerate(headers):
                cell = self.get_cell(row_idx, col_idx)
                row_dict[header] = cell.text if cell else ""
            rows.append(row_dict)

        return {
            "headers": headers,
            "rows": rows,
            "num_rows": self.num_rows - len(self.header_rows),
            "num_cols": self.num_cols,
        }


class ChartDataPoint(BaseModel):
    """A data point in a chart."""
    label: Optional[str] = None
    value: Optional[float] = None
    category: Optional[str] = None
    series: Optional[str] = None
    confidence: float = Field(default=1.0, ge=0.0, le=1.0)


class ChartChunk(DocumentChunk):
    """
    Specialized chunk for charts/graphs with structured interpretation.

    Extracts title, axes, series, and key values from visualizations.
    """
    chunk_type: ChunkType = Field(default=ChunkType.CHART)

    # Chart metadata
    chart_type: Optional[str] = Field(default=None, description="bar, line, pie, scatter, etc.")
    title: Optional[str] = Field(default=None)

    # Axes
    x_axis_label: Optional[str] = Field(default=None)
    y_axis_label: Optional[str] = Field(default=None)
    x_axis_unit: Optional[str] = Field(default=None)
    y_axis_unit: Optional[str] = Field(default=None)

    # Data
    series_names: List[str] = Field(default_factory=list)
    data_points: List[ChartDataPoint] = Field(default_factory=list)

    # Interpretation
    key_values: Dict[str, Any] = Field(default_factory=dict, description="Key numeric values")
    trends: List[str] = Field(default_factory=list, description="Identified trends")
    summary: Optional[str] = Field(default=None, description="Natural language summary")

    def to_structured_json(self) -> Dict[str, Any]:
        """Export chart data as structured JSON."""
        return {
            "chart_type": self.chart_type,
            "title": self.title,
            "axes": {
                "x": {"label": self.x_axis_label, "unit": self.x_axis_unit},
                "y": {"label": self.y_axis_label, "unit": self.y_axis_unit},
            },
            "series": self.series_names,
            "data_points": [dp.model_dump() for dp in self.data_points],
            "key_values": self.key_values,
            "trends": self.trends,
            "summary": self.summary,
        }


class FormFieldChunk(DocumentChunk):
    """
    Specialized chunk for form fields.

    Handles text fields, checkboxes, radio buttons, signatures.
    """
    chunk_type: ChunkType = Field(default=ChunkType.FORM_FIELD)

    # Field metadata
    field_name: Optional[str] = Field(default=None, description="Field label/name")
    field_value: Optional[str] = Field(default=None, description="Extracted value")
    field_type: str = Field(default="text", description="text, checkbox, signature, date, etc.")

    # For checkboxes/radio
    is_checked: Optional[bool] = Field(default=None)
    options: List[str] = Field(default_factory=list)

    # Validation
    is_required: bool = Field(default=False)
    is_filled: bool = Field(default=False)


# =============================================================================
# Evidence References
# =============================================================================

class EvidenceRef(BaseModel):
    """
    Evidence reference for grounding extractions.

    Links every extracted value back to its source in the document.
    Required for auditability and trust.
    """
    # Source identification
    chunk_id: str = Field(..., description="Source chunk ID")
    doc_id: str = Field(..., description="Document ID")
    page: int = Field(..., ge=0, description="Page number (0-indexed)")
    bbox: BoundingBox = Field(..., description="Bounding box of evidence")

    # Evidence content
    source_type: str = Field(..., description="text, table, chart, form_field, etc.")
    snippet: str = Field(..., max_length=1000, description="Text snippet as evidence")

    # Quality
    confidence: float = Field(..., ge=0.0, le=1.0, description="Evidence confidence")

    # Optional cell reference for tables
    cell_id: Optional[str] = Field(default=None, description="Table cell ID if applicable")

    # Optional visual evidence
    crop_path: Optional[str] = Field(default=None, description="Path to cropped image")
    image_base64: Optional[str] = Field(default=None, description="Base64 encoded crop")

    # Warnings
    warnings: List[str] = Field(default_factory=list)

    @property
    def needs_review(self) -> bool:
        return self.confidence < 0.7 or len(self.warnings) > 0

    def to_citation(self, include_bbox: bool = False) -> str:
        """Format as human-readable citation."""
        citation = f"[Page {self.page + 1}, {self.source_type}]"
        if include_bbox:
            citation += f" @ ({self.bbox.x_min:.0f}, {self.bbox.y_min:.0f})"
        citation += f': "{self.snippet[:100]}..."' if len(self.snippet) > 100 else f': "{self.snippet}"'
        return citation


# =============================================================================
# Parse Results
# =============================================================================

class PageResult(BaseModel):
    """Result of parsing a single page."""
    page_num: int = Field(..., ge=0, description="Page number (0-indexed)")
    width: int = Field(..., gt=0, description="Page width in pixels")
    height: int = Field(..., gt=0, description="Page height in pixels")

    # Page content
    chunks: List[DocumentChunk] = Field(default_factory=list)
    markdown: str = Field(default="", description="Page content as Markdown")

    # Quality metrics
    ocr_confidence: Optional[float] = Field(default=None)
    layout_confidence: Optional[float] = Field(default=None)

    # Image path
    image_path: Optional[str] = Field(default=None, description="Path to rendered page image")


class ParseResult(BaseModel):
    """
    Complete result of document parsing.

    Contains all parsed content with metadata for downstream processing.
    """
    # Document identification
    doc_id: str = Field(..., description="Unique document identifier")
    source_path: str = Field(..., description="Original file path")
    filename: str = Field(..., description="Original filename")

    # File metadata
    file_type: str = Field(..., description="pdf, png, jpg, tiff, etc.")
    file_size_bytes: int = Field(default=0, ge=0)
    file_hash: Optional[str] = Field(default=None, description="SHA256 of file content")

    # Page data
    num_pages: int = Field(..., ge=1)
    pages: List[PageResult] = Field(default_factory=list)

    # Aggregated chunks (all pages)
    chunks: List[DocumentChunk] = Field(default_factory=list)

    # Full document markdown
    markdown_full: str = Field(default="", description="Full document as Markdown")
    markdown_by_page: Dict[int, str] = Field(default_factory=dict)

    # Processing metadata
    parsed_at: datetime = Field(default_factory=datetime.utcnow)
    processing_time_ms: float = Field(default=0.0)

    # Quality metrics
    avg_ocr_confidence: Optional[float] = Field(default=None)
    avg_layout_confidence: Optional[float] = Field(default=None)

    # Language detection
    detected_language: Optional[str] = Field(default=None)

    # Processing info
    models_used: Dict[str, str] = Field(default_factory=dict, description="Model name -> version")

    # Warnings and errors
    warnings: List[str] = Field(default_factory=list)
    errors: List[str] = Field(default_factory=list)

    # Additional metadata
    metadata: Dict[str, Any] = Field(default_factory=dict)

    @property
    def is_successful(self) -> bool:
        return len(self.errors) == 0 and len(self.chunks) > 0

    @property
    def has_tables(self) -> bool:
        return any(c.chunk_type == ChunkType.TABLE for c in self.chunks)

    @property
    def has_charts(self) -> bool:
        return any(c.chunk_type == ChunkType.CHART for c in self.chunks)

    def get_chunks_by_type(self, chunk_type: ChunkType) -> List[DocumentChunk]:
        return [c for c in self.chunks if c.chunk_type == chunk_type]

    def get_chunks_by_page(self, page: int) -> List[DocumentChunk]:
        return [c for c in self.chunks if c.page == page]

    def get_tables(self) -> List[TableChunk]:
        return [c for c in self.chunks if isinstance(c, TableChunk)]

    def get_charts(self) -> List[ChartChunk]:
        return [c for c in self.chunks if isinstance(c, ChartChunk)]

    def to_json(self, indent: int = 2) -> str:
        """Serialize to JSON."""
        return self.model_dump_json(indent=indent)

    @classmethod
    def from_json(cls, json_str: str) -> ParseResult:
        """Deserialize from JSON."""
        return cls.model_validate_json(json_str)

    def save(self, path: Union[str, Path]):
        """Save to JSON file."""
        Path(path).write_text(self.to_json(), encoding="utf-8")

    @classmethod
    def load(cls, path: Union[str, Path]) -> ParseResult:
        """Load from JSON file."""
        return cls.from_json(Path(path).read_text(encoding="utf-8"))


# =============================================================================
# Extraction Results
# =============================================================================

class FieldExtraction(BaseModel):
    """
    Single extracted field with evidence.
    """
    field_name: str = Field(..., description="Schema field name")
    value: Any = Field(..., description="Extracted value")
    value_type: str = Field(..., description="string, number, boolean, array, object")

    # Grounding
    evidence: List[EvidenceRef] = Field(default_factory=list)
    confidence: float = Field(default=1.0, ge=0.0, le=1.0)

    # Validation
    is_valid: bool = Field(default=True)
    validation_errors: List[str] = Field(default_factory=list)

    # Abstention
    abstained: bool = Field(default=False)
    abstain_reason: Optional[str] = Field(default=None)


class ExtractionResult(BaseModel):
    """
    Complete extraction result with data, evidence, and validation.
    """
    # Extracted data
    data: Dict[str, Any] = Field(default_factory=dict)
    fields: List[FieldExtraction] = Field(default_factory=list)

    # Grounding
    evidence: List[EvidenceRef] = Field(default_factory=list)

    # Quality
    overall_confidence: float = Field(default=1.0, ge=0.0, le=1.0)

    # Validation
    validation_passed: bool = Field(default=True)
    validation_errors: List[str] = Field(default_factory=list)
    validation_warnings: List[str] = Field(default_factory=list)

    # Abstention
    abstained_fields: List[str] = Field(default_factory=list)

    # Processing
    processing_time_ms: float = Field(default=0.0)
    model_used: Optional[str] = Field(default=None)

    @property
    def is_grounded(self) -> bool:
        """Check if all fields have evidence."""
        return all(f.evidence for f in self.fields if not f.abstained)

    @property
    def needs_review(self) -> bool:
        """Check if result needs human review."""
        return (
            self.overall_confidence < 0.7 or
            len(self.abstained_fields) > 0 or
            not self.validation_passed
        )


# =============================================================================
# Document Classification
# =============================================================================

class DocumentType(str, Enum):
    """Document type classifications."""
    INVOICE = "invoice"
    CONTRACT = "contract"
    AGREEMENT = "agreement"
    PATENT = "patent"
    RESEARCH_PAPER = "research_paper"
    REPORT = "report"
    LETTER = "letter"
    FORM = "form"
    RECEIPT = "receipt"
    BANK_STATEMENT = "bank_statement"
    TAX_DOCUMENT = "tax_document"
    ID_DOCUMENT = "id_document"
    MEDICAL_RECORD = "medical_record"
    LEGAL_DOCUMENT = "legal_document"
    TECHNICAL_SPEC = "technical_spec"
    PRESENTATION = "presentation"
    SPREADSHEET = "spreadsheet"
    EMAIL = "email"
    OTHER = "other"
    UNKNOWN = "unknown"


class ClassificationResult(BaseModel):
    """Document classification result."""
    doc_id: str
    doc_type: DocumentType
    confidence: float = Field(ge=0.0, le=1.0)

    # Alternative classifications
    alternatives: List[Tuple[DocumentType, float]] = Field(default_factory=list)

    # Evidence
    evidence: List[EvidenceRef] = Field(default_factory=list)
    reasoning: Optional[str] = Field(default=None)

    # Confidence threshold check
    is_confident: bool = Field(default=True)