"""
Fashion Q&A Agent Controller
Langfuse will auto-trace via LangChain integration (no code changes needed).
"""

import asyncio
import json
import logging
import random
import time
import uuid

from fastapi import BackgroundTasks
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from langchain_core.runnables import RunnableConfig

from common.conversation_manager import ConversationManager, get_conversation_manager
from common.langfuse_client import get_callback_handler, langfuse_trace_context
from common.llm_factory import create_llm
from config import DEFAULT_MODEL

from .graph import build_graph
from .models import AgentState, get_config
from .tools.get_tools import get_all_tools

logger = logging.getLogger(__name__)

# --- MOCK LLM RESPONSES (không gọi OpenAI) ---
MOCK_AI_RESPONSES = [
    "Dựa trên tìm kiếm của bạn, tôi tìm thấy các sản phẩm phù hợp với nhu cầu của bạn. Những mặt hàng này có chất lượng tốt và giá cả phải chăng.",
    "Tôi gợi ý cho bạn những sản phẩm sau. Chúng đều là những lựa chọn phổ biến và nhận được đánh giá cao từ khách hàng.",
    "Dựa trên tiêu chí tìm kiếm của bạn, đây là những sản phẩm tốt nhất mà tôi có thể giới thiệu.",
    "Những sản phẩm này hoàn toàn phù hợp với yêu cầu của bạn. Hãy xem chi tiết để chọn sản phẩm yêu thích nhất.",
    "Tôi đã tìm được các mặt hàng tuyệt vời cho bạn. Hãy kiểm tra chúng để tìm ra lựa chọn tốt nhất.",
]


async def chat_controller(
    query: str,
    user_id: str,
    background_tasks: BackgroundTasks,
    model_name: str = DEFAULT_MODEL,
    images: list[str] | None = None,
) -> dict:
    """
    Controller main logic for non-streaming chat requests.
    Langfuse will automatically trace all LangChain operations.
    """
    logger.info(f"▶️ Starting chat_controller with model: {model_name} for user: {user_id}")

    config = get_config()
    config.model_name = model_name

    # Enable JSON mode to ensure structured output
    llm = create_llm(model_name=model_name, streaming=False, json_mode=True)
    tools = get_all_tools()

    graph = build_graph(config, llm=llm, tools=tools)

    # Init ConversationManager (Singleton)
    memory = await get_conversation_manager()

    # LOAD HISTORY & Prepare State (Optimize: history logic remains solid)
    history_dicts = await memory.get_chat_history(user_id, limit=20)

    history = []
    for h in reversed(history_dicts):
        msg_cls = HumanMessage if h["is_human"] else AIMessage
        history.append(msg_cls(content=h["message"]))

    initial_state, exec_config = _prepare_execution_context(
        query=query, user_id=user_id, history=history, images=images
    )

    try:
        # 🔥 Wrap graph execution với langfuse_trace_context để set user_id cho tất cả observations
        with langfuse_trace_context(user_id=user_id, session_id=user_id):
            # TỐI ƯU: Chạy Graph
            result = await graph.ainvoke(initial_state, config=exec_config)

            # TỐI ƯU: Extract IDs từ Tool Messages một lần duy nhất
            all_product_ids = _extract_product_ids(result.get("messages", []))

            # TỐI ƯU: Xử lý AI Response
            ai_raw_content = result.get("ai_response").content if result.get("ai_response") else ""
            logger.info(f"💾 [RAW AI OUTPUT]:\n{ai_raw_content}")

            # Parse JSON để lấy text response và product_ids từ AI
            ai_text_response = ai_raw_content
            try:
                # Vì json_mode=True, OpenAI sẽ nhả raw JSON
                ai_json = json.loads(ai_raw_content)

                # Extract text response từ JSON
                ai_text_response = ai_json.get("ai_response", ai_raw_content)

                # Merge product_ids từ AI JSON (nếu có) - KHÔNG dùng set() vì dict unhashable
                explicit_ids = ai_json.get("product_ids", [])
                if explicit_ids and isinstance(explicit_ids, list):
                    # Merge và deduplicate by SKU
                    seen_skus = {p["sku"] for p in all_product_ids if "sku" in p}
                    for product in explicit_ids:
                        if isinstance(product, dict) and product.get("sku") not in seen_skus:
                            all_product_ids.append(product)
                            seen_skus.add(product.get("sku"))
            except (json.JSONDecodeError, Exception) as e:
                # Nếu AI trả về text thường (hiếm khi xảy ra trong JSON mode) thì ignore
                logger.warning(f"Could not parse AI response as JSON: {e}")
                pass

            # BACKGROUND TASK: Lưu history nhanh gọn
            background_tasks.add_task(
                _handle_post_chat_async,
                memory=memory,
                user_id=user_id,
                human_query=query,
                ai_msg=AIMessage(content=ai_text_response),
            )

            return {
                "ai_response": ai_text_response,  # CHỈ text, không phải JSON
                "product_ids": all_product_ids,  # Array of product objects
            }

    except Exception as e:
        logger.error(f"💥 Chat error for user {user_id}: {e}", exc_info=True)
        raise


