models.py

"""
Legal Document Intelligence System - Shared Data Models

All Pydantic models used across layers. Single source of truth for data shapes.
"""

from __future__ import annotations

from datetime import datetime
from pathlib import Path
from typing import Optional

from pydantic import BaseModel, Field


# ---------------------------------------------------------------------------
# Ingestion Models
# ---------------------------------------------------------------------------

class RawDocument(BaseModel):
    """Represents raw text extracted from a single document."""
    file_path: str
    file_name: str
    raw_text: str
    pages: int
    extraction_method: str  # pdfplumber | tesseract | gemini_vision | tesseract_low_conf | none
    confidence: float
    error: Optional[str] = None

    @classmethod
    def failed(cls, file_path: str, error: str) -> "RawDocument":
        return cls(
            file_path=file_path,
            file_name=Path(file_path).name,
            raw_text="",
            pages=0,
            extraction_method="none",
            confidence=0.0,
            error=error,
        )


class LegalDocumentFields(BaseModel):
    """Structured fields extracted from a legal document via Gemini."""
    document_type: str = "unknown"
    parties: list[str] = Field(default_factory=list)
    dates: list[str] = Field(default_factory=list)
    jurisdiction: Optional[str] = None
    case_number: Optional[str] = None
    key_obligations: list[str] = Field(default_factory=list)
    key_terms: list[str] = Field(default_factory=list)
    red_flags: list[str] = Field(default_factory=list)
    summary_one_liner: str = ""
    extraction_confidence: float = 0.0

    @classmethod
    def empty(cls) -> "LegalDocumentFields":
        return cls(
            document_type="unknown",
            parties=[],
            dates=[],
            key_obligations=[],
            key_terms=[],
            red_flags=[],
            summary_one_liner="Extraction failed.",
            extraction_confidence=0.0,
        )


# ---------------------------------------------------------------------------
# Chunk Models
# ---------------------------------------------------------------------------

class DocumentChunk(BaseModel):
    """A single chunk of text with metadata for indexing."""
    chunk_id: str
    text: str
    source_file: str
    doc_type: str
    parties: list[str] = Field(default_factory=list)
    jurisdiction: Optional[str] = None
    extraction_method: str = ""
    chunk_index: int = 0
    total_chunks: int = 0


# ---------------------------------------------------------------------------
# Retrieval Models
# ---------------------------------------------------------------------------

class EvidenceChunk(BaseModel):
    """A retrieved evidence chunk with relevance score."""
    chunk_id: str
    text: str
    source_file: str
    score: float
    rrf_score: Optional[float] = None


class Citation(BaseModel):
    """A parsed citation reference from a generated draft."""
    short_id: str
    full_chunk_id: str
    chunk: Optional[EvidenceChunk] = None


# ---------------------------------------------------------------------------
# Generation Models
# ---------------------------------------------------------------------------

class Draft(BaseModel):
    """A generated legal draft with evidence and citations."""
    draft_id: str
    draft_text: str
    draft_type: str
    evidence_used: list[EvidenceChunk] = Field(default_factory=list)
    citations: list[Citation] = Field(default_factory=list)
    doc_ids: list[str] = Field(default_factory=list)
    created_at: datetime = Field(default_factory=datetime.utcnow)


# ---------------------------------------------------------------------------
# Feedback Models
# ---------------------------------------------------------------------------

class EditDiff(BaseModel):
    """Computed diff between original and edited draft text."""
    additions: list[str] = Field(default_factory=list)
    deletions: list[str] = Field(default_factory=list)
    original: str
    edited: str
    change_ratio: float


class DedupeResult(BaseModel):
    """Result of deduplication check for a new pattern."""
    action: str  # "store" | "skip" | "merge"
    similarity: Optional[float] = None
    merge_target_id: Optional[str] = None
    merged_pattern: Optional[str] = None


class LearnedPattern(BaseModel):
    """A learned editing preference stored in the pattern store."""
    id: str
    pattern: str
    doc_type: Optional[str] = None
    created_at: str = ""
    use_count: int = 0
    active: bool = True


# ---------------------------------------------------------------------------
# Ingestion Result (combines RawDocument + LegalDocumentFields + Chunks)
# ---------------------------------------------------------------------------

class IngestResult(BaseModel):
    """Full result of ingesting a single document."""
    document: RawDocument
    fields: LegalDocumentFields
    chunks: list[DocumentChunk] = Field(default_factory=list)
    chunk_count: int = 0