import json
import logging
import time
from typing import Dict, Any
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, StylistOutput
from agent.nodes.utils import _extract_skus_from_text, _dedupe_preserve_order, _format_tool_result, parse_llm_json
from agent.prompt_module.stylist_pro_prompts import STYLIST_PRO_RESPONDER_PROMPT, format_stylist_pro_injection

logger = logging.getLogger(__name__)

async def stylist_node(state: StylistProState, config: RunnableConfig):
    """
    Stylist Node: Sinh câu trả lời tư vấn và cập nhật Insight khách hàng.
    """
    llm = create_llm(model_name=config.get("configurable", {}).get("model_name"), streaming=False)
    
    insight_dict = json.loads(state.get("user_insight") or "{}")
    injection = format_stylist_pro_injection(insight_dict)
    
    system_content = STYLIST_PRO_RESPONDER_PROMPT.format(stage_injection=injection)
    messages = [SystemMessage(content=system_content)]
    
    # History context (3 turns)
    human_count = 0
    recent = []
    for msg in reversed(state["messages"][:-1]):
        if isinstance(msg, HumanMessage):
            recent.insert(0, msg)
            human_count += 1
        elif isinstance(msg, AIMessage) and not (hasattr(msg, "tool_calls") and msg.tool_calls):
            recent.insert(0, msg)
        if human_count >= 3: break
    messages.extend(recent)
    messages.append(state["messages"][-1]) 

    # Tool context
    tool_context = _format_tool_result(state.get("tool_result"))
    if tool_context:
        messages.append(HumanMessage(content=f"=== KẾT QUẢ TRA CỨU ===\n{tool_context}"))

    # Extra instructions
    messages.append(HumanMessage(content="REMINDER: Output ONLY JSON. List ALL SKUs in product_ids. Ensure Top+Bottom for combos."))

    start_time = time.time()
    response = await llm.ainvoke(messages, config=config)
    elapsed = (time.time() - start_time) * 1000
    
    data = parse_llm_json(response.content)
    output = StylistOutput(**data)
    
    # Enrichment & Deduplication
    text_skus = _extract_skus_from_text(output.ai_response)
    final_product_ids = _dedupe_preserve_order((output.product_ids or []) + text_skus)
    
    # Insight mapping
    insight_raw = output.user_insight
    insight_dict = insight_raw.model_dump() if hasattr(insight_raw, 'model_dump') else insight_raw

    lead_stage = {
        "stage": insight_dict.get("STAGE_NUM", 1),
        "stage_name": insight_dict.get("STAGE", "BROWSE"),
        "tone_directive": insight_dict.get("TONE", "Friendly"),
        "behavioral_hints": insight_dict.get("BEHAVIORAL_HINTS", []),
    }

    diag = {
        "step": "stylist",
        "label": "💬 Stylist Response",
        "content": output.ai_response[:200],
        "elapsed_ms": round(elapsed)
    }
    
    return {
        "updated_insight": insight_dict,
        "lead_stage": lead_stage,
        "stage_injection": format_stylist_pro_injection(insight_dict),
        "product_ids": final_product_ids,
        "messages": [AIMessage(content=output.ai_response)],
        "diagnostics": [diag]
    }
