"""
SKU Search Graph — 2-Agent LangGraph Architecture cho tra cứu mã SKU.

Flow: User query → Planner AI → gọi sku_search_tool → Responder AI → answer

2 con AI:
  - Planner:   Phân tích câu hỏi, extract mã SKU, gọi tool.
  - Responder: Nhận kết quả tool, format câu trả lời tự nhiên.
"""

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

from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
from langgraph.graph import END, StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode

from common.llm_factory import create_llm
from config import DEFAULT_MODEL

from .sku_search_tool import sku_search_tool

logger = logging.getLogger(__name__)


def _extract_text(content) -> str:
    """Extract plain text from msg.content — handles both str and Gemini list format."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        # Gemini returns [{'type': 'text', 'text': '...'}, ...]
        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 "")

# ═══════════════════════════════════════════════
# System Prompts — 2 con AI riêng biệt
# ═══════════════════════════════════════════════

PLANNER_PROMPT = """Bạn là AI PLANNER của CANIFA. Nhiệm vụ DUY NHẤT: Phân tích câu hỏi khách → tìm MÃ SKU → gọi tool.

## LUỒNG XỬ LÝ
1. Đọc câu hỏi khách
2. Tìm MÃ SKU trong câu (VD: "6OT24S001-SG055", "TS092", "ABCD123")
3. Gọi sku_search_tool với mã tìm được

## CÁCH NHẬN DIỆN MÃ SKU
- Mã SKU thường gồm chữ cái + số, có thể có dấu gạch ngang
- VD: "6OT24S001-SG055", "TS092", "1AT25S008", "6QT24W008-SK090"
- Nếu khách hỏi "mã 6OT24S001" → sku_code = "6OT24S001"
- Nếu khách paste nguyên mã dài → lấy nguyên văn, BAO GỒM phần sau dấu gạch ngang
- VD: "6OT24S001-SG055" → sku_code = "6OT24S001-SG055" (KHÔNG cắt bỏ -SG055)

## KHI NÀO HỎI LẠI
- Khách không cung cấp mã SKU nào → Hỏi: "Bạn vui lòng cho mình mã sản phẩm (SKU) cần tra cứu nhé!"
- Khách hỏi chung chung như "tìm áo" → Hỏi: "Bạn có mã SKU cụ thể không ạ? Nếu muốn tìm theo tiêu chí, bạn dùng tính năng Tìm kiếm sản phẩm nhé."

## PARAMS
- sku_code: Mã SKU CHÍNH XÁC khách cung cấp (bắt buộc)
- master_color: Màu sắc nếu khách NÓI RÕ (VD: "mã 6OT24S001 màu đỏ" → master_color="đỏ"). Để null nếu không nhắc.

## QUY TẮC
1. TUYỆT ĐỐI KHÔNG trả lời về sản phẩm nếu chưa gọi tool.
2. Chỉ trả lời text thuần khi HỎI LẠI khách.
3. Tool sẽ tự động tìm tất cả biến thể (màu, size) cùng mã gốc.
4. Nếu khách hỏi "mã này có màu gì" → gọi tool, tool trả hết variants.
5. Nếu khách hỏi "mã này màu đỏ" → sku_code=mã, master_color="đỏ".

## VÍ DỤ
- "Tra cứu mã 6OT24S001-SG055" → sku_code="6OT24S001-SG055"
- "Mã TS092 màu trắng" → sku_code="TS092", master_color="trắng"
- "6QT24W008" → sku_code="6QT24W008"
- "Tìm áo đẹp" → HỎI LẠI mã SKU
"""

RESPONDER_PROMPT = """Bạn là AI RESPONDER của CANIFA. Nhiệm vụ: Format kết quả tra cứu SKU thành câu trả lời tự nhiên.

## FORMAT KẾT QUẢ CHÍNH
Với mỗi sản phẩm tìm được:
- **Tên sản phẩm** (Mã: SKU)
- 🎨 Màu: [color_name]
- 💰 Giá: xxx₫ ~~(giá gốc)~~ (-xx%)
- 📐 Size: [sizes]
- 🔗 [Xem chi tiết tại đây](url)

## NẾU CÓ BIẾN THỂ (VARIANTS)
Liệt kê các biến thể theo dạng bảng ngắn gọn:
| Màu | Size | Giá | Link |
|-----|------|-----|------|

## NGUYÊN TẮC
- Trả lời bằng tiếng Việt, thân thiện, chuyên nghiệp.
- Hiển thị thông tin CHÍNH XÁC từ tool, KHÔNG bịa.
- Nếu 0 kết quả → "Rất tiếc, mình không tìm thấy sản phẩm với mã này. Bạn vui lòng kiểm tra lại mã SKU nhé!"
- Nếu sản phẩm hết hàng hoặc giá = 0 → ghi chú rõ.
- Cuối cùng hỏi nhẹ: "Bạn cần tra cứu mã nào khác không?"
"""


# ═══════════════════════════════════════════════
# State
# ═══════════════════════════════════════════════
class SkuSearchState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]


# ═══════════════════════════════════════════════
# Graph Builder — 2-Agent Architecture
# ═══════════════════════════════════════════════
class SkuSearchGraph:
    """2-Agent LangGraph: Planner → Tools → Responder."""

    def __init__(self, model_name: str | None = None):
        self.model_name = model_name or DEFAULT_MODEL
        self.tools = [sku_search_tool]

        # 2 con LLM riêng biệt
        self.planner_llm = create_llm(model_name=self.model_name, streaming=True)
        self.planner_with_tools = self.planner_llm.bind_tools(self.tools)

        self.responder_llm = create_llm(model_name=self.model_name, streaming=True)

        self._compiled = None
        logger.info(f"✅ SkuSearchGraph 2-Agent initialized with model: {self.model_name}")

    # ─────────────────────────────────────
    # Node 1: PLANNER — Phân tích + gọi tool
    # ─────────────────────────────────────
    async def _planner_node(self, state: SkuSearchState) -> dict:
        """Planner AI: phân tích câu hỏi, extract SKU, gọi tool."""
        messages = state["messages"]

        planner_messages = [SystemMessage(content=PLANNER_PROMPT)]
        for msg in messages:
            if not isinstance(msg, SystemMessage):
                planner_messages.append(msg)

        response = await self.planner_with_tools.ainvoke(planner_messages)
        return {"messages": [response]}

    # ─────────────────────────────────────
    # Node 2: RESPONDER — Format kết quả
    # ─────────────────────────────────────
    async def _responder_node(self, state: SkuSearchState) -> dict:
        """Responder AI: nhận kết quả tool, format câu trả lời đẹp."""
        messages = state["messages"]

        responder_messages = [SystemMessage(content=RESPONDER_PROMPT)]
        for msg in messages:
            if not isinstance(msg, SystemMessage):
                responder_messages.append(msg)

        response = await self.responder_llm.ainvoke(responder_messages)
        return {"messages": [response]}

    # ─────────────────────────────────────
    # Router: Planner → tool hoặc end
    # ─────────────────────────────────────
    def _after_planner(self, state: SkuSearchState) -> str:
        """Sau Planner: nếu có tool_calls → tools, nếu không → end (hỏi lại)."""
        last = state["messages"][-1]
        if hasattr(last, "tool_calls") and last.tool_calls:
            return "tools"
        return "end"

    def build(self):
        """Build 2-agent graph."""
        if self._compiled:
            return self._compiled

        workflow = StateGraph(SkuSearchState)

        # 3 Nodes
        workflow.add_node("planner", self._planner_node)
        workflow.add_node("tools", ToolNode(self.tools))
        workflow.add_node("responder", self._responder_node)

        # Flow: Entry → Planner
        workflow.set_entry_point("planner")

        # Planner → (tools | end)
        workflow.add_conditional_edges(
            "planner",
            self._after_planner,
            {"tools": "tools", "end": END},
        )

        # Tools → Responder → END
        workflow.add_edge("tools", "responder")
        workflow.add_edge("responder", END)

        self._compiled = workflow.compile()
        logger.info("✅ SkuSearchGraph 2-Agent compiled")
        return self._compiled

    async def chat(self, user_message: str, history: list[BaseMessage] | None = None) -> dict:
        """
        Main entry point — gửi message và nhận response.
        Returns full pipeline trace for UI display.
        """
        start = time.time()
        graph = self.build()

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

        result = await graph.ainvoke({"messages": messages})

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

        ai_response = ""
        tool_calls_log: list[dict[str, Any]] = []
        has_tool_calls = False

        # Build full pipeline trace for UI
        pipeline: list[dict[str, Any]] = []

        for msg in all_messages:
            content_str = _extract_text(msg.content)

            if isinstance(msg, HumanMessage):
                pipeline.append({
                    "step": "user",
                    "label": "👤 User",
                    "content": content_str,
                })
            elif isinstance(msg, AIMessage):
                if msg.tool_calls:
                    has_tool_calls = True
                    tool_calls_log.extend(msg.tool_calls)
                    pipeline.append({
                        "step": "planner",
                        "label": "🧠 Planner AI",
                        "content": content_str if content_str and content_str != "[]" else "(gọi tool...)",
                        "tool_calls": [
                            {"name": tc["name"], "args": tc["args"]}
                            for tc in msg.tool_calls
                        ],
                    })
                elif msg.content:
                    ai_response = content_str
                    pipeline.append({
                        "step": "responder",
                        "label": "💬 Responder AI",
                        "content": content_str[:500] + ("..." if len(content_str) > 500 else ""),
                    })
            else:
                # ToolMessage — result from tool execution
                tool_content = content_str
                # Try to parse and summarize
                summary = tool_content[:600]
                try:
                    import json as _json
                    parsed = _json.loads(tool_content)
                    if parsed.get("status") == "success":
                        n_colors = parsed.get('unique_colors', parsed.get('total_variants', '?'))
                        summary = f"✅ Tìm thấy {n_colors} sản phẩm ({parsed.get('elapsed_ms', '?')}ms)"
                        # Color filter info
                        if parsed.get("color_filter"):
                            if parsed.get("color_found") is False:
                                summary += f"\n⚠️ master_color '{parsed['color_filter']}' KHÔNG có!"
                                ce = parsed.get("color_exception", {})
                                if ce.get("matched_product_color_code"):
                                    summary += f"\n🏷️ product_color_code {ce['matched_product_color_code']} thực tế là: {ce['matched_master_color']}"
                                
                                recs = ce.get("recommend_colors", [])
                                if recs:
                                    rec_strs = [f"{r['master_color']} ({r['product_color_code'].split('-')[-1] if '-' in r['product_color_code'] else r['product_color_code']})" for r in recs[:3]]
                                    summary += f"\n🎨 Recommend: {', '.join(rec_strs)}"
                            else:
                                summary += f"\n🎨 Lọc màu: {parsed['color_filter']} → {parsed.get('matched_count', '?')} kết quả"
                        elif parsed.get("available_colors"):
                            summary += f"\n🎨 Màu có: {', '.join(parsed['available_colors'])}"
                        # Product preview
                        products = parsed.get("products", [])
                        if products:
                            names = [f"• {p.get('name','')} - {p.get('color','')} ({p.get('sku_color',p.get('sku',''))})" for p in products[:5]]
                            summary += "\n" + "\n".join(names)
                            if len(products) > 5:
                                summary += f"\n... và {len(products) - 5} sản phẩm khác"
                        if parsed.get("message"):
                            summary += f"\n💡 {parsed['message']}"
                    elif parsed.get("status") == "error":
                        summary = f"❌ {parsed.get('message', 'Error')}"
                except Exception:
                    pass

                # Pretty-print raw JSON for UI display
                raw_json = tool_content
                try:
                    raw_json = _json.dumps(_json.loads(tool_content), ensure_ascii=False, indent=2, default=str)
                except Exception:
                    pass

                pipeline.append({
                    "step": "tool_result",
                    "label": f"📦 Tool Result ({getattr(msg, 'name', 'tool')})",
                    "content": summary,
                    "raw_json": raw_json,
                })

        agent_path = "planner→tools→responder" if has_tool_calls else "planner_only"

        logger.info(
            "🔍 [SKU SEARCH] query='%s' | path=%s | tools=%d | time=%.2fms",
            user_message[:50], agent_path, len(tool_calls_log), elapsed_ms,
        )

        return {
            "response": ai_response,
            "elapsed_ms": elapsed_ms,
            "tool_calls": [
                {"name": tc["name"], "args": tc["args"]}
                for tc in tool_calls_log
            ],
            "agent_path": agent_path,
            "pipeline": pipeline,
        }


# ═══════════════════════════════════════════════
# Singleton
# ═══════════════════════════════════════════════
_instance: SkuSearchGraph | None = None


def get_sku_search_agent(model_name: str | None = None) -> SkuSearchGraph:
    """Get or create SkuSearchGraph singleton."""
    global _instance
    if _instance is None:
        _instance = SkuSearchGraph(model_name)
    return _instance
