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

import hashlib
import logging
from datetime import datetime
from typing import Any

from langchain_core.language_models import BaseChatModel
from langchain_core.messages import HumanMessage
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_utils import get_system_prompt_template
from .tools.get_tools import get_all_tools, get_collection_tools

logger = logging.getLogger(__name__)


class CANIFAGraph:
    """
    Fashion Q&A 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)
        self.cache = InMemoryCache()

        # Chain caching: avoid rebuilding ChatPromptTemplate every turn
        self._cached_chain = None
        self._cached_prompt_hash: str | None = None

    def _build_chain(self, system_prompt_template: str):
        """Build chain with dynamic system prompt (fetched from Langfuse per request).
        Caches the chain and only rebuilds when prompt content changes.
        """
        # Check if prompt changed via hash comparison
        prompt_hash = hashlib.md5(system_prompt_template.encode()).hexdigest()

        if self._cached_chain is not None and self._cached_prompt_hash == prompt_hash:
            return self._cached_chain

        logger.info("🔄 Building new chain (prompt hash changed)" if self._cached_chain else "🔧 Building initial chain")

        prompt_template = ChatPromptTemplate.from_messages(
            [
                ("system", system_prompt_template),
                (
                    "system",
                    "===== USER INSIGHT (TỪ TURN TRƯỚC) =====\n⚡ BẮT BUỘC: Đọc [LAST_ACTION] bên dưới để hiểu context turn trước, rồi tự suy ra bước tiếp theo!\n\n{user_insight}\n=====================================",
                ),
                ("system", "Chat History:"),
                MessagesPlaceholder(variable_name="history"),
                ("system", "Current Query:"),
                MessagesPlaceholder(variable_name="user_query"),
                MessagesPlaceholder(variable_name="messages"),
            ]
        )
        self._cached_chain = prompt_template | self.llm_with_tools
        self._cached_prompt_hash = prompt_hash
        return self._cached_chain

    async def _agent_node(self, state: AgentState, config: RunnableConfig) -> dict:
        """Agent node - Chỉ việc đổ dữ liệu riêng vào khuôn đã có sẵn."""
        messages = state.get("messages", [])
        history = state.get("history", [])
        user_query = state.get("user_query")

        transient_images = config.get("configurable", {}).get("transient_images", [])
        if transient_images and user_query:
            # Inject images into user_query as multimodal content
            # GPT-4o / GPT-5 vision: text + image_url blocks
            text_content = user_query.content if isinstance(user_query.content, str) else str(user_query.content)
            # Append hint so LLM knows images are attached
            text_content += "\n\n[📸 Có ảnh sản phẩm kèm theo - Hãy mô tả CHI TIẾT sản phẩm trong ảnh (loại, màu, kiểu dáng, hình in, chất liệu) rồi gọi data_retrieval_tool để tìm sản phẩm tương tự.]"
            multimodal_content = [{"type": "text", "text": text_content}]
            for img in transient_images:
                # Support both base64 and URL
                if img.startswith("data:"):
                    image_url = img
                elif img.startswith("http"):
                    image_url = img
                else:
                    image_url = f"data:image/jpeg;base64,{img}"
                multimodal_content.append({
                    "type": "image_url",
                    "image_url": {"url": image_url, "detail": "auto"}
                })
                logger.info(f"📸 [IMAGE] Image size: {len(img)} chars")
            user_query = HumanMessage(content=multimodal_content)
            logger.info(f"📸 [IMAGE] Injected {len(transient_images)} image(s) into user_query")
        # Invoke chain with history, user_query, messages (scratchpad), and user_insight
        user_insight_text = (
            state.get("user_insight")
            or "⚠️ TRẠNG THÁI KHỞI TẠO: Chưa có User Insight từ lịch sử. Hãy bắt đầu thu thập thông tin mới (Nếu thiếu thông tin thì ghi 'Chưa rõ')."
        )

        # Fetch prompt from Langfuse (cached 5min by SDK) + build chain (cached by hash)
        current_date_str = datetime.now().strftime("%d/%m/%Y")
        system_prompt_template = get_system_prompt_template()
        chain = self._build_chain(system_prompt_template)

        response = await chain.ainvoke(
            {
                "date_str": current_date_str,
                "history": history,
                "user_query": [user_query] if user_query else [],
                "messages": messages,
                "user_insight": user_insight_text,
            }
        )
        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: list[CANIFAGraph | None] = [None]


def build_graph(config: AgentConfig | None = None, llm: BaseChatModel | None = None, tools: list | None = None) -> Any:
    """Get compiled graph (Singleton usage)."""
    # Use singleton to avoid rebuilding graph on every request
    manager = get_graph_manager(config, llm, tools)
    return manager.build()


def get_graph_manager(
    config: AgentConfig | None = None, llm: BaseChatModel | None = None, tools: list | None = None
) -> CANIFAGraph:
    """Get CANIFAGraph instance (Auto-rebuild if model config changes).
    
    Prompt is now fetched dynamically per request from Langfuse,
    so no need to rebuild graph when prompt changes.
    """
    # 1. New Instance if Empty
    if _instance[0] is None:
        _instance[0] = CANIFAGraph(config, llm, tools)
        logger.info(f"✨ Graph Created: {_instance[0].config.model_name} (prompts from Langfuse)")
        return _instance[0]

    # 2. Check for Model Config Changes only
    is_model_changed = config and config.model_name != _instance[0].config.model_name

    if is_model_changed:
        logger.info(f"🔄 Rebuilding Graph: Model ({_instance[0].config.model_name}->{config.model_name})")
        _instance[0] = CANIFAGraph(config, llm, tools)
        return _instance[0]

    return _instance[0]


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


def reset_chain_cache() -> None:
    """Reset only the cached chain (when prompt changes).
    Keeps the graph/LLM/tools intact, only forces chain rebuild on next request.
    """
    if _instance[0] is not None:
        _instance[0]._cached_chain = None
        _instance[0]._cached_prompt_hash = None
        logger.info("🔄 Chain cache cleared — will rebuild on next request")
