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):
    CREATE_DASHBOARD = auto()
    OTHER = auto()


class IntentClassificationState(BaseModel):
    intent: IntentEnum
    analysis: str
    message_to_user: str = Field("")
    suggestions: list[str] = Field(default_factory=list)


class ChartSuggestion(BaseModel):
    chart_title: str
    chart_type: str
    metrics: list[str] = Field(default_factory=list)
    dimensions: list[str] = Field(default_factory=list)
    description: str = Field("")


class DashboardSuggestChartsState(BaseModel):
    dashboard_reasoning: str = Field("")
    charts: list[ChartSuggestion] = Field(default_factory=list)


class DashboardState(BaseModel):
    user_input: str
    database_schemas: str

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

    intent_classification: Optional[IntentClassificationState] = None
    dashboard_suggest_charts: Optional[DashboardSuggestChartsState] = None
