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

Flow: User query → Planner AI → (hỏi lại | gọi tool) → Responder AI → answer

2 con AI:
  - Planner:   Phân tích ý định, extract params, gọi tool. Nếu mơ hồ → hỏi lại.
  - Responder: Nhận kết quả tool, format câu trả lời tự nhiên, thân thiệ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 .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
# ═══════════════════════════════════════════════
class TagSearchState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]


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

    # 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]

        # 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"✅ TagSearchGraph 2-Agent initialized with model: {self.model_name}")

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

        # Inject planner system prompt
        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: TagSearchState) -> dict:
        """Responder AI: nhận kết quả tool, format câu trả lời đẹp."""
        messages = state["messages"]

        # Inject responder system prompt
        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: TagSearchState) -> str:
        """Sau Planner: nếu có tool_calls → đi 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"  # Planner hỏi lại khách → end luôn

    def build(self):
        """Build 2-agent graph."""
        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 | 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("✅ TagSearchGraph 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
        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:
                tool_content = content_str
                summary = tool_content[:600]
                try:
                    import json as _json
                    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", [])
                        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')}"
                except Exception:
                    pass

                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,
                })

        # Collect products from tool results for frontend rendering
        all_products = []
        for step in pipeline:
            if step.get("step") == "tool_result" and step.get("raw_json"):
                try:
                    import json as _json2
                    parsed_result = _json2.loads(step["raw_json"])
                    if parsed_result.get("status") == "success":
                        all_products.extend(parsed_result.get("products", []))
                except Exception:
                    pass

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

        logger.info(
            "🏷️ [TAG 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,
            "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




