"""
Lead Stage Graph — 2-Agent Architecture (Classifier NHẸ + Stylist NẶNG).

Flow: User → Classifier (route tool) → [Tool Exec] → Stylist (response + insight) → END

2 con AI:
  - Classifier (NHẸ): with_structured_output(ClassifierOutput) → chỉ quyết định tool
  - Stylist (NẶNG): with_structured_output(StylistOutput) → sinh response + InsightJSON
"""

import json
import logging
import re
import time
from typing import Annotated, Any, TypedDict

from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
from langchain_core.runnables import RunnableConfig
from langgraph.graph import END, StateGraph
from langgraph.graph.message import add_messages
from pydantic import BaseModel, Field

from agent.tools.add_to_cart_tool import add_to_cart
from agent.tools.brand_knowledge_tool import canifa_knowledge_search
from agent.tools.check_is_stock import check_is_stock
from agent.tools.customer_info_tool import collect_customer_info
from agent.tools.promotion_canifa_tool import canifa_get_promotions
from agent.tools.store_search_tool import canifa_store_search
from common.llm_factory import create_llm
from config import DEFAULT_MODEL

# Tools
from .lead_search_tool import lead_search_tool
from .prompts import LEAD_STAGE_CLASSIFIER_PROMPT, STYLIST_SYSTEM_PROMPT, format_insight_injection

logger = logging.getLogger(__name__)


# ═══════════════════════════════════════════════
# Pydantic Models
# ═══════════════════════════════════════════════


class LiteralSearchArgs(BaseModel):
    model_config = {"extra": "ignore"}
    raw_text: str = Field(
        default="",
        description=(
            "Câu search NGUYÊN VĂN từ user. Copy y hệt phần liên quan đến sản phẩm. "
            "VD: user nói 'tìm áo 30/4' -> raw_text='áo 30/4'. "
            "VD: user nói 'có quần jean ống rộng ko' -> raw_text='quần jean ống rộng'. "
            "KHÔNG suy luận, KHÔNG thêm bớt. Giữ nguyên ý user."
        ),
    )


class InferredSearchArgs(BaseModel):
    model_config = {"extra": "ignore"}

    product_line_vn: list[str] = Field(
        default=[],
        description=(
            "BẮT BUỘC ĐIỀN nêú xác định được loại SP. Đây là bộ lọc SQL QUAN TRỌNG NHẤT. "
            "Sử dụng các từ khóa chuẩn: 'Áo phông', 'Áo Polo', 'Áo sơ mi', 'Quần jean', 'Váy liền', 'Chân váy', 'Áo khoác', 'Áo len', 'Áo nỉ', 'Áo giữ nhiệt', 'Bộ mặc nhà', 'Quần Khaki', 'Quần dài', 'Bộ thể thao', 'Áo kiểu'..."
        ),
    )
    gender_by_product: str | None = Field(
        default=None,
        description="Suy luan tu ngu canh: women/men/boy/girl/unisex. Null neu khong co manh moi.",
    )
    age_by_product: str | None = Field(default=None, description="adult | kid | others")
    master_color: str | None = Field(
        default=None,
        description="Màu sắc AI SUY LUẬN từ ngữ cảnh. VD: '30/4' -> color='đỏ'. 'áo trắng' -> color='trắng'.",
    )
    tags: list[str] = Field(
        default=[],
        description=(
            "AI suy luận ý định khách -> chọn từ DANH SÁCH TAGS CỐ ĐỊNH: "
            "Trục 1 (Màu & Tông): Màu đen, Màu trắng, Màu xám, Màu xanh navy, Màu xanh da trời, Màu be, Màu nâu, Màu hồng, Màu đỏ, Màu vàng, Tông trung tính (Neutral), Tông màu pastel, Tông sáng (Light), Tông nổi bật (Dark). "
            "Trục 2 (Dịp): Đi làm / Công sở, Đi chơi / dạo phố, Mặc nhà, Tập thể thao, Đi tiệc, Đi biển, Du lịch, Đi học. "
            "Trục 3 (Mùa): Mùa hè, Mùa đông, Cả 2 mùa, Giao mùa. "
            "Trục 4 (Phong cách): Cơ bản (Basic), Năng động, Thanh lịch, Dễ thương, Cá tính, Tối giản, Smart Casual, Thoáng mát, Giữ ấm, Thấm hút, Nhanh khô, Chống UV, Cản gió. "
            "Trục 5 (Dáng): Oversize, Slim, Regular, Wide leg, Cropped, Relaxed. "
            "VD: 'đi tiệc' -> tags=['Đi tiệc']. 'mùa đông' -> tags=['Mùa đông', 'Giữ ấm']. "
            "KHÔNG tự nghĩ tag mới! Tối đa 3."
        ),
    )
    keywords: list[str] = Field(
        default=[],
        description=(
            "Từ khóa BỔ TRỢ AI suy luận thêm (chất liệu, tính năng...). "
            "VD: '30/4' -> keywords=['cờ đỏ sao vàng']. 'đi biển' -> keywords=['thoáng mát']. "
            "Tối đa 2."
        ),
    )
    price_min: int | None = Field(default=None)
    price_max: int | None = Field(default=None)
    discount_min: int | None = Field(default=None, description="% giam gia toi thieu. VD: hang sale -> 1")
    discount_max: int | None = Field(default=None)
    discovery_mode: str | None = Field(default=None, description="new | best_seller. Chi dung khi khach noi ro.")
    size: str | None = Field(default=None, description="XS/S/M/L/XL/XXL/3XL... Chi dien khi khach noi ro size.")


class ClassifierLeadSearchArgs(BaseModel):
    """Dual-Lane Search Args — buộc AI tách bạch rõ Literal vs Inferred."""

    model_config = {"extra": "ignore"}

    literal: LiteralSearchArgs = Field(
        default_factory=LiteralSearchArgs,
        description="Lane 1: Câu search nguyên văn từ user. KHÔNG filter, KHÔNG suy luận.",
    )
    inferred: InferredSearchArgs = Field(
        default_factory=InferredSearchArgs,
        description="Lane 2: AI suy luận filter structured.",
    )
    magento_ref_code: str | None = Field(default=None, description="Ma SKU cu the neu khach cung cap.")
    reasoning: str | None = Field(default=None, description="Giai thich ngan gon tai sao chon params nay.")


class ClassifierOutput(BaseModel):
    """Classifier output — Định tuyến hoặc Trả lời Sớm."""

    model_config = {"extra": "ignore"}

    reasoning: str = Field(description="Ly luan goi tool hay tra loi truc tiep")

    # ── TH 1: GOI TOOL ──
    tool_name: str | None = Field(
        default=None,
        description="Ten tool (lead_search_tool). Bat buoc de null neu khach chi dang tam su/chao hoi/ko can DB.",
    )
    lead_search_args: ClassifierLeadSearchArgs | None = Field(
        default=None,
        description=(
            "Args cho lead_search_tool. PHAI dien khi tool_name='lead_search_tool'. Null neu tool_name khac hoac null."
        ),
    )
    tool_args: dict | None = Field(
        default=None,
        description="Args cho cac tool khac (knowledge_search, check_is_stock, v.v.). Null neu dung lead_search_tool.",
    )

    # ── TH 2: EARLY EXIT ──
    ai_response: str | None = Field(
        default=None,
        description="Cau tra loi cho khach. CHI nha ra khi tool_name=null. Neu goi tool, bat buoc de null.",
    )
    product_ids: list[str] | None = Field(default_factory=list, description="Ma SKU (neu khach hoi trung ma).")


class InsightJSON(BaseModel):
    """Bộ nhớ khách hàng — 12 trường."""

    model_config = {"extra": "ignore"}

    USER: str = Field(default="Chưa rõ")
    TARGET: str = Field(default="Chưa rõ")
    GOAL: str = Field(default="Chưa rõ")
    CONSTRAINS: str = Field(
        default="", description="CHỈ điền khi khách NÓI RÕ. VD: khách nói 'dưới 500k' → 'Budget: <500k'. KHÔNG TỰ BỊA!"
    )
    STAGE: str = Field(default="BROWSE")
    STAGE_NUM: int = Field(default=1)
    TONE: str = Field(default="Friendly")
    BEHAVIORAL_HINTS: list[str] = Field(default_factory=list)
    LATEST_PRODUCT_INTEREST: str = Field(default="")
    LAST_ACTION: str = Field(default="")
    SUMMARY_HISTORY: str = Field(default="")


class StylistOutput(BaseModel):
    """Stylist output — NẶNG, sinh cả response + insight."""

    model_config = {"extra": "ignore"}

    ai_response: str = Field(description="Câu trả lời cho khách (max 200 từ)")
    product_ids: list[str] = Field(default_factory=list)
    user_insight: InsightJSON = Field(description="Updated user insight sau turn này")


# ═══════════════════════════════════════════════
# State
# ═══════════════════════════════════════════════


class LeadStageState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]

    # Input context
    user_insight: str | None

    # Internal pipeline
    tool_result: str | None
    tool_name_used: str | None

    # Output
    updated_insight: dict | None
    lead_stage: dict | None
    stage_injection: str | None
    product_ids: list[str]
    early_exit: bool | None

    diagnostics: Annotated[list[dict[str, Any]], lambda x, y: x + y]


