"""Define API models for agent endpoints."""

from __future__ import annotations

from typing import Dict, List, Optional

from pydantic import BaseModel, Field

# Group agent transport schemas together so route modules stay focused on HTTP flow.


class AskRequest(BaseModel):
    """Define the Ask Request API model."""

    query: str = Field(..., min_length=1, max_length=2000)
    filters: Optional[str] = None
    provider: Optional[str] = None
    model: Optional[str] = None
    top_k: Optional[int] = Field(default=None, ge=1, le=20)
    prompt_template: Optional[str] = None


class AskResponse(BaseModel):
    """Define the Ask Response API model."""

    answer: str
    sources: List[Dict] = Field(default_factory=list)
    model: str = ""
    latency_ms: int = 0
    retrieved_count: int = 0
    context_tokens: int = 0
    truncated: bool = False
    error: Optional[str] = None


class ChatRequest(BaseModel):
    """Define the Chat Request API model."""

    query: str = Field(..., min_length=1, max_length=2000)
    session_id: Optional[str] = None
    filters: Optional[str] = None
    provider: Optional[str] = None
    model: Optional[str] = None
    top_k: Optional[int] = Field(default=None, ge=1, le=20)


class ChatResponse(BaseModel):
    """Define the Chat Response API model."""

    session_id: str
    answer: str
    sources: List[Dict] = Field(default_factory=list)
    model: str = ""
    latency_ms: int = 0
    message_count: int = 0


class SessionInfo(BaseModel):
    """Define the Session Info API model."""

    session_id: str
    message_count: int
    created_at: int
    updated_at: int


class FeedbackRequest(BaseModel):
    """Define the Feedback Request API model."""

    session_id: str
    message_index: int
    rating: int = Field(..., ge=1, le=5)
    comment: Optional[str] = None


class TemplateListResponse(BaseModel):
    """Define the Template List Response API model."""

    templates: Dict[str, str]