"""
Tag Search Graph — 2-Agent LangGraph Architecture cho tư vấn sản phẩm.

Flow: User → Planner ⇄ Tools → Responder → END

2 con AI:
  - Planner:   Phân tích ý định, extract params, gọi tool. Nếu mơ hồ → hỏi lại.
               Sau khi tool trả kết quả → pass-through → route tới Responder.
  - Responder: Nhận kết quả tool, format câu trả lời tự nhiên, thân thiện.
               KHÔNG gọi tools.

Architecture aligned with LeadStageGraph pattern (loop-back, diagnostics, RunnableConfig).
"""

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

from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
from langchain_core.runnables import RunnableConfig
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 .tag_search_tool import tag_search_tool
from .prompts import PLANNER_PROMPT, RESPONDER_PROMPT


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


# ═══════════════════════════════════════════════
# State — enriched with diagnostics & routing flags
# ═══════════════════════════════════════════════
class TagSearchState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]

    # Routing flag: True khi Planner trả lời trực tiếp (hỏi lại / greeting)
    planner_responded: bool

    # Pipeline diagnostics cho UI trace display
    diagnostics: Annotated[list[dict[str, Any]], lambda x, y: x + y]


# ═══════════════════════════════════════════════
# Graph Builder — 2-Agent Architecture (with loop-back)
# Planner ⇄ Tools → Responder → END
# ═══════════════════════════════════════════════
class TagSearchGraph:
    """2-Agent LangGraph: Planner ⇄ Tools → Responder → END."""

    # Tag Search dùng GPT-5-nano — nhanh nhất, tool calling tốt
    TAG_SEARCH_MODEL = "gpt-5.4-nano"

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

        # AI 1: Planner + Tool Caller (streaming off for tool calling reliability)
        self.planner_llm = create_llm(model_name=self.model_name, streaming=False)
        self.planner_with_tools = self.planner_llm.bind_tools(self.tools)

        # AI 2: Responder — response only (streaming ON, NO tools)
        self.responder_llm = create_llm(model_name=self.model_name, streaming=True)

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

    # ─────────────────────────────────────
    # Node 1: PLANNER — Phân tích + gọi tool (with loop-back)
    # ─────────────────────────────────────
    async def _planner_node(self, state: TagSearchState, config: RunnableConfig | None = None) -> dict:
        """
        Planner AI: phân tích câu hỏi, quyết định gọi tool hoặc hỏi lại.

        Pass 1: Nhận câu hỏi user → gọi tool hoặc trả lời trực tiếp.
        Pass 2+: Tool results đã có → pass-through → route tới Responder.
        """
        messages = state["messages"]

        has_tool_results = any(isinstance(m, ToolMessage) for m in messages)

        # ── PASS 2+: Tool results back → skip LLM, go to Responder ──
        if has_tool_results:
            logger.info("🔄 Planner pass-through → Responder")
            return {
                "planner_responded": False,
                "diagnostics": [],
            }

        # ── PASS 1: Single LLM call — analyze + tools ──
        planner_messages = [SystemMessage(content=PLANNER_PROMPT)]
        for msg in messages:
            if not isinstance(msg, SystemMessage):
                planner_messages.append(msg)

        start = time.time()
        response = await self.planner_with_tools.ainvoke(planner_messages, config=config)
        elapsed_ms = (time.time() - start) * 1000

        response_text = _extract_text(response.content)
        has_tool_calls = hasattr(response, "tool_calls") and response.tool_calls

        # ── Build diagnostic ──
        diag: dict[str, Any] = {"elapsed_ms": round(elapsed_ms)}

        if has_tool_calls:
            tool_names = ', '.join(tc['name'] for tc in response.tool_calls)
            diag["step"] = "planner_tool_call"
            diag["label"] = "🧠 Planner AI (Tool Call)"
            diag["content"] = f"Tools: {tool_names}"
            diag["tool_calls"] = [
                {"name": tc["name"], "args": tc["args"]}
                for tc in response.tool_calls
            ]
        else:
            diag["step"] = "planner_direct"
            diag["label"] = "💬 Planner AI (Direct)"
            diag["content"] = response_text[:500]

        # planner_responded = True → Planner hỏi lại hoặc greeting → skip Responder
        is_direct = not has_tool_calls

        return {
            "messages": [response],
            "planner_responded": is_direct,
            "diagnostics": [diag],
        }

    # ─────────────────────────────────────
    # Router: Planner → tools | responder | end
    # ─────────────────────────────────────
    def _after_planner(self, state: TagSearchState) -> str:
        """
        Sau Planner:
        - Có tool_calls → đi tools
        - Planner trả lời trực tiếp (hỏi lại, greeting) → end
        - Mọi trường hợp khác (sau tools loop-back) → responder
        """
        last = state["messages"][-1]
        if hasattr(last, "tool_calls") and last.tool_calls:
            return "tools"
        if state.get("planner_responded"):
            return "end"
        return "responder"

    # ─────────────────────────────────────
    # Node 2: RESPONDER — Format kết quả (NO tools)
    # ─────────────────────────────────────
    async def _responder_node(self, state: TagSearchState, config: RunnableConfig | None = None) -> dict:
        """
        Responder AI: Nhận kết quả tool, format câu trả lời tư vấn đẹp.
        KHÔNG gọi tools — chỉ viết response.
        """
        messages = state["messages"]

        # Collect tool results for cleaner context
        tool_context_parts = []
        for msg in messages:
            if isinstance(msg, ToolMessage):
                tool_name = getattr(msg, "name", "tool")
                tool_content = _extract_text(msg.content)
                try:
                    parsed = json.loads(tool_content)
                    if parsed.get("status") == "success":
                        products = parsed.get("products", [])
                        if products:
                            product_lines = []
                            for p in products[:8]:
                                line = f"- {p.get('name', '')} | {p.get('price', '')} | SKU: {p.get('sku_code', '')}"
                                if p.get("url"):
                                    line += f" | Link: {p['url']}"
                                product_lines.append(line)
                            tool_context_parts.append(
                                f"[{tool_name}] Tìm thấy {len(products)} sản phẩm:\n" + "\n".join(product_lines)
                            )
                        else:
                            tool_context_parts.append(f"[{tool_name}] {json.dumps(parsed, ensure_ascii=False)[:500]}")
                    else:
                        tool_context_parts.append(f"[{tool_name}] {parsed.get('message', tool_content[:300])}")
                except Exception:
                    tool_context_parts.append(f"[{tool_name}] {tool_content[:500]}")

        # Build responder messages
        responder_messages = [SystemMessage(content=RESPONDER_PROMPT)]

        # Add user's original question
        for msg in messages:
            if isinstance(msg, HumanMessage):
                responder_messages.append(msg)
                break  # Only first human message (the actual query)

        # Add tool results as context
        if tool_context_parts:
            tool_summary = "=== KẾT QUẢ TRA CỨU TỪ TOOLS ===\n" + "\n\n".join(tool_context_parts)
            responder_messages.append(HumanMessage(content=tool_summary))

        start = time.time()
        response = await self.responder_llm.ainvoke(responder_messages, config=config)
        elapsed_ms = (time.time() - start) * 1000

        response_text = _extract_text(response.content)

        diag = {
            "step": "responder",
            "label": "💬 Responder AI (Response)",
            "content": response_text[:500] + ("..." if len(response_text) > 500 else ""),
            "elapsed_ms": round(elapsed_ms),
        }

        return {
            "messages": [response],
            "diagnostics": [diag],
        }

    # ─────────────────────────────────────
    # Build Graph
    # ─────────────────────────────────────
    def build(self):
        """Build 2-agent graph: Planner ⇄ Tools → Responder → END."""
        if self._compiled:
            return self._compiled

        workflow = StateGraph(TagSearchState)

        # 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 | responder | end)
        workflow.add_conditional_edges(
            "planner",
            self._after_planner,
            {"tools": "tools", "responder": "responder", "end": END},
        )

        # Tools → Planner (LOOP BACK — planner sees tool results, routes to responder)
        workflow.add_edge("tools", "planner")

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

        self._compiled = workflow.compile()
        logger.info("✅ TagSearchGraph compiled: Planner ⇄ Tools → Responder → END")
        return self._compiled

    # ─────────────────────────────────────
    # Main Entry Point
    # ─────────────────────────────────────
    async def chat(self, user_message: str, history: list[BaseMessage] | None = None, config: RunnableConfig | None = None) -> dict:
        """
        Main entry point — gửi message và nhận response.
        Returns 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))

        initial_state = {
            "messages": messages,
            "planner_responded": False,
            "diagnostics": [],
        }

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

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

        # Extract AI response (last AI message without tool_calls)
        ai_response = ""
        for msg in reversed(all_messages):
            if isinstance(msg, AIMessage) and not (hasattr(msg, "tool_calls") and msg.tool_calls):
                ai_response = _extract_text(msg.content)
                break

        # Build pipeline trace
        pipeline = [{"step": "user", "label": "👤 User", "content": user_message}]

        if result.get("diagnostics"):
            for diag in result["diagnostics"]:
                pipeline.append(diag)

        # Parse tool results for product cards
        all_products = []
        for msg in all_messages:
            if isinstance(msg, ToolMessage):
                tool_content = _extract_text(msg.content)
                summary = tool_content[:600]
                raw_json = tool_content
                try:
                    parsed = json.loads(tool_content)
                    if parsed.get("status") == "success":
                        result_count = parsed.get("count", len(parsed.get("products", [])))
                        tier = parsed.get("tier", 0)
                        tier_info = f"[Tier {tier}]" if tier > 0 else "[SKU Direct]"
                        summary = f"✅ Tìm thấy {result_count} sản phẩm ({parsed.get('elapsed_ms', '?')}ms) {tier_info}"
                        products = parsed.get("products", [])
                        all_products.extend(products)
                        if products:
                            names = [f"• {p.get('name', '')}" 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"
                    elif parsed.get("status") == "error":
                        summary = f"❌ {parsed.get('message', 'Error')}"
                    raw_json = json.dumps(parsed, 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,
                })

        # Determine agent path for logging
        has_tool_calls = any(
            isinstance(msg, AIMessage) and hasattr(msg, "tool_calls") and msg.tool_calls
            for msg in all_messages
        )
        agent_path = "planner→tools→responder" if has_tool_calls else "planner_only"

        logger.info(
            "🏷️ [TAG SEARCH] query='%s' | path=%s | products=%d | time=%.2fms",
            user_message[:50], agent_path, len(all_products), elapsed_ms,
        )

        return {
            "response": ai_response,
            "elapsed_ms": elapsed_ms,
            "tool_calls": [
                {"name": tc["name"], "args": tc["args"]}
                for msg in all_messages
                if isinstance(msg, AIMessage) and hasattr(msg, "tool_calls") and msg.tool_calls
                for tc in msg.tool_calls
            ],
            "agent_path": agent_path,
            "pipeline": pipeline,
            "products": all_products,
        }


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


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