from typing import Annotated, Any, TypedDict, Optional, List
from pydantic import BaseModel, Field
from langchain_core.messages import BaseMessage
from langgraph.graph.message import add_messages

# ==============================================================================
# SEARCH MODELS
# ==============================================================================

class LiteralSearchArgs(BaseModel):
    raw_text: str = ""

class InferredSearchArgs(BaseModel):
    product_line_vn: List[str] = Field(default_factory=list)
    gender_target: Optional[str] = None
    age_group: Optional[str] = None
    master_color: Optional[str] = None
    tags: List[str] = Field(default_factory=list)
    keywords: List[str] = Field(default_factory=list)
    price_min: Optional[int] = None
    price_max: Optional[int] = None

class ClassifierSearchArgs(BaseModel):
    literal: LiteralSearchArgs = Field(default_factory=LiteralSearchArgs)
    inferred: InferredSearchArgs = Field(default_factory=InferredSearchArgs)
    magento_ref_code: Optional[str] = None

# ==============================================================================
# NODE OUTPUT MODELS
# ==============================================================================

class ClassifierOutput(BaseModel):
    reasoning: str
    tool_name: Optional[str] = None
    search_args: Optional[ClassifierSearchArgs] = None
    tool_args: Optional[dict] = None
    ai_response: Optional[str] = None
    product_ids: List[str] = Field(default_factory=list)

class StylistProInsight(BaseModel):
    USER: Optional[str] = "Chưa rõ"
    TARGET: Optional[str] = "Chưa rõ"
    GOAL: Optional[str] = "Chưa rõ"
    CONSTRAINTS: Optional[str] = None
    STAGE: Optional[str] = "BROWSE"
    STAGE_NUM: int = 1
    TONE: Optional[str] = "Friendly"
    LATEST_PRODUCT_INTEREST: Optional[str] = ""
    SUMMARY_HISTORY: Optional[str] = ""
    BEHAVIORAL_HINTS: List[str] = Field(default_factory=list)

class StylistOutput(BaseModel):
    ai_response: str
    product_ids: List[str] = Field(default_factory=list)
    user_insight: StylistProInsight

# ==============================================================================
# STATE DEFINITION
# ==============================================================================

class StylistProState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    user_insight: Optional[str]
    tool_name_used: Optional[str]
    tool_args: Optional[dict] # Added for 3-node flow
    tool_result: Optional[str]
    updated_insight: Optional[dict]
    lead_stage: Optional[dict]
    stage_injection: Optional[str]
    product_ids: List[str]
    early_exit: Optional[bool]
    diagnostics: Annotated[list[dict[str, Any]], lambda x, y: x + y]
