"""
HybridRAG models for multi-source search with RRF fusion.

Combines vector search (Qdrant), knowledge graph (Neo4j), and web search (SearXNG)
with Reciprocal Rank Fusion and LLM re-ranking.
"""

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


class HybridRAGConfig(BaseModel):
    """Configuration for HybridRAG query."""
    vector_limit: int = Field(default=10, ge=1, le=50, description="Max vector results")
    graph_limit: int = Field(default=10, ge=1, le=50, description="Max graph results")
    web_limit: int = Field(default=5, ge=1, le=20, description="Max web results")
    volatile_limit: int = Field(default=1, ge=1, le=5, description="Max volatile results (typically 1)")
    document_limit: int = Field(default=5, ge=1, le=20, description="Max Paperless document results")
    enable_vector: bool = Field(default=True, description="Enable vector search")
    enable_graph: bool = Field(default=True, description="Enable graph search")
    enable_web: bool = Field(default=True, description="Enable web search")
    enable_volatile: bool = Field(default=True, description="Enable volatile cache search")
    enable_documents: bool = Field(default=True, description="Enable Paperless document search")
    enable_reranking: bool = Field(default=True, description="Enable LLM re-ranking")
    enable_enrichment: bool = Field(default=True, description="Enable graph enrichment")
    final_result_count: int = Field(default=10, ge=1, le=50, description="Final results to return")
    rrf_k: int = Field(default=60, ge=1, le=100, description="RRF constant")
    volatile_threshold: float = Field(default=0.8, ge=0.5, le=1.0, description="Volatile similarity threshold")
    document_threshold: float = Field(default=0.6, ge=0.3, le=1.0, description="Document similarity threshold")


class RelatedDossier(BaseModel):
    """Related document metadata from graph enrichment."""
    page_id: int
    title: str
    path: str
    tag: str
    shared_entities: int


class HybridRAGResult(BaseModel):
    """Single result from HybridRAG query."""
    source_type: str = Field(..., description="Source: 'wiki', 'web', 'volatile', 'document'")
    title: str
    content: str
    url: Optional[str] = Field(None, description="URL for web results")
    page_id: Optional[int] = Field(None, description="Page ID for wiki results")
    page_path: Optional[str] = Field(None, description="Wiki page path")
    paperless_id: Optional[int] = Field(None, description="Paperless document ID")
    rrf_score: float = Field(..., description="Reciprocal Rank Fusion score")
    final_rank: int = Field(..., description="Final rank after re-ranking")
    sources: List[str] = Field(..., description="Which sources included this result")
    related_dossiers: List[RelatedDossier] = Field(default=[], description="Related documents via shared entities")
    metadata: Dict[str, Any] = Field(default={}, description="Additional metadata")


class TimingBreakdown(BaseModel):
    """Performance timing breakdown for each phase."""
    query_enhancement_ms: float = Field(..., description="Phase 0: Keyword/synonym extraction")
    vector_ms: float = Field(..., description="Phase 1: Vector search")
    graph_ms: float = Field(..., description="Phase 1: Graph search")
    web_ms: float = Field(..., description="Phase 1: Web search")
    volatile_ms: float = Field(default=0, description="Phase 1: Volatile cache search")
    document_ms: float = Field(default=0, description="Phase 1: Paperless document search")
    fusion_ms: float = Field(..., description="Phase 2: RRF fusion")
    enrichment_ms: float = Field(..., description="Phase 3: Graph enrichment")
    reranking_ms: float = Field(..., description="Phase 4: LLM re-ranking")
    persistence_ms: float = Field(..., description="Phase 6: Search persistence")
    total_ms: float = Field(..., description="Total end-to-end time")


class KeywordExtraction(BaseModel):
    """Extracted keywords and synonyms from query enhancement."""
    core_keywords: List[str] = Field(default=[], description="Primary keywords")
    entities: List[str] = Field(default=[], description="Named entities")
    synonyms: Dict[str, List[str]] = Field(default={}, description="Synonyms map")
    expansions: Dict[str, List[str]] = Field(default={}, description="Abbreviation expansions")


class HybridRAGResponse(BaseModel):
    """Response from HybridRAG query."""
    query: str = Field(..., description="Original search query")
    keywords: KeywordExtraction = Field(..., description="Extracted keywords/synonyms")
    results: List[HybridRAGResult] = Field(..., description="Ranked search results")
    context: str = Field(..., description="Formatted context for LLM consumption")
    source_counts: Dict[str, int] = Field(..., description="Result counts by source")
    total_results: int = Field(..., description="Total number of results")
    timing: TimingBreakdown = Field(..., description="Performance breakdown")
    config_used: HybridRAGConfig = Field(..., description="Configuration used")
    search_id: Optional[str] = Field(None, description="Search ID for Librarian tracking")


class HybridRAGRequest(BaseModel):
    """Request for HybridRAG query."""
    query: str = Field(..., min_length=1, max_length=500, description="Search query")
    config: Optional[HybridRAGConfig] = Field(None, description="Custom configuration")