"""
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 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 common.llm_factory import create_llm
from config import DEFAULT_MODEL

from .prompts import LEAD_STAGE_CLASSIFIER_PROMPT, STYLIST_SYSTEM_PROMPT, format_insight_injection

# Tools
from .lead_search_tool import tag_search_tool
from agent.tools.brand_knowledge_tool import canifa_knowledge_search
from agent.tools.promotion_canifa_tool import canifa_get_promotions
from agent.tools.store_search_tool import canifa_store_search
from agent.tools.check_is_stock import check_is_stock
from agent.tools.customer_info_tool import collect_customer_info
from agent.tools.add_to_cart_tool import add_to_cart

logger = logging.getLogger(__name__)


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

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

    reasoning: str = Field(description="Lý luận gọi tool hay trả lời trực tiếp")
    
    # ── TH 1: GỌI TOOL ──
    tool_name: str | None = Field(description="Tên tool (tag_search_tool). Bắt buộc để null nếu khách chỉ đang tâm sự/chào hỏi/ko cần DB.")
    tool_args: dict | None = Field(description="Tham số tool. Null nếu tool_name=null.")
    
    # ── TH 2: EARLY EXIT ──
    ai_response: str | None = Field(description="Câu trả lời cho khách. CHỈ nhả ra khi tool_name=null. Nếu gọi tool, bắt buộc để null.")
    product_ids: list[str] = Field(default_factory=list, description="Mã SKU (nếu khách hỏi trúng mã). Bắt buộc phải có để hiện trên frontend.")


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 "")


# ═══════════════════════════════════════════════
# 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 = {
            "tag_search_tool": tag_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, NHẸ (streaming=False)
        _classifier_base = create_llm(model_name=self.model_name, streaming=False)
        self.classifier_llm = _classifier_base.with_structured_output(
            ClassifierOutput, method="function_calling"
        )

        # AI 2: Stylist — structured output, NẶNG (streaming=False vì structured)
        _stylist_base = create_llm(model_name=self.model_name, streaming=False)
        self.stylist_llm = _stylist_base.with_structured_output(
            StylistOutput, method="function_calling"
        )

        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 = 4) -> str:
        """Format recent messages cho context."""
        lines = []
        human_count = 0
        for msg in reversed(messages):
            if isinstance(msg, HumanMessage):
                lines.insert(0, f"Khách: {_extract_text(msg.content)[:200]}")
                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)[:200]}")
            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]}"
                        lines.append(line)
                    return "\n".join(lines)
                else:
                    return parsed.get("message", "Không tìm thấy sản phẩm phù hợp.")
            else:
                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])
        context_block = (
            f"=== TIN NHẮN HIỆN TẠI ===\n{user_query}\n\n"
            f"=== USER INSIGHT ===\n{user_insight}\n\n"
            f"=== LỊCH SỬ ===\n{history_text}"
        )

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

        start = time.time()
        try:
            output: ClassifierOutput = await self.classifier_llm.ainvoke(
                classifier_messages, config=config
            )
        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
        tool_args = output.tool_args or {}
        ai_response = output.ai_response

        # 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."

        stylist_sys = STYLIST_SYSTEM_PROMPT.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: Structured Output ──
        start = time.time()
        try:
            output: StylistOutput = await self.stylist_llm.ainvoke(
                stylist_messages, config=config
            )
        except Exception as e:
            logger.error(f"❌ Stylist structured output error: {e}")
            # Fallback: trả response text thuần
            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

        # Extract insight
        insight_dict = output.user_insight.model_dump()
        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)

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

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

        return {
            "messages": [AIMessage(content=output.ai_response)],
            "updated_insight": insight_dict,
            "lead_stage": lead_stage,
            "stage_injection": stage_injection,
            "product_ids": output.product_ids or [],
            "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,
            "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
                
        # Filter products by AI selection
        if "product_ids" in result and result["product_ids"] is not None:
            ai_chosen_ids = result["product_ids"]
            filtered_products = []
            for pid in ai_chosen_ids:
                for p in all_products:
                    psku = str(p.get("sku_code") or p.get("sku") or p.get("internal_ref_code") or "").upper()
                    if psku and psku in str(pid).upper():
                        filtered_products.append(p)
                        break
            # LUÔN ghi đè. Nếu Stylist chốt mảng rỗng [] thì Frontend cũng nhận list rỗng (chứ k fallback nhét bừa)
            all_products = filtered_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(all_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,
            "products": all_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"),
    }
