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, description="Dòng sản phẩm (Vd: Áo phông, Quần short).")
    gender_target: Optional[str] = Field(default=None, description="Giới tính mục tiêu (male, female, unisex).")
    age_group: Optional[str] = Field(default=None, description="Độ tuổi mục tiêu. BẮT BUỘC chỉ chọn: adult, kid, others.")
    master_color: Optional[str] = Field(default=None, description="Màu sắc chủ đạo.")
    tags: List[str] = Field(
        default_factory=list, 
        description=(
            "AI suy luận ý định khách -> chọn từ 4 TRỤC CỐ ĐỊNH (BẮT BUỘC có prefix!): "
            "Trục 1 (occ:): occ:di_lam, occ:di_choi, occ:di_tiec, occ:di_hoc, occ:mac_nha, occ:the_thao, occ:di_bien, occ:du_lich, occ:da_ngoai, occ:di_ngu. "
            "Trục 2a (wthr:): wthr:mua_he, wthr:mua_dong, wthr:giao_mua, wthr:troi_mua, wthr:troi_nang. "
            "Trục 2b (func:): func:thoang_mat, func:giu_am, func:tham_hut, func:nhanh_kho, func:chong_uv, func:can_gio. "
            "Trục 3 (style:): style:thanh_lich, style:nang_dong, style:basic, style:ca_tinh, style:de_thuong, style:tre_trung, style:toi_gian, style:smart_casual. "
            "Trục 4 (fit:): fit:oversize, fit:slim, fit:regular, fit:wide_leg, fit:cropped, fit:relaxed. "
            "KHÔNG tự nghĩ tag mới! PHẢI giữ prefix! Tối đa 3."
        )
    )
    keywords: List[str] = Field(default_factory=list, description="Từ khóa bổ trợ (chất liệu, kiểu dáng...).")
    price_min: Optional[int] = Field(default=None, description="Giá tối thiểu (VND).")
    price_max: Optional[int] = Field(default=None, description="Giá tối đa (VND).")

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[Any] = "Chưa rõ"
    TARGET: Optional[Any] = "Chưa rõ"
    STAGE: Optional[Any] = "BROWSE"
    LATEST_PRODUCT_INTEREST: Optional[Any] = ""
    SUMMARY_HISTORY: Optional[Any] = ""

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]
