import json
import logging
import time
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_core.runnables import RunnableConfig

from common.llm_factory import create_llm
from agent.nodes.models_node import StylistProState, ClassifierOutput
from agent.nodes.utils import _extract_text, _format_history, parse_llm_json
from agent.prompt_module.stylist_pro_prompts import STYLIST_PRO_CLASSIFIER_PROMPT

logger = logging.getLogger(__name__)

async def classifier_node(state: StylistProState, config: RunnableConfig):
    """
    Classifier Node: Định tuyến intent và thực thi Tool inline.
    Sử dụng data_retrieval_tool làm công cụ tìm kiếm chính.
    """
    llm = create_llm(model_name=config.get("configurable", {}).get("model_name"), streaming=False)
    
    messages = state["messages"]
    user_query = _extract_text(messages[-1].content)
    user_insight = state.get("user_insight") or "{}"
    history_text = _format_history(messages[:-1])

    # Context preparation
    products_context = ""
    try:
        if isinstance(user_insight, str) and user_insight.startswith("{"):
            _ins = json.loads(user_insight)
            latest = _ins.get("LATEST_PRODUCT_INTEREST", "")
            summary = _ins.get("SUMMARY_HISTORY", "")
            if latest: products_context += f"\n=== SP ĐÃ GIỚI THIỆU ===\n{latest}\n"
            if summary: products_context += f"\n=== TÓM TẮT HỘI THOẠI ===\n{summary}\n"
    except Exception: pass

    context_block = (
        f"=== TIN NHẮN HIỆN TẠI ===\n{user_query}\n\n"
        f"=== USER INSIGHT ===\n{user_insight}\n\n"
        f"{products_context}"
        f"=== LỊCH SỬ ===\n{history_text}"
    )

    prompt = [
        SystemMessage(content=STYLIST_PRO_CLASSIFIER_PROMPT),
        HumanMessage(content=context_block),
    ]

    start_time = time.time()
    response = await llm.ainvoke(prompt, config=config)
    elapsed = (time.time() - start_time) * 1000
    
    data = parse_llm_json(response.content)
    
    # --- PHÒNG THỦ: Dọn dẹp dữ liệu trước khi nạp vào Pydantic ---
    if "search_args" in data and "inferred" in data["search_args"]:
        inferred = data["search_args"]["inferred"]
        # Các trường mong đợi String/Int nhưng hay bị AI trả về List rỗng []
        str_fields = ["gender_target", "age_group", "master_color", "price_min", "price_max"]
        for field in str_fields:
            if field in inferred and (isinstance(inferred[field], list) or isinstance(inferred[field], dict)):
                if isinstance(inferred[field], list):
                    inferred[field] = inferred[field][0] if len(inferred[field]) > 0 else None
                else:
                    inferred[field] = None
        
        # Xử lý product_line_vn (đảm bảo luôn là List)
        if "product_line_vn" in inferred and isinstance(inferred["product_line_vn"], str):
            inferred["product_line_vn"] = [inferred["product_line_vn"]]
    
    output = ClassifierOutput(**data)
    
    # Tool Registry
    tool_name = output.tool_name
    tool_result = None
    diagnostics = [{
        "step": "classifier",
        "label": f"🎯 Classifier -> {tool_name or 'Early Exit'}",
        "content": output.reasoning,
        "elapsed_ms": round(elapsed)
    }]

    # Prepare tool_args
    tool_args = {}
    if tool_name:
        # Case A: Search Tools (Data Retrieval) - UPDATED FOR DUAL-LANE
        if (tool_name == "lead_search_tool" or tool_name == "data_retrieval_tool") and output.search_args:
            inf = output.search_args.inferred
            tool_args = {
                "searches": [
                    {
                        "literal": {"raw_text": output.search_args.literal.raw_text},
                        "inferred": {
                            "product_line_vn": inf.product_line_vn or [],
                            "gender_by_product": inf.gender_target,
                            "age_by_product": inf.age_group,
                            "master_color": inf.master_color,
                            "price_min": inf.price_min,
                            "price_max": inf.price_max,
                            "discount_min": 0,
                            "discount_max": 100,
                            "discovery_mode": None
                        }
                    }
                ]
            }
        
        # Case B: Store Search (Ensure 'location' exists)
        elif tool_name == "canifa_store_search":
            loc = (output.tool_args or {}).get("location") or user_query
            tool_args = {"location": loc}
        
        # Case C: Knowledge Search (Ensure 'query' exists)
        elif tool_name == "canifa_knowledge_search":
            q = (output.tool_args or {}).get("query") or user_query
            tool_args = {"query": q}
        
        # Case D: Other Specific Tools
        elif output.tool_args:
            tool_args = output.tool_args

    return {
        "tool_name_used": tool_name,
        "tool_args": tool_args,
        "early_exit": tool_name is None and output.ai_response is not None,
        "product_ids": output.product_ids,
        "messages": [AIMessage(content=output.ai_response)] if output.ai_response else [],
        "diagnostics": diagnostics
    }
