"""
Diagram Agent Graph — 2-Agent LangGraph: Planner → Tool → Responder.
Planner generates Mermaid code, Responder explains the result.
"""

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

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

from config import OPENAI_API_KEY

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


# ═══════════════════════════════════════════════
# Tool: Generate Diagram
# ═══════════════════════════════════════════════
@tool
def generate_diagram(mermaid_code: str, title: str = "", diagram_type: str = "flowchart") -> str:
    """Generate a Mermaid diagram. Returns the Mermaid code for frontend rendering.

    Args:
        mermaid_code: Complete Mermaid syntax code for the diagram.
        title: Optional title for the diagram.
        diagram_type: Type of diagram (flowchart, sequence, class, gantt, pie, mindmap, er, state).

    Returns:
        JSON string with diagram data for frontend rendering.
    """
    import json

    # Validate mermaid_code is not empty
    if not mermaid_code or not mermaid_code.strip():
        return json.dumps({
            "status": "error",
            "message": "Mermaid code is empty",
        }, ensure_ascii=False)

    return json.dumps({
        "status": "success",
        "mermaid_code": mermaid_code.strip(),
        "title": title,
        "diagram_type": diagram_type,
    }, ensure_ascii=False)


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


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

    DIAGRAM_MODEL = "gpt-5.4-nano"

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

        # Planner needs higher max_tokens (4000) for complex Mermaid code generation
        self.planner_llm = ChatOpenAI(
            model=self.model_name, streaming=True, api_key=OPENAI_API_KEY,
            temperature=0, max_tokens=4000,
        )
        self.planner_with_tools = self.planner_llm.bind_tools(self.tools)

        # Responder only explains — default tokens fine
        self.responder_llm = ChatOpenAI(
            model=self.model_name, streaming=True, api_key=OPENAI_API_KEY,
            temperature=0, max_tokens=1500,
        )

        self._compiled = None
        logger.info(f"✅ DiagramGraph initialized with model: {self.model_name} (planner: 4000 tokens)")

    # ─── Node 1: PLANNER ───
    async def _planner_node(self, state: DiagramState) -> dict:
        """Planner AI: analyze request, generate Mermaid code via 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 ───
    async def _responder_node(self, state: DiagramState) -> dict:
        """Responder AI: explain the generated diagram."""
        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 ───
    def _after_planner(self, state: DiagramState) -> str:
        """After Planner: if tool_calls → go tools, else → end (ask clarification)."""
        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(DiagramState)

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

        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("✅ DiagramGraph 2-Agent compiled")
        return self._compiled

    async def chat(self, user_message: str, history: list[BaseMessage] | None = None) -> dict:
        """Main entry — send message and get response with diagram."""
        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]] = []
        diagram_data = None  # Mermaid code for frontend

        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 "(generating diagram...)",
                        "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 result — extract Mermaid code
                tool_content = content_str
                try:
                    import json
                    parsed = json.loads(tool_content)
                    if parsed.get("status") == "success":
                        diagram_data = {
                            "mermaid_code": parsed.get("mermaid_code", ""),
                            "title": parsed.get("title", ""),
                            "diagram_type": parsed.get("diagram_type", "flowchart"),
                        }
                        pipeline.append({
                            "step": "tool_result",
                            "label": "📊 Diagram Generated",
                            "content": f"✅ Đã tạo {parsed.get('diagram_type', 'diagram')}: {parsed.get('title', 'Untitled')}",
                        })
                    else:
                        pipeline.append({
                            "step": "tool_result",
                            "label": "❌ Tool Error",
                            "content": parsed.get("message", "Unknown error"),
                        })
                except Exception:
                    pipeline.append({
                        "step": "tool_result",
                        "label": "📊 Tool Result",
                        "content": tool_content[:300],
                    })

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

        logger.info(
            "📊 [DIAGRAM] query='%s' | path=%s | time=%.2fms",
            user_message[:50], agent_path, 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,
            "diagram": diagram_data,
        }


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


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