from typing import Annotated, Any, TypedDict

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

import config as global_config


class QueryRequest(BaseModel):
    """API Request model cho Fashion Q&A Chat"""

    user_id: str | None = None
    user_query: str
    images: list[str] | None = None
    image_analysis: dict[str, Any] | None = None


class AgentState(TypedDict):
    """Trạng thái của Agent trong LangGraph."""

    user_query: BaseMessage
    history: list[BaseMessage]
    user_id: str | None
    user_insight: str | None
    ai_response: BaseMessage | None
    images_embedding: list[str] | None
    messages: Annotated[list[BaseMessage], add_messages]


class AgentConfig:
    """Class chứa cấu hình runtime cho Agent."""

    def __init__(self, **kwargs):
        self.model_name = kwargs.get("model_name") or global_config.DEFAULT_MODEL
        self.openai_api_key = kwargs.get("openai_api_key")
        self.google_api_key = kwargs.get("google_api_key")
        self.groq_api_key = kwargs.get("groq_api_key")
        self.supabase_url = kwargs.get("supabase_url")
        self.supabase_key = kwargs.get("supabase_key")
        self.langfuse_public_key = kwargs.get("langfuse_public_key")
        self.langfuse_secret_key = kwargs.get("langfuse_secret_key")
        self.langfuse_base_url = kwargs.get("langfuse_base_url")


def get_config() -> AgentConfig:
    """Khởi tạo cấu hình Agent từ các biến môi trường."""
    return AgentConfig(
        model_name=global_config.DEFAULT_MODEL,
        openai_api_key=global_config.OPENAI_API_KEY,
        google_api_key=global_config.GOOGLE_API_KEY,
        groq_api_key=global_config.GROQ_API_KEY,
        supabase_url=global_config.AI_SUPABASE_URL,
        supabase_key=global_config.AI_SUPABASE_KEY,
        langfuse_public_key=global_config.LANGFUSE_PUBLIC_KEY,
        langfuse_secret_key=global_config.LANGFUSE_SECRET_KEY,
        langfuse_base_url=global_config.LANGFUSE_BASE_URL,
    )
# ==============================================================================
# STYLIST PRO MODELS (Architecture 2.0)
# ==============================================================================

class RawSearchArgs(BaseModel):
    """Lane 1: Trích xuất từ khóa gốc, giữ nguyên ý định khách, sửa lỗi typo."""
    raw_text: str = Field(default="", description="Core keywords from user message")

class StructuredSearchArgs(BaseModel):
    """Lane 2: Suy luận bộ lọc có cấu trúc từ ngữ cảnh."""
    product_line_vn: list[str] = Field(default_factory=list, description="Áo phông, Quần jean...")
    gender_target: str | None = Field(default=None, description="men/women/boy/girl/unisex")
    age_group: str | None = Field(default=None, description="adult/kid")
    master_color: str | None = Field(default=None, description="Màu sắc suy luận")
    tags: list[str] = Field(default_factory=list)
    keywords: list[str] = Field(default_factory=list)
    price_min: int | None = Field(default=None)
    price_max: int | None = Field(default=None)
    discount_min: int | None = Field(default=None)
    discovery_mode: str | None = Field(default=None)

class ClassifierSearchArgs(BaseModel):
    """Tổng hợp 2 Lane Search."""
    literal: RawSearchArgs = Field(default_factory=RawSearchArgs)
    inferred: StructuredSearchArgs = Field(default_factory=StructuredSearchArgs)
    magento_ref_code: str | None = Field(default=None)

class StylistProInsight(BaseModel):
    """Bộ nhớ khách hàng - 12 trường cốt lõi."""
    USER: str | None = Field(default="Chưa rõ")
    TARGET: str | None = Field(default="Chưa rõ")
    GOAL: str | None = Field(default="Chưa rõ")
    CONSTRAINTS: str | None = Field(default="")
    STAGE: str | None = Field(default="BROWSE")
    STAGE_NUM: int = Field(default=1)
    TONE: str | None = Field(default="Friendly")
    BEHAVIORAL_HINTS: list[str] = Field(default_factory=list)
    LATEST_PRODUCT_INTEREST: str | None = Field(default="")
    LAST_ACTION: str | None = Field(default="")
    SUMMARY_HISTORY: str | None = Field(default="")

class ClassifierOutput(BaseModel):
    """Kết quả từ node Classifier."""
    reasoning: str
    tool_name: str | None = None
    search_args: ClassifierSearchArgs | None = None
    tool_args: dict[str, Any] | None = None
    ai_response: str | None = None
    product_ids: list[str] = Field(default_factory=list)

class StylistOutput(BaseModel):
    """Kết quả từ node Stylist."""
    ai_response: str
    product_ids: list[str] = Field(default_factory=list)
    user_insight: StylistProInsight

class StylistProState(TypedDict):
    """Trạng thái luồng chạy của Stylist Pro."""
    messages: Annotated[list[BaseMessage], add_messages]
    user_insight: str | None
    tool_name_used: str | None
    tool_result: str | None
    updated_insight: dict[str, Any] | None
    lead_stage: dict[str, Any] | None
    stage_injection: str | None
    product_ids: list[str]
    early_exit: bool | None
    diagnostics: Annotated[list[dict[str, Any]], lambda x, y: x + y]
