"""
Image Search Graph — 2-Agent LangGraph Architecture cho Image Search.
Hỗ trợ Multi-modal vision.
"""

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.lead_stage_agent.lead_search_tool import lead_search_tool as canifa_tag_search
from .prompts import VISION_PLANNER_PROMPT, VISION_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 ImageSearchState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]


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

    def __init__(self, model_name: str | None = None):
        self.model_name = model_name or DEFAULT_MODEL
        self.tools = [canifa_tag_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"✅ ImageSearchGraph 2-Agent initialized with model: {self.model_name}")

    async def _planner_node_async(self, state: ImageSearchState) -> dict:
        messages = state["messages"]
        planner_messages = [SystemMessage(content=VISION_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: ImageSearchState) -> dict:
        messages = state["messages"]
        responder_messages = [SystemMessage(content=VISION_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: ImageSearchState) -> 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(ImageSearchState)
        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("✅ ImageSearchGraph 2-Agent compiled")
        return self._compiled

    async def chat(self, user_message: str, images: list[str] | None = None, history: list[BaseMessage] | None = None) -> dict:
        start = time.time()
        graph = self.build()

        messages: list[BaseMessage] = []
        if history:
            messages.extend(history)
            
        # Xử lý multimodal (ảnh + text)
        if images and len(images) > 0:
            text_content = user_message + "\n\n[📸 Xin hãy mô tả rõ ràng, cụ thể các đặc điểm của trang phục trong ảnh để tìm kiếm món đồ tương tự.]" 
            multimodal_content = [{"type": "text", "text": text_content}]
            for img in images:
                image_url = img if img.startswith("data:") or img.startswith("http") else f"data:image/jpeg;base64,{img}"
                multimodal_content.append({
                    "type": "image_url",
                    "image_url": {"url": image_url, "detail": "auto"}
                })
            messages.append(HumanMessage(content=multimodal_content))
            logger.info(f"📸 [IMAGE SEARCH] Injected {len(images)} image(s)")
        else:
            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):
                # Ẩn bớt base64 dài dòng trong log
                if isinstance(msg.content, list):
                    clean_content = []
                    for it in msg.content:
                        if it.get("type") == "image_url":
                            clean_content.append("[📸 Bức ảnh đã đính kèm]")
                        elif it.get("type") == "text":
                            clean_content.append(it.get("text", ""))
                    content_str = "\n".join(clean_content)
                
                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": "🧠 Vision Planner",
                        "content": content_str if content_str and content_str != "[]" else "(phân tích ảnh & 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_result",
                    "label": f"🛠️ Tool ({msg.name})",
                    "content": summary,
                })

        end_idx = "responder" if has_tool_calls else "planner"
        path_str = f"planner→tools→responder" if has_tool_calls else "planner"
        
        logger.info(f"🖼️ [IMAGE 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_image_search_agent() -> ImageSearchGraph:
    global _agent_instance
    if _agent_instance is None:
        _agent_instance = ImageSearchGraph()
    return _agent_instance
