import json
import logging
import time
from langchain_core.runnables import RunnableConfig
from agent.nodes.models_node import StylistProState

from agent.tools.data_retrieval_tool import data_retrieval_tool

logger = logging.getLogger(__name__)


def _get_tool_registry():
    from agent.tools.brand_knowledge_tool import canifa_knowledge_search
    from agent.tools.check_is_stock import check_is_stock
    from agent.tools.promotion_canifa_tool import canifa_get_promotions
    from agent.tools.store_search_tool import canifa_store_search

    return {
        "lead_search_tool": data_retrieval_tool,
        "data_retrieval_tool": data_retrieval_tool,
        "canifa_knowledge_search": canifa_knowledge_search,
        "canifa_get_promotions": canifa_get_promotions,
        "canifa_store_search": canifa_store_search,
        "check_is_stock": check_is_stock,
    }

async def tools_node(state: StylistProState, config: RunnableConfig):
    """
    Tools Node: Thực thi Tool dựa trên tool_name_used và tool_args từ Classifier.
    """
    tool_name = state.get("tool_name_used")
    tool_args = state.get("tool_args") or {}
    tool_result = None
    diagnostics = []

    if not tool_name:
        return {"tool_result": None, "diagnostics": []}

    tool_registry = _get_tool_registry()

    tool_fn = tool_registry.get(tool_name)
    if tool_fn:
        t_start = time.time()
        try:
            logger.info(f"🛠️ [ToolsNode] Executing '{tool_name}' with args: {tool_args}")
            
            # Execute Tool
            tool_result = await tool_fn.ainvoke(tool_args)
            tool_elapsed = (time.time() - t_start) * 1000
            
            diagnostics.append({
                "step": "tool_result",
                "label": f"📦 {tool_name}",
                "content": f"Success: {str(tool_result)[:100]}...",
                "raw_json": tool_result,
                "elapsed_ms": round(tool_elapsed)
            })
            logger.info(f"✅ [ToolsNode] '{tool_name}' finished in {round(tool_elapsed)}ms")
            
        except Exception as e:
            logger.error(f"❌ [ToolsNode] Execution Error in '{tool_name}': {e}")
            tool_result = json.dumps({"status": "error", "message": str(e)})
            diagnostics.append({
                "step": "tool_error",
                "label": f"❌ {tool_name} Error",
                "content": str(e),
                "elapsed_ms": 0
            })
    else:
        logger.warning(f"⚠️ [ToolsNode] Unknown tool: '{tool_name}'")

    return {
        "tool_result": tool_result,
        "diagnostics": diagnostics
    }
