"""Pydantic v2 response models for the query API."""

from __future__ import annotations

from pydantic import BaseModel, Field


class PersonaInfo(BaseModel):
    name: str
    anchor_nodes: list[str]
    edge_types: list[str]
    chunk_types: list[str]


class CaseSummary(BaseModel):
    mnc: str
    court: str = ""
    date: str = ""
    jurisdiction: str = ""
    matter_type: str = ""
    verdict: str = ""


class ChunkResult(BaseModel):
    chunk_id: str
    text: str | None = None
    metadata: dict = Field(default_factory=dict)
    distance: float | None = None


class JurorContextResponse(BaseModel):
    persona: str
    case_mnc: str
    context_text: str
    source_chunk_ids: list[str]
    total_tokens: int


class CaseGraphNode(BaseModel):
    node_id: str
    label: str
    properties: dict = Field(default_factory=dict)


class CaseGraphRelationship(BaseModel):
    rel_id: str
    from_id: str
    to_id: str
    rel_type: str
    properties: dict = Field(default_factory=dict)


class CaseGraphSummary(BaseModel):
    case: CaseGraphNode
    judges: list[CaseGraphNode] = Field(default_factory=list)
    charges: list[CaseGraphNode] = Field(default_factory=list)
    rulings: list[CaseGraphNode] = Field(default_factory=list)
    witnesses: list[CaseGraphNode] = Field(default_factory=list)
    chunks: list[CaseGraphNode] = Field(default_factory=list)
    relationships: list[CaseGraphRelationship] = Field(default_factory=list)


class HybridResult(BaseModel):
    case_mnc: str
    chunks: list[ChunkResult] = Field(default_factory=list)
    juror_context: JurorContextResponse | None = None


class HealthResponse(BaseModel):
    status: str
    nodes: int
    edges: int
    cases: int
    chunks: int