from enum import StrEnum, auto
from typing import Annotated, Optional

from langchain_core.messages import AnyMessage
from langgraph.graph import add_messages
from pydantic import BaseModel, Field


class MetadataStatus(StrEnum):
    SUCCESS = auto()
    ERROR = auto()


class MetadataState(BaseModel):
    node: str
    status: MetadataStatus = Field(MetadataStatus.SUCCESS)
    error: Optional[str] = None


class IntentEnum(StrEnum):
    DIRECT_SQL = auto()
    AMBIGUOUS_SQL = auto()
    IRRELEVANT = auto()


class Clarification(BaseModel):
    ambiguity_reason: str
    suggestions: list[str] = Field(default_factory=list)


class IntentClassificationState(BaseModel):
    intent: IntentEnum = Field(IntentEnum.IRRELEVANT)
    message_to_user: Optional[str] = Field(None)
    analysis: str
    clarification: Optional[Clarification] = Field(None)


class SqlState(BaseModel):
    user_input: str
    database_schemas: str
    dialect: str

    query_sql: Optional[str] = None
    query_error: Optional[str] = None

    history_messages: Annotated[list[AnyMessage], add_messages] = Field(
        default_factory=list
    )
    metadata: Optional[MetadataState] = None

    intent_classification: Optional[IntentClassificationState] = None
