"""
Data models for search functionality.

This module defines Pydantic models for search operations including search tasks
and results used in the search_indicators functionality.
"""

from enum import Enum
from typing import Literal

from pydantic import BaseModel, Field


class NodeInfo(BaseModel):
    """Represents node information with name and types."""

    model_config = {"populate_by_name": True}

    name: str = Field(description="Human-readable name of the node")
    type_of: list[str] = Field(
        description="All types of the node (e.g., ['State', 'AdministrativeArea1'])",
        alias="typeOf",
    )


class ResolvedPlace(BaseModel):
    """Represents a place that has been successfully resolved to a DCID."""

    model_config = {"populate_by_name": True}

    dcid: str = Field(description="The resolved DCID of the place")
    name: str = Field(description="Human-readable name of the node")
    type_of: list[str] = Field(
        description="All types of the node (e.g., ['State', 'AdministrativeArea1'])",
        alias="typeOf",
    )


class SearchMode(str, Enum):
    """Enumeration of search modes for the search_indicators tool."""

    BROWSE = "browse"
    LOOKUP = "lookup"


# Type alias for use in function signatures and validation
SearchModeType = Literal["browse", "lookup"]


class SearchTask(BaseModel):
    """Represents a single search task with query and place filters."""

    query: str = Field(description="The search query string")
    place_dcids: list[str] = Field(
        default_factory=list, description="List of place DCIDs to filter by"
    )


class SearchIndicator(BaseModel):
    """Base model for a search indicator, which can be a topic or a variable."""

    dcid: str
    description: str | None = None
    alternate_descriptions: list[str] | None = Field(
        None, description="Alternate descriptions or matched sentences."
    )


class SearchVariable(SearchIndicator):
    """Represents a variable object in search results."""

    places_with_data: list[str] = Field(
        default_factory=list, description="Place DCIDs where data exists"
    )


class SearchTopic(SearchIndicator):
    """Represents a topic object in search results."""

    member_topics: list[str] = Field(
        default_factory=list, description="Direct member topic DCIDs"
    )
    member_variables: list[str] = Field(
        default_factory=list, description="Direct member variable DCIDs"
    )
    places_with_data: list[str] | None = Field(
        None,
        description="Place DCIDs where data exists (if place filtering was performed)",
    )


class SearchResult(BaseModel):
    """Represents intermediate results from DCClient."""

    topics: dict[str, SearchTopic] = Field(
        default_factory=dict, description="Map of topic DCID to SearchTopic"
    )
    variables: dict[str, SearchVariable] = Field(
        default_factory=dict, description="Map of variable DCID to SearchVariable"
    )


class SearchResponse(BaseModel):
    """Unified response model for search operations."""

    topics: list[SearchTopic] | None = Field(
        None, description="List of topic objects (browse mode only)"
    )
    variables: list[SearchVariable] = Field(
        description="List of variable objects with dcid and places_with_data"
    )
    dcid_name_mappings: dict[str, str] = Field(
        default_factory=dict, description="DCID to human-readable name mappings"
    )
    dcid_place_type_mappings: dict[str, list[str]] = Field(
        default_factory=dict, description="Place DCID to type mappings"
    )
    resolved_parent_place: ResolvedPlace | None = Field(
        None,
        description="The resolved node information for the parent place, if one was provided.",
    )

    status: str = Field(default="SUCCESS", description="Status of the search operation")
