"""
CuCu Note Agent Graph
LangGraph workflow với clean architecture.
Tất cả resources (LLM, Tools) khởi tạo trong __init__.
Sử dụng ConversationManager (MongoDB) để lưu history thay vì checkpoint.
"""

import logging
from typing import Any

from langchain_core.language_models import BaseChatModel
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableConfig
from langgraph.cache.memory import InMemoryCache
from langgraph.graph import END, StateGraph
from langgraph.prebuilt import ToolNode
from langgraph.types import CachePolicy

from common.llm_factory import create_llm

from .models import AgentConfig, AgentState, get_config
from .prompt import get_system_prompt
from .tools.get_tools import get_all_tools, get_collection_tools

logger = logging.getLogger(__name__)


class CuCuGraph:
    """
    CuCu Note Agent Graph Manager.
    """

    def __init__(
        self,
        config: AgentConfig | None = None,
        llm: BaseChatModel | None = None,
        tools: list | None = None,
    ):
        self.config = config or get_config()
        self._compiled_graph: Any | None = None
        self.llm: BaseChatModel = llm or create_llm(
            model_name=self.config.model_name, api_key=self.config.openai_api_key, streaming=True
        )
        self.all_tools = tools or get_all_tools()
        self.collection_tools = get_collection_tools()  # Vẫn lấy list name để routing
        self.retrieval_tools = self.all_tools

        self.llm_with_tools = self.llm.bind_tools(self.all_tools, strict=True)
        # NOTE: prompt is NOT cached here — fetched fresh each request
        # so Langfuse updates take effect immediately.
        self.cache = InMemoryCache()

    def _build_chain(self):
        """Build chain with fresh system prompt (from Langfuse or local fallback)."""
        system_prompt = get_system_prompt()
        prompt_template = ChatPromptTemplate.from_messages(
            [
                ("system", system_prompt),
                MessagesPlaceholder(variable_name="history"),
                MessagesPlaceholder(variable_name="user_query"),
                MessagesPlaceholder(variable_name="messages"),
            ]
        )
        return prompt_template | self.llm_with_tools

    async def _agent_node(self, state: AgentState, config: RunnableConfig) -> dict:
        """Agent node — rebuilds chain each call for realtime prompt updates."""
        messages = state.get("messages", [])
        history = state.get("history", [])
        user_query = state.get("user_query")

        chain = self._build_chain()
        response = await chain.ainvoke({
            "user_query": [user_query] if user_query else [],
            "history": history,
            "messages": messages
        })
        return {"messages": [response], "ai_response": response}


    def _should_continue(self, state: AgentState) -> str:
        """Routing: tool nodes hoặc end."""
        last_message = state["messages"][-1]

        if not hasattr(last_message, "tool_calls") or not last_message.tool_calls:
            logger.info("🏁 Agent finished")
            return "end"

        tool_names = [tc["name"] for tc in last_message.tool_calls]
        collection_names = [t.name for t in self.collection_tools]

        if any(name in collection_names for name in tool_names):
            logger.info(f"🔄 → collect_tools: {tool_names}")
            return "collect_tools"

        logger.info(f"🔄 → retrieve_tools: {tool_names}")
        return "retrieve_tools"

    def build(self) -> Any:
        """Build và compile LangGraph workflow."""
        if self._compiled_graph is not None:
            return self._compiled_graph
        workflow = StateGraph(AgentState)

        # Nodes
        workflow.add_node("agent", self._agent_node)
        workflow.add_node("retrieve_tools", ToolNode(self.retrieval_tools), cache_policy=CachePolicy(ttl=3600))
        workflow.add_node("collect_tools", ToolNode(self.collection_tools))

        # Edges
        workflow.set_entry_point("agent")
        workflow.add_conditional_edges(
            "agent",
            self._should_continue,
            {"retrieve_tools": "retrieve_tools", "collect_tools": "collect_tools", "end": END},
        )
        workflow.add_edge("retrieve_tools", "agent")
        workflow.add_edge("collect_tools", "agent")

        self._compiled_graph = workflow.compile(cache=self.cache)  # No Checkpointer
        logger.info("✅ Graph compiled (Langfuse callback will be per-run)")

        return self._compiled_graph

    @property
    def graph(self) -> Any:
        return self.build()


# --- Singleton & Public API ---
_instance: CuCuGraph | None = None

def build_graph(config: AgentConfig | None = None, llm: BaseChatModel | None = None, tools: list | None = None) -> Any:
    """Get compiled graph (singleton)."""
    global _instance
    if _instance is None:
        _instance = CuCuGraph(config, llm, tools)
    return _instance.build()


def get_graph_manager(
    config: AgentConfig | None = None, llm: BaseChatModel | None = None, tools: list | None = None
) -> CuCuGraph:
    """Get CuCuGraph instance."""
    global _instance
    if _instance is None:
        _instance = CuCuGraph(config, llm, tools)
    return _instance


def reset_graph() -> None:
    """Reset singleton for testing."""
    global _instance
    _instance = None
