src/models.py

from __future__ import annotations

from typing import Dict, List, Optional

from pydantic import BaseModel, Field, field_validator


class UserRequest(BaseModel):
    """Validated user request for planning a Qiddiya visit."""

    visit_date: str = Field(
        ...,
        description="Date of visit in YYYY-MM-DD format.",
        examples=["2026-02-15"],
    )
    start_time: str = Field(
        ...,
        description="Start time in HH:MM 24h format.",
        examples=["10:00"],
    )
    end_time: str = Field(
        ...,
        description="End time in HH:MM 24h format.",
        examples=["20:00"],
    )
    must_do_attractions: List[str] = Field(
        default_factory=list,
        description="List of must-do attraction names.",
    )
    intensity_preference: int = Field(
        3,
        ge=1,
        le=5,
        description="1=relaxed, 5=extreme thrills.",
    )
    walking_tolerance: int = Field(
        3,
        ge=1,
        le=5,
        description="1=min walking, 5=okay with longer walks.",
    )

    @field_validator("visit_date")
    @classmethod
    def validate_date(cls, v: str) -> str:
        parts = v.split("-")
        if len(parts) != 3:
            raise ValueError("visit_date must be in YYYY-MM-DD format")
        year, month, day = parts
        if len(year) != 4 or len(month) != 2 or len(day) != 2:
            raise ValueError("visit_date must be in YYYY-MM-DD format")
        return v

    @field_validator("start_time", "end_time")
    @classmethod
    def validate_time(cls, v: str) -> str:
        parts = v.split(":")
        if len(parts) != 2:
            raise ValueError("time must be in HH:MM format")
        hour, minute = parts
        if not hour.isdigit() or not minute.isdigit():
            raise ValueError("time must be numeric HH:MM")
        h = int(hour)
        m = int(minute)
        if not (0 <= h <= 23 and 0 <= m <= 59):
            raise ValueError("time must be a valid 24h time")
        return v


class Attraction(BaseModel):
    id: str
    name: str
    node_id: str
    thrill_level: int = Field(ge=1, le=5)
    family_friendly: bool
    average_duration_minutes: int = Field(ge=5, le=240)
    base_popularity: int = Field(ge=1, le=10)


class ItineraryStop(BaseModel):
    attraction_id: str
    attraction_name: str
    node_id: str
    start_time: str
    end_time: str
    estimated_wait_minutes: int = Field(ge=0)
    walking_distance_m: int = Field(ge=0)
    is_suggested: bool = Field(default=False, description="True if added by system to boost enjoyment, not selected by user.")


class ItineraryPlan(BaseModel):
    visit_date: str
    total_wait_minutes: int
    total_walking_m: int
    coverage_score: float = Field(ge=0.0)
    enjoyment_score: float = Field(ge=0.0)
    stops: List[ItineraryStop]
    logs: List[str] = Field(default_factory=list)


class EvaluationResult(BaseModel):
    visit_date: str
    request_summary: str
    total_wait_minutes: int
    total_walking_m: int
    estimated_fatigue_score: float = Field(ge=0.0)
    must_do_coverage_ratio: float = Field(ge=0.0, le=1.0)
    constraint_violations: List[str] = Field(default_factory=list)
    notes: Optional[str] = None


class SystemResponse(BaseModel):
    """Debug view of the multi-agent run (logs, reflection, critique)."""

    logs: List[str] = Field(default_factory=list)
    reflection_round: int = 0
    critique: str = ""
    wait_time_forecast: Optional[Dict[str, int]] = None


class PlanResponse(BaseModel):
    """Plan plus system debug info for the UI."""

    plan: ItineraryPlan
    system: SystemResponse