"""
Pydantic models for the Helix API request/response schemas.
"""

from __future__ import annotations

from datetime import datetime
from typing import Any, Literal

from pydantic import BaseModel, Field


# ============================================================
# Query Schemas
# ============================================================

class QueryRequest(BaseModel):
    """Standard RAG query request."""
    query: str = Field(..., description="The question to ask.")
    mode: Literal["local", "global", "hybrid", "naive", "mix"] = Field(
        default="hybrid",
        description="Retrieval mode. 'hybrid' combines local entity search with global community summaries.",
    )
    top_k: int = Field(default=10, ge=1, le=100, description="Number of results to retrieve.")
    stream: bool = Field(default=False, description="Whether to stream the response (SSE).")


class TemporalQueryRequest(QueryRequest):
    """Query with temporal (time-aware) parameters."""
    valid_at: datetime | None = Field(
        default=None,
        description="Point-in-time for historical queries. E.g. '2024-06-01T00:00:00Z' to ask 'what was true on June 1 2024?'",
    )
    valid_from: datetime | None = Field(
        default=None,
        description="Start of time range filter.",
    )
    valid_to: datetime | None = Field(
        default=None,
        description="End of time range filter.",
    )


class MultiHopQueryRequest(QueryRequest):
    """Query for multi-hop reasoning across the knowledge graph."""
    max_hops: int = Field(default=3, ge=1, le=6, description="Maximum graph traversal depth.")
    max_paths: int = Field(default=10, ge=1, le=50, description="Maximum reasoning paths to explore.")
    seed_entities: list[str] | None = Field(
        default=None,
        description="Starting entity names. Auto-extracted from query if omitted.",
    )


class QueryResponse(BaseModel):
    """Standard query response."""
    answer: str
    query: str
    mode: str
    temporal_context: list[dict[str, Any]] | None = None
    valid_at: str | None = None


class MultiHopQueryResponse(BaseModel):
    """Multi-hop query response with reasoning paths."""
    answer: str
    query: str
    mode: str
    paths: list[dict[str, Any]] = Field(default_factory=list, description="Reasoning paths through the KG.")
    seed_entities: list[str] = Field(default_factory=list)


# ============================================================
# Insert Schemas
# ============================================================

class InsertRequest(BaseModel):
    """Standard document insertion."""
    document: str = Field(..., description="Text content to insert into the RAG pipeline.")
    doc_id: str | None = Field(default=None, description="Optional document ID. Auto-generated if omitted.")


class InsertWithMetadataRequest(InsertRequest):
    """Document insertion with temporal metadata and grouping."""
    source_description: str = Field(
        default="document",
        description="Description of the document source (e.g. 'annual_report', 'news_article').",
    )
    reference_time: datetime | None = Field(
        default=None,
        description="When the facts in the document were true (event time). Uses current time if omitted.",
    )
    group_id: str = Field(
        default="default",
        description="Group identifier for organizing related documents.",
    )


class InsertBatchRequest(BaseModel):
    """Batch insertion of multiple documents."""
    documents: list[InsertWithMetadataRequest] = Field(
        ..., description="List of documents to insert.", min_length=1, max_length=100,
    )


class InsertResponse(BaseModel):
    """Response from document insertion."""
    doc_id: str
    status: str = "success"
    entities_extracted: int = 0
    relations_extracted: int = 0
    entity_names: list[str] = Field(default_factory=list)
    message: str = ""


# ============================================================
# Validation Schemas
# ============================================================

class ValidateRequest(BaseModel):
    """Request to validate an LLM response against the knowledge graph."""
    response: str = Field(..., description="The LLM-generated answer to validate.")
    query: str = Field(..., description="The original query that produced the response.")
    context: dict[str, Any] | None = Field(
        default=None,
        description="Optional retrieval context used during generation.",
    )
    grounding_threshold: float = Field(
        default=0.7, ge=0.0, le=1.0,
        description="Minimum CFI score to consider the response 'grounded'.",
    )


class ValidateResponse(BaseModel):
    """Hallucination validation result."""
    is_grounded: bool = Field(description="Whether the response is grounded in the knowledge graph.")
    confidence_score: float = Field(description="Composite Fidelity Index (0.0–1.0).")
    entity_coverage: float
    relation_coverage: float
    temporal_consistency: float
    ungrounded_claims: list[str] = Field(default_factory=list)
    verification_details: dict[str, Any] = Field(default_factory=dict)


# ============================================================
# Health / Info Schemas
# ============================================================

class HealthResponse(BaseModel):
    """Health check response."""
    status: str = "ok"
    version: str
    components: dict[str, str] = Field(default_factory=dict)


class KnowledgeGraphRequest(BaseModel):
    """Request to explore the knowledge graph."""
    entity_name: str = Field(default="*", description="Starting entity name, or '*' for overview.")
    max_depth: int = Field(default=3, ge=1, le=10)
    max_nodes: int = Field(default=100, ge=1, le=1000)


class KnowledgeGraphResponse(BaseModel):
    """Knowledge graph exploration result."""
    nodes: list[dict[str, Any]] = Field(default_factory=list)
    edges: list[dict[str, Any]] = Field(default_factory=list)
    is_truncated: bool = False