def _extract_product_ids(messages: list) -> list[dict]:
    """
    Extract full product info from tool messages (data_retrieval_tool results).
    Returns list of product objects with: sku, name, price, sale_price, url, thumbnail_image_url.
    """
    products = []
    seen_skus = set()

    for msg in messages:
        if isinstance(msg, ToolMessage):
            try:
                # Tool result is JSON string
                tool_result = json.loads(msg.content)

                # Check if tool returned products
                if tool_result.get("status") == "success" and "products" in tool_result:
                    for product in tool_result["products"]:
                        sku = product.get("internal_ref_code")
                        if sku and sku not in seen_skus:
                            seen_skus.add(sku)

                            # Extract full product info
                            product_obj = {
                                "sku": sku,
                                "name": product.get("magento_product_name", ""),
                                "price": product.get("price_vnd", 0),
                                "sale_price": product.get("sale_price_vnd"),  # null nếu không sale
                                "url": product.get("magento_url_key", ""),
                                "thumbnail_image_url": product.get("thumbnail_image_url", ""),
                            }
                            products.append(product_obj)
            except (json.JSONDecodeError, KeyError, TypeError) as e:
                logger.debug(f"Could not parse tool message for products: {e}")
                continue

    return products


def _prepare_execution_context(query: str, user_id: str, history: list, images: list | None):
    """Prepare initial state and execution config for the graph run."""
    initial_state: AgentState = {
        "user_query": HumanMessage(content=query),
        "messages": [HumanMessage(content=query)],
        "history": history,
        "user_id": user_id,
        "images_embedding": [],
        "ai_response": None,
    }
    run_id = str(uuid.uuid4())

    # Metadata for LangChain (tags for logging/filtering)
    metadata = {
        "run_id": run_id,
        "tags": "chatbot,production",
    }

    # 🔥 CallbackHandler - sẽ được wrap trong langfuse_trace_context để set user_id
    # Per Langfuse docs: propagate_attributes() handles user_id propagation
    langfuse_handler = get_callback_handler()

    exec_config = RunnableConfig(
        configurable={
            "user_id": user_id,
            "transient_images": images or [],
            "run_id": run_id,
        },
        run_id=run_id,
        metadata=metadata,
        callbacks=[langfuse_handler] if langfuse_handler else [],
    )
    return initial_state, exec_config


async def _handle_post_chat_async(
    memory: ConversationManager, user_id: str, human_query: str, ai_msg: AIMessage | None
):
    """Save chat history in background task after response is sent."""
    if ai_msg:
        try:
            await memory.save_conversation_turn(user_id, human_query, ai_msg.content)
            logger.debug(f"Saved conversation for user {user_id}")
        except Exception as e:
            logger.error(f"Failed to save conversation for user {user_id}: {e}", exc_info=True)


# ========================================
# MOCK CONTROLLER (Fake LLM - Real Tools)
# ========================================


async def mock_chat_controller(
    query: str,
    user_id: str,
    background_tasks: BackgroundTasks,
    images: list[str] | None = None,
) -> dict:
    """
    Mock Agent Controller với FAKE LLM (không gọi OpenAI):
    - Sử dụng toàn bộ graph flow từ chat_controller
    - data_retrieval_tool THẬT (retriever thật, embedding thật, products thật)
    - LLM fake (return mock response nhanh, tiết kiệm chi phí OpenAI)
    - Dùng để STRESS TEST + testing mà không tốn tiền API

    Similarities với chat_controller:
    ✅ Sử dụng graph pipeline
    ✅ Lấy history từ ConversationManager
    ✅ Extract products từ tool messages
    ✅ Save conversation history in background

    Differences từ chat_controller:
    ✅ Dùng fake LLM response thay vì gọi OpenAI
    ✅ Không cần JSON parsing (response là plain text)
    ✅ Nhanh hơn (~1-3ms giả lập LLM thay vì 1-3s real LLM)
    """
    logger.info(f"🚀 [MOCK Chat Controller] Starting with query: {query} for user: {user_id}")
    start_time = time.time()

    config = get_config()

    # KHÔNG gọi OpenAI - dùng tools THẬT nhưng fake LLM response
    tools = get_all_tools()
    graph = build_graph(config, llm=None, tools=tools)  # llm=None để skip LLM node

    # Init ConversationManager (Singleton)
    memory = await get_conversation_manager()

    # LOAD HISTORY & Prepare State
    history_dicts = await memory.get_chat_history(user_id, limit=20)

    history = []
    for h in reversed(history_dicts):
        msg_cls = HumanMessage if h["is_human"] else AIMessage
        history.append(msg_cls(content=h["message"]))

    initial_state, exec_config = _prepare_execution_context(
        query=query, user_id=user_id, history=history, images=images
    )

    try:
        with langfuse_trace_context(user_id=user_id, session_id=user_id):
            # Chạy Graph với tools THẬT
            result = await graph.ainvoke(initial_state, config=exec_config)

            # Extract products từ tool messages (tools THẬT)
            all_product_ids = _extract_product_ids(result.get("messages", []))

            # Generate FAKE LLM response (không gọi OpenAI)
            logger.info("🤖 [FAKE LLM] Generating mock response...")
            fake_llm_time = random.uniform(0.001, 0.003)  # 1-3ms fake latency
            await asyncio.sleep(fake_llm_time)  # ✅ NON-BLOCKING

            ai_text_response = random.choice(MOCK_AI_RESPONSES)
            logger.info(f"💾 [MOCK RESPONSE]: {ai_text_response}")

            # BACKGROUND TASK: Lưu history
            background_tasks.add_task(
                _handle_post_chat_async,
                memory=memory,
                user_id=user_id,
                human_query=query,
                ai_msg=AIMessage(content=ai_text_response),
            )

            elapsed_time = time.time() - start_time
            logger.info(f"✅ Mock Chat Controller completed in {elapsed_time:.3f}s")

            return {
                "status": "success",
                "ai_response": ai_text_response,  # Plain text mock response
                "product_ids": all_product_ids,  # Real products từ tools
                "total_products_found": len(all_product_ids),
                "is_mock": True,
                "processing_time_ms": round(elapsed_time * 1000, 2),
            }

    except Exception as e:
        logger.error(f"💥 Mock Chat Controller error for user {user_id}: {e}", exc_info=True)
        raise
