"""
Store Search Graph — 2-Agent LangGraph Architecture cho Store Search.
"""

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 agent.tools.store_search_tool import canifa_store_search
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 StoreSearchState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]


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

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

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

    def _planner_node(self, state: StoreSearchState) -> dict:
        messages = state["messages"]

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

        # synchronous run (we use ainvoke in async wrapper)
        pass

    async def _planner_node_async(self, state: StoreSearchState) -> dict:
        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]}


    async def _responder_node_async(self, state: StoreSearchState) -> dict:
        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]}


    def _after_planner(self, state: StoreSearchState) -> str:
        last = state["messages"][-1]
        if hasattr(last, "tool_calls") and last.tool_calls:
            return "tools"
        return "end"

    def build(self):
        if self._compiled:
            return self._compiled

        workflow = StateGraph(StoreSearchState)
        workflow.add_node("planner", self._planner_node_async)
        workflow.add_node("tools", ToolNode(self.tools))
        workflow.add_node("responder", self._responder_node_async)

        workflow.set_entry_point("planner")
        workflow.add_conditional_edges(
            "planner",
            self._after_planner,
            {"tools": "tools", "end": END},
        )
        workflow.add_edge("tools", "responder")
        workflow.add_edge("responder", END)

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

    async def chat(self, user_message: str, history: list[BaseMessage] | None = None) -> dict:
        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]
                
                pipeline.append({
                    "step": "tool",
                    "label": f"🛠️ Tool ({msg.name})",
                    "content": summary,
                })

        # Cắt ngắn step if needed
        end_idx = "responder" if has_tool_calls else "planner"
        path_str = f"planner→tools→responder" if has_tool_calls else "planner"
        
        logger.info(f"🏬 [STORE SEARCH] query='{user_message[:50]}' | path={path_str} | tools={len(tool_calls_log)} | time={elapsed_ms}ms")

        return {
            "response": ai_response,
            "elapsed_ms": elapsed_ms,
            "agent_path": path_str,
            "tool_calls": tool_calls_log,
            "pipeline": pipeline
        }

# --- Lazy Init ---
_agent_instance = None

def get_store_search_agent() -> StoreSearchGraph:
    global _agent_instance
    if _agent_instance is None:
        _agent_instance = StoreSearchGraph()
    return _agent_instance
