"""
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 agent.tools.data_retrieval_tool import SearchItem, data_retrieval_tool
from common.conversation_manager import ConversationManager, get_conversation_manager
from common.langfuse_client import get_callback_handler
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

    tools = get_all_tools()
    graph = build_graph(config, llm=None, 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:
        # 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}")

        # 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 for Langfuse (if enabled)
    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()

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

    try:
        # Gọi tool trực tiếp (không qua LLM) để tránh bottleneck
        search_item = SearchItem(
            query=query or "sản phẩm",
            magento_ref_code=None,
            price_min=None,
            price_max=None,
            action="search",
            # Metadata fields - all required with None default
            gender_by_product=None,
            age_by_product=None,
            product_name=None,
            style=None,
            master_color=None,
            season=None,
            material_group=None,
            fitting=None,
            form_neckline=None,
            form_sleeve=None,
        )

        result_json = await data_retrieval_tool.ainvoke({"searches": [search_item]})
        result = json.loads(result_json)
        all_product_ids = result.get("results", [{}])[0].get("products", [])

        # 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