# ═══════════════════════════════════════════════
# Helpers
# ═══════════════════════════════════════════════


def _extract_text(content) -> str:
    """Extract plain text từ msg.content."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                parts.append(item.get("text", ""))
            elif isinstance(item, str):
                parts.append(item)
        return "\n".join(parts) if parts else str(content)
    return str(content or "")


_SKU_REGEX = re.compile(r"\b\d[A-Z0-9]{5,}(?:-[A-Z0-9]{2,}){0,2}\b", re.IGNORECASE)


def _dedupe_preserve_order(items: list[str]) -> list[str]:
    seen: set[str] = set()
    out: list[str] = []
    for raw in items:
        val = str(raw or "").upper().strip()
        if not val or val in seen:
            continue
        seen.add(val)
        out.append(val)
    return out


def _extract_skus_from_text(text: str | None) -> list[str]:
    if not text:
        return []
    matches = _SKU_REGEX.findall(text)
    return _dedupe_preserve_order([m for m in matches if m])


# ═══════════════════════════════════════════════
# Graph Builder
# ═══════════════════════════════════════════════


class LeadStageGraph:
    """2-Agent: Classifier (NHẸ, route tool) → Stylist (NẶNG, response + insight)."""

    def __init__(self, model_name: str | None = None):
        self.model_name = model_name or DEFAULT_MODEL

        # Tool registry
        self._tool_registry = {
            "lead_search_tool": lead_search_tool,
            "canifa_knowledge_search": canifa_knowledge_search,
            "canifa_get_promotions": canifa_get_promotions,
            "canifa_store_search": canifa_store_search,
            "check_is_stock": check_is_stock,
            "collect_customer_info": collect_customer_info,
            "add_to_cart": add_to_cart,
        }

        # AI 1: Classifier — structured output bypass for Claude API (zunef gateway)
        _classifier_base = create_llm(model_name=self.model_name, streaming=False)
        self.classifier_llm = _classifier_base
        # Note: Claude models bypass structured output, parse JSON manually

        # AI 2: Stylist — structured output bypass for Claude API (zunef gateway)
        _stylist_base = create_llm(model_name=self.model_name, streaming=False)
        self.stylist_llm = _stylist_base
        # Note: Claude models bypass structured output, parse JSON manually

        self._compiled = None
        logger.info(f"✅ LeadStageGraph initialized | model: {self.model_name}")

    # ─────────────────────────────────────
    # Helper: Format history
    # ─────────────────────────────────────
    def _format_history(self, messages: list[BaseMessage], max_turns: int = 10) -> str:
        """Format recent messages cho context — giữ 10 turns, 500 chars/msg."""
        lines = []
        human_count = 0
        for msg in reversed(messages):
            if isinstance(msg, HumanMessage):
                lines.insert(0, f"Khách: {_extract_text(msg.content)[:500]}")
                human_count += 1
            elif isinstance(msg, AIMessage) and not (hasattr(msg, "tool_calls") and msg.tool_calls):
                lines.insert(0, f"Bot: {_extract_text(msg.content)[:500]}")
            if human_count >= max_turns:
                break
        return "\n".join(lines) if lines else "Chưa có lịch sử."

    # ─────────────────────────────────────
    # Helper: Format tool result cho Stylist
    # ─────────────────────────────────────
    def _format_tool_result(self, tool_result: str | None) -> str:
        if not tool_result:
            return ""
        try:
            parsed = json.loads(tool_result)
            if parsed.get("status") == "success":
                products = parsed.get("products", [])
                if products:
                    # Deduplicate by SKU to prevent duplicate context
                    dedup = {}
                    for p in products:
                        sku = str(p.get("sku_code") or p.get("sku") or p.get("internal_ref_code") or "").upper().strip()
                        if sku and sku not in dedup:
                            dedup[sku] = p
                    unique_products = list(dedup.values())

                    lines = [f"✅ Tìm thấy {len(unique_products)} sản phẩm:"]
                    for p in unique_products[:8]:
                        price = p.get("price", 0)
                        discount = p.get("discount", "")
                        line = f"- [{p.get('sku', '')}] {p.get('name', '')} | Giá: {price:,}đ"
                        if discount:
                            line += f" ({discount})"
                        if p.get("sizes"):
                            line += f" | Sizes: {p['sizes']}"
                        desc = (p.get("description") or "").strip()
                        if desc:
                            line += f"\n  📝 {desc[:200]}"

                        outfit = p.get("outfit_recommendations")
                        if outfit and isinstance(outfit, list):
                            outfit_lines = []
                            for o in outfit[:3]:
                                outfit_lines.append(
                                    f"{o.get('role', 'match')}: {o.get('match_product_name')} ({o.get('match_product_code')})"
                                )
                            if outfit_lines:
                                line += f"\n  👗 outfit_recommendations: {', '.join(outfit_lines)}"

                        lines.append(line)
                    return "\n".join(lines)
                return parsed.get("message", "Không tìm thấy sản phẩm phù hợp.")
            content = parsed.get("content") or parsed.get("message") or str(parsed)
            return str(content)[:800]
        except Exception:
            return (tool_result or "")[:500]

    # ─────────────────────────────────────
    # Node 1: CLASSIFIER — TOOL ROUTER (NHẸ)
    # ─────────────────────────────────────
    async def _classifier_node(self, state: LeadStageState, config: RunnableConfig | None = None) -> dict:
        """Classifier: Đọc query + insight → quyết định gọi tool nào. KHÔNG sinh insight."""
        messages = state["messages"]
        user_query = _extract_text(messages[-1].content) if messages else ""
        user_insight = state.get("user_insight") or "Chưa có thông tin."

        history_text = self._format_history(messages[:-1])

        # Extract previously shown products from insight
        products_context = ""
        try:
            if isinstance(user_insight, str) and user_insight.startswith("{"):
                _insight = json.loads(user_insight)
                latest = _insight.get("LATEST_PRODUCT_INTEREST", "")
                summary = _insight.get("SUMMARY_HISTORY", "")
                if latest:
                    products_context += f"\n=== SP ĐÃ GIỚI THIỆU (QUAN TRỌNG) ===\n{latest}\n"
                if summary:
                    products_context += f"\n=== TÓM TẮT HỘI THOẠI ===\n{summary}\n"
        except Exception:
            pass

        context_block = (
            f"=== TIN NHẮN HIỆN TẠI ===\n{user_query}\n\n"
            f"=== USER INSIGHT ===\n{user_insight}\n\n"
            f"{products_context}"
            f"=== LỊCH SỬ (10 cặp gần nhất) ===\n{history_text}"
        )

        classifier_messages = [
            SystemMessage(content=LEAD_STAGE_CLASSIFIER_PROMPT),
            HumanMessage(content=context_block),
        ]

        start = time.time()
        try:
            response = await self.classifier_llm.ainvoke(
                classifier_messages, config=config
            )
            # Parse JSON from text response (Claude models return text, not structured output)
            import json as json_lib
            text = response.content if hasattr(response, "content") else str(response)
            logger.info(f"[CLASSIFIER] Raw response repr: {repr(text[:500])}")
            # Extract JSON from text (may have extra text)
            text = text.strip()
            if text.startswith('```json'):
                text = text[7:]
            if text.startswith('```'):
                text = text[3:]
            if text.endswith('```'):
                text = text[:-3]
            text = text.strip()

            try:
                output = json_lib.loads(text)
            except json_lib.JSONDecodeError as je:
                logger.error(f"[CLASSIFIER] JSON parse error: {je}. Text was: {repr(text[:500])}")
                raise
            # Convert dict to object-like access
            class DictToObj:
                def __init__(self, data):
                    self.__dict__.update(data)
            output = DictToObj(output)
        except Exception as e:
            logger.error(f"❌ Classifier error: {e}")
            return {
                "tool_result": None,
                "tool_name_used": None,
                "diagnostics": [
                    {"step": "classifier_error", "label": "❌ Classifier", "content": str(e)[:300], "elapsed_ms": 0}
                ],
            }

        elapsed_ms = (time.time() - start) * 1000
        tool_name = output.tool_name
        ai_response = output.ai_response

        # ── Resolve tool_args: typed lead_search_args hoac raw dict ──
        if tool_name == "lead_search_tool" and output.lead_search_args:
            # Handle both dict and Pydantic model cases
            if isinstance(output.lead_search_args, dict):
                tool_args = output.lead_search_args
            elif hasattr(output.lead_search_args, 'model_dump'):
                tool_args = output.lead_search_args.model_dump(exclude_none=False)
            else:
                tool_args = dict(output.lead_search_args) if output.lead_search_args else {}
        elif tool_name == "lead_search_tool" and output.tool_args:
            # fallback: model dung tool_args thay vi lead_search_args
            tool_args = output.tool_args
        elif tool_name and tool_name != "lead_search_tool":
            tool_args = output.tool_args or {}
        else:
            tool_args = {}

        # ── PASS USER_INSIGHT TO LEAD_SEARCH_TOOL ──
        if tool_name == "lead_search_tool" and user_insight:
            tool_args["user_insight"] = user_insight

        # ── Log Classifier parsed output ──
        logger.info(
            "[CLASSIFIER] tool=%s | product_line=%s | keywords=%s | gender=%s | color=%s | reasoning=%s",
            tool_name,
            tool_args.get("product_line_vn") if tool_args else None,
            tool_args.get("keywords") if tool_args else None,
            tool_args.get("gender_by_product") if tool_args else None,
            tool_args.get("master_color") if tool_args else None,
            (output.reasoning or "")[:120],
        )

        # product_line_vn guard removed — keywords LIKE description_text handles category matching
        # Tool call proceeds even with empty product_line_vn

        # Build diagnostic
        diag = {
            "step": "classifier",
            "label": f"🔧 Classifier → {tool_name or 'Early Exit'}",
            "content": output.reasoning,
            "elapsed_ms": round(elapsed_ms),
        }

        # ── EARLY EXIT (KHÔNG GỌI TOOL) ──
        if not tool_name and ai_response:
            logger.info(f"⚡ Early Exit: {ai_response[:50]}...")
            diag["label"] = "⚡ Classifier (Early Exit)"
            diag["content"] = output.reasoning + f"\n[AI Response: {ai_response[:100]}]"
            return {
                "messages": [AIMessage(content=ai_response)],
                "tool_result": None,
                "tool_name_used": None,
                "product_ids": output.product_ids or [],
                "early_exit": True,
                "diagnostics": [diag],
            }

        # ── CALL TOOL ──
        if tool_name:
            diag["tool_calls"] = [{"name": tool_name, "args": tool_args}]

        diagnostics = [diag]

        # ── Execute Tool Inline ──
        tool_result = None
        if tool_name:
            tool_fn = self._tool_registry.get(tool_name)
            if tool_fn:
                try:
                    t_start = time.time()
                    tool_result = await tool_fn.ainvoke(tool_args)
                    tool_elapsed = round((time.time() - t_start) * 1000)
                    logger.info(f"📦 Tool '{tool_name}' OK | {tool_elapsed}ms")

                    # Tool result diagnostic
                    try:
                        parsed = json.loads(tool_result)
                        count = parsed.get("count", len(parsed.get("products", [])))
                        summary = f"✅ {count} sản phẩm"
                        products = parsed.get("products", [])
                        if products:
                            names = [f"• {p.get('name', '')}" for p in products[:4]]
                            summary += "\n" + "\n".join(names)
                    except Exception:
                        summary = (tool_result or "")[:200]

                    diagnostics.append(
                        {
                            "step": "tool_result",
                            "label": f"📦 {tool_name}",
                            "content": summary,
                            "raw_json": tool_result,
                            "elapsed_ms": tool_elapsed,
                        }
                    )
                except Exception as te:
                    logger.error(f"❌ Tool '{tool_name}' error: {te}")
                    tool_result = json.dumps({"status": "error", "message": str(te)})
            else:
                logger.warning(f"⚠️ Unknown tool: '{tool_name}'")

        return {
            "tool_result": tool_result,
            "tool_name_used": tool_name,
            "diagnostics": diagnostics,
        }

    # ─────────────────────────────────────
    # Node 2: STYLIST — Response + Insight (NẶNG, LUÔN CHẠY)
    # ─────────────────────────────────────
    async def _stylist_node(self, state: LeadStageState, config: RunnableConfig | None = None) -> dict:
        """Stylist: Nhận tool result + history + insight cũ → sinh response + InsightJSON mới."""
        messages = state["messages"]
        tool_result = state.get("tool_result")
        user_insight = state.get("user_insight") or "Chưa có."

        # Format insight cũ → inject vào system prompt
        try:
            old_insight = (
                json.loads(user_insight) if isinstance(user_insight, str) and user_insight.startswith("{") else {}
            )
        except Exception:
            old_insight = {}

        injection = format_insight_injection(old_insight) if old_insight else "Chưa có insight — đây là lần đầu."

        # Escape curly braces in prompt, then substitute only stage_injection
        stylist_sys = STYLIST_SYSTEM_PROMPT.replace("{", "{{").replace("}", "}}").replace("{{stage_injection}}", "{stage_injection}").format(stage_injection=injection)

        # Build messages cho Stylist
        stylist_messages: list[BaseMessage] = [SystemMessage(content=stylist_sys)]

        # Recent history (3 turns)
        human_count = 0
        recent: list[BaseMessage] = []
        for msg in reversed(messages[:-1]):
            if isinstance(msg, HumanMessage):
                recent.insert(0, msg)
                human_count += 1
            elif isinstance(msg, AIMessage) and not (hasattr(msg, "tool_calls") and msg.tool_calls):
                recent.insert(0, msg)
            if human_count >= 3:
                break
        stylist_messages.extend(recent)

        # Current query (LAST HumanMessage — tin nhắn hiện tại)
        current_query = None
        for msg in reversed(messages):
            if isinstance(msg, HumanMessage):
                current_query = msg
                break
        if current_query:
            stylist_messages.append(current_query)

        # Tool results
        tool_context = self._format_tool_result(tool_result)
        if tool_context:
            stylist_messages.append(HumanMessage(content=f"=== KẾT QUẢ TRA CỨU ===\n{tool_context}"))
        else:
            stylist_messages.append(
                HumanMessage(content="=== KHÔNG CÓ TOOL RESULT === (Khách chào hỏi hoặc câu hỏi chung)")
            )

        # ── LLM Call: Parse JSON from text (Claude returns text, not structured) ──
        # ── LLM Call: Parse JSON from text (Claude returns text, not structured) ──
        start = time.time()
        try:
            response = await self.stylist_llm.ainvoke(
                stylist_messages, config=config
            )
            # Parse JSON from text response
            text = response.content if hasattr(response, "content") else str(response)
            logger.info(f"[STYLIST] Raw response text (first 500 chars): {repr(text[:500])}")
            text = text.strip()
            if text.startswith('```json'):
                text = text[7:]
            if text.startswith('```'):
                text = text[3:]
            if text.endswith('```'):
                text = text[:-3]
            text = text.strip()

            output = json.loads(text)
            # Convert dict to object-like access
            class DictToObj:
                def __init__(self, data):
                    self.__dict__.update(data)
                def __getattr__(self, name):
                    return self.__dict__.get(name)
            output = DictToObj(output)

            # Handle nested user_insight dict
            if hasattr(output, "user_insight") and isinstance(output.user_insight, dict):
                output.user_insight = DictToObj(output.user_insight)
        except Exception as e:
            logger.error(f"❌ Stylist JSON parse error: {e}")
            fallback = "Dạ bạn cho mình hỏi thêm để tư vấn chính xác hơn nhé!"
            return {
                "messages": [AIMessage(content=fallback)],
                "updated_insight": None,
                "lead_stage": {"stage": 1, "stage_name": "BROWSE", "tone_directive": "Friendly"},
                "stage_injection": "",
                "product_ids": [],
                "diagnostics": [
                    {"step": "stylist_error", "label": "❌ Stylist Error", "content": str(e)[:300], "elapsed_ms": 0}
                ],
            }

        elapsed_ms = (time.time() - start) * 1000

        # Ensure product_ids covers every SKU mentioned in ai_response (deterministic safeguard)
        from_text = _extract_skus_from_text(output.ai_response)
        merged_product_ids = _dedupe_preserve_order((getattr(output, "product_ids", []) or []) + from_text)

        # Extract insight - handle both dict and object cases
        insight_raw = getattr(output, "user_insight", None)
        if isinstance(insight_raw, dict):
            insight_dict = insight_raw
        elif hasattr(insight_raw, "__dict__"):
            insight_dict = insight_raw.__dict__
        else:
            insight_dict = {}

        lead_stage = {
            "stage": insight_dict.get("STAGE_NUM", 1),
            "stage_name": insight_dict.get("STAGE", "BROWSE"),
            "tone_directive": insight_dict.get("TONE", "Friendly"),
            "behavioral_hints": insight_dict.get("BEHAVIORAL_HINTS", []),
        }
        stage_injection = format_insight_injection(insight_dict)

        ai_response = getattr(output, "ai_response", "Dạ mình chưa hiểu, bạn nói lại nhé!")
        product_ids = getattr(output, "product_ids", []) or []

        logger.info(
            f"💬 Stylist: stage={insight_dict.get('STAGE')} | goal={insight_dict.get('GOAL')} | "
            f"response_len={len(ai_response)} | products={merged_product_ids} | {elapsed_ms:.0f}ms"
        )

        diag = {
            "step": "stylist",
            "label": "💬 Stylist (Response + Insight)",
            "content": ai_response[:500],
            "elapsed_ms": round(elapsed_ms),
            "raw_json": json.dumps(insight_dict, ensure_ascii=False, indent=2),
        }

        return {
            "messages": [AIMessage(content=ai_response)],
            "updated_insight": insight_dict,
            "lead_stage": lead_stage,
            "stage_injection": stage_injection,
            "product_ids": merged_product_ids,
            "diagnostics": [diag],
        }

    # ─────────────────────────────────────
    # Router: Classifier → Stylist or END
    # ─────────────────────────────────────
    def _after_classifier(self, state: LeadStageState) -> str:
        if state.get("early_exit"):
            return END
        return "stylist"

    # ─────────────────────────────────────
    # Build Graph: Classifier → Stylist → END (always)
    # ─────────────────────────────────────
    def build(self):
        if self._compiled:
            return self._compiled

        workflow = StateGraph(LeadStageState)

        workflow.add_node("classifier", self._classifier_node)
        workflow.add_node("stylist", self._stylist_node)

        workflow.set_entry_point("classifier")

        workflow.add_conditional_edges("classifier", self._after_classifier, {"stylist": "stylist", END: END})

        workflow.add_edge("stylist", END)

        self._compiled = workflow.compile()
        logger.info("✅ LeadStageGraph compiled: Classifier → Stylist → END")
        return self._compiled

    # ─────────────────────────────────────
    # Main Entry
    # ─────────────────────────────────────
    async def chat(
        self,
        user_message: str,
        user_insight: str | None = None,
        history: list[BaseMessage] | None = None,
        config: RunnableConfig | None = None,
    ) -> dict:
        start = time.time()
        graph = self.build()

        messages: list[BaseMessage] = []
        if history:
            messages.extend(history)
        messages.append(HumanMessage(content=user_message))

        initial_state = {
            "messages": messages,
            "user_insight": user_insight,
            "tool_result": None,
            "tool_name_used": None,
            "updated_insight": None,
            "lead_stage": None,
            "stage_injection": None,
            "product_ids": [],
            "diagnostics": [],
        }

        result = await graph.ainvoke(initial_state, config=config)

        elapsed_ms = round((time.time() - start) * 1000, 2)
        all_messages = result["messages"]

        # Extract AI response
        ai_response = ""
        for msg in reversed(all_messages):
            if isinstance(msg, AIMessage):
                ai_response = _extract_text(msg.content)
                break

        # Build pipeline trace
        pipeline = [{"step": "user", "label": "👤 User", "content": user_message}]
        for diag in result.get("diagnostics", []):
            pipeline.append(diag)

        # Extract products from tool_result
        all_products = []
        tool_result_raw = result.get("tool_result")
        if tool_result_raw:
            try:
                parsed = json.loads(tool_result_raw)
                if parsed.get("status") == "success":
                    all_products = parsed.get("products", [])
            except Exception:
                pass

        # ════════════════════════════════════════════════════════════════
        # LỌC & FLATTEN SẢN PHẨM:
        # Lấy đúng các sản phẩm AI đã chọn trả ra trong product_ids,
        # bao gồm cả sản phẩm chính VÀ sản phẩm phối (nằm trong outfit_recommendations).
        # ════════════════════════════════════════════════════════════════
        final_products = []
        ai_product_ids = result.get("product_ids", [])

        if ai_product_ids and all_products:
            # Tạo dictionary lưu trữ mọi sản phẩm (chính + outfit) theo dạng in hoa, không khoảng trắng
            product_dict = {}
            for p in all_products:
                # Tìm đúng key của mã SKU cho SP chính
                p_sku = (
                    str(
                        p.get("magento_ref_code")
                        or p.get("sku_code")
                        or p.get("sku")
                        or p.get("internal_ref_code")
                        or ""
                    )
                    .upper()
                    .strip()
                )
                if p_sku:
                    # Đảm bảo frontend có trường 'sku' để render
                    if "sku" not in p:
                        p["sku"] = p_sku
                    product_dict[p_sku] = p

                # Duyệt luôn outfit_recommendations để gom data
                outfit_recs = p.get("outfit_recommendations", [])
                if isinstance(outfit_recs, list):
                    for rec in outfit_recs:
                        # Món phối thường nằm ở match_product_code hoặc magento_ref_code
                        rec_sku = (
                            str(rec.get("match_product_code") or rec.get("magento_ref_code") or rec.get("sku") or "")
                            .upper()
                            .strip()
                        )
                        if not rec_sku:
                            continue

                        # Tạo object tối thiểu để FE render được nếu chưa enrich
                        if "sku" not in rec:
                            rec["sku"] = rec_sku
                        if "name" not in rec:
                            rec["name"] = rec.get("match_product_name") or rec.get("product_name") or ""

                        product_dict[rec_sku] = rec

            # Enrich: nếu product_ids có món phối chưa có đủ data (price/image), lookup DB theo magento_ref_code
            wanted_ids = [str(x).upper().strip() for x in (ai_product_ids or []) if str(x).strip()]
            missing_ids = [
                pid
                for pid in wanted_ids
                if pid not in product_dict
                or not isinstance(product_dict.get(pid), dict)
                or not product_dict.get(pid, {}).get("name")
                or not product_dict.get(pid, {}).get("image")
            ]

            if missing_ids:
                try:
                    from agent.lead_stage_agent.product_search_engine import SELECT_COLUMNS, TABLE_NAME
                    from common.starrocks_connection import get_db_connection

                    db = get_db_connection()
                    # Query cả magento_ref_code (full) + internal_ref_code (base)
                    skus = list(dict.fromkeys([pid for pid in missing_ids if pid]))
                    base_skus = list(dict.fromkeys([pid.split("-")[0] for pid in skus if pid]))

                    keys = list(dict.fromkeys([*skus, *base_skus]))
                    placeholders = ",".join(["%s"] * len(keys))
                    sql = f"""
                    SELECT {SELECT_COLUMNS}
                    FROM {TABLE_NAME}
                    WHERE UPPER(magento_ref_code) IN ({placeholders})
                       OR UPPER(internal_ref_code) IN ({placeholders})
                    ORDER BY quantity_sold DESC NULLS LAST
                    LIMIT 200
                    """
                    rows = await db.execute_query_async(sql, params=tuple(keys + keys))

                    # Map kết quả
                    by_magento: dict[str, dict] = {}
                    by_internal: dict[str, dict] = {}
                    for r in rows or []:
                        mag = str(r.get("magento_ref_code") or "").upper().strip()
                        internal = str(r.get("internal_ref_code") or "").upper().strip()
                        if mag and mag not in by_magento:
                            by_magento[mag] = r
                        if internal and internal not in by_internal:
                            by_internal[internal] = r

                    def _row_to_card(row: dict) -> dict:
                        sale = float(row.get("sale_price") or 0)
                        orig = float(row.get("original_price") or 0)
                        has_discount = sale < orig and orig > 0
                        disc_pct = int(row.get("discount_percent") or 0)
                        sku = (row.get("magento_ref_code") or row.get("internal_ref_code") or "").strip()
                        return {
                            "sku": sku,
                            "name": row.get("product_name", ""),
                            "price": int(sale),
                            "original_price": int(orig),
                            "discount": f"-{disc_pct}%" if has_discount else None,
                            "color": row.get("master_color", ""),
                            "gender": row.get("gender_by_product", ""),
                            "product_line": row.get("product_line_vn", ""),
                            "image": row.get("product_image_url_thumbnail", ""),
                            "url": row.get("product_web_url", ""),
                            "sizes": row.get("size_scale", ""),
                        }

                    for pid in skus:
                        row = by_magento.get(pid) or by_internal.get(pid.split("-")[0])
                        if row:
                            product_dict[pid] = _row_to_card(row)
                except Exception as e:
                    logger.warning(f"⚠️ Product enrich by SKU failed: {str(e)[:200]}")

            # Gom đúng thứ tự AI trả ra, tránh duplicate
            seen_skus = set()
            for raw_pid in ai_product_ids:
                pid = str(raw_pid).upper().strip()
                if pid in product_dict and pid not in seen_skus:
                    final_products.append(product_dict[pid])
                    seen_skus.add(pid)
        else:
            # Nếu AI không trả product_ids (hoặc không có tool_result), trả về mặc định
            final_products = all_products

        logger.info(
            "🏷️ [LEAD] query='%s' | stage=%s | products=%d | time=%.0fms",
            user_message[:50],
            result.get("lead_stage", {}).get("stage_name", "?") if result.get("lead_stage") else "?",
            len(final_products),
            elapsed_ms,
        )

        return {
            "response": ai_response,
            "elapsed_ms": elapsed_ms,
            "lead_stage": result.get("lead_stage"),
            "updated_insight": result.get("updated_insight"),
            "pipeline": pipeline,
            "product_ids": result.get("product_ids", []),
            "products": final_products,
        }


# ═══════════════════════════════════════════════
# Singleton & Public API
# ═══════════════════════════════════════════════

_instance: LeadStageGraph | None = None


def get_lead_stage_agent(model_name: str | None = None) -> LeadStageGraph:
    global _instance
    if _instance is None:
        _instance = LeadStageGraph(model_name)
    return _instance


async def run_lead_stage_classifier(
    *,
    user_query: str,
    user_insight: str | None = None,
    chat_history_summary: str | None = None,
) -> dict:
    """Backward compat API."""
    agent = get_lead_stage_agent()
    state = {
        "messages": [HumanMessage(content=user_query)],
        "user_insight": user_insight,
        "tool_result": None,
        "tool_name_used": None,
        "updated_insight": None,
        "lead_stage": None,
        "stage_injection": None,
        "diagnostics": [],
    }
    res = await agent._classifier_node(state)
    return {
        "lead_stage": res.get("lead_stage"),
        "classifier_latency_ms": res.get("diagnostics", [{}])[0].get("elapsed_ms", 0),
        "tool_result": res.get("tool_result"),
    }
