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

import asyncio
import json
import logging
import re
import time
import uuid

from fastapi import BackgroundTasks
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.runnables import RunnableConfig
from langfuse import get_client as get_langfuse, observe, propagate_attributes

from common.cache import redis_cache
from common.conversation_manager import get_conversation_manager
from config import DEFAULT_MODEL, REDIS_CACHE_TURN_ON

from .controller_helpers import (
    extract_and_save_user_insight,
    finalize_user_insight_after_stream,
    load_user_insight_from_redis,
)
from .graph import build_graph
from .helper import extract_product_ids, handle_post_chat_async
from .models import AgentState, get_config
from .streaming_callback import ProductIDStreamingCallback

logger = logging.getLogger(__name__)

@observe(name="chat-request")
async def chat_controller(
    query: str,
    identity_key: str,
    background_tasks: BackgroundTasks,
    model_name: str = DEFAULT_MODEL,
    images: list[str] | None = None,
    return_user_insight: bool = False,
    is_authenticated: bool = False,  # New: để track user vs guest
    device_id: str | None = None,  # Device ID từ request (có cả khi authenticated)
) -> dict:
    """
    Controller main logic for non-streaming chat requests.

    Flow:
    1. Check cache (if enabled) → HIT: return cached response
    2. MISS: Call LLM → Save to cache → Return response
    3. user_insight extract:
       - return_user_insight=False (prod): background task (non-blocking)
       - return_user_insight=True (dev): extract ngay và return luôn

    Args:
        identity_key: Key for saving/loading history (identity.history_key)
                      Guest: device_id, User: user_id
        return_user_insight: If True, extract and return user_insight immediately (dev mode)
        is_authenticated: True if user logged in, False if guest
    """

    logger.info("chat_controller start: model=%s, identity_key=%s, auth=%s", model_name, identity_key, is_authenticated)
    if images:
        logger.info("📸 [CONTROLLER] Received %d image(s), sizes: %s", len(images), [len(img) for img in images])
    else:
        logger.debug("📸 [CONTROLLER] No images in request")

    # ====================== CACHE LAYER ======================
    if REDIS_CACHE_TURN_ON:
        cached_response = await redis_cache.get_response(user_id=identity_key, query=query)
        if cached_response:
            logger.info(f"⚡ CACHE HIT for identity_key={identity_key}")
            cached_insight = None
            if identity_key:
                cached_insight = await load_user_insight_from_redis(identity_key)
                if cached_insight:
                    cached_response = {**cached_response, "user_insight": cached_insight}
            memory = await get_conversation_manager()
            background_tasks.add_task(
                handle_post_chat_async,
                memory=memory,
                identity_key=identity_key,
                human_query=query,
                ai_response=cached_response,
            )
            return {**cached_response, "cached": True}

    config = get_config()
    config.model_name = model_name

    graph = build_graph(config)

    memory = await get_conversation_manager()

    history_dicts = await memory.get_chat_history(identity_key, limit=15, include_product_ids=False)
    messages = [
        HumanMessage(content=m["message"]) if m["is_human"] else AIMessage(content=m["message"]) for m in history_dicts
    ][::-1]

    # ====================== GRAPH EXECUTION ======================
    user_insight = None
    if identity_key:
        user_insight = await load_user_insight_from_redis(identity_key)

    initial_state: AgentState = {
        "user_query": HumanMessage(content=query),
        "history": messages,
        "user_id": identity_key,
        "user_insight": user_insight,
        "images_embedding": [],
        "ai_response": None,
    }

    run_id = str(uuid.uuid4())
    session_id = f"{identity_key}-{run_id[:8]}"

    # Set Langfuse trace metadata (nests under @observe parent)
    langfuse = get_langfuse()
    langfuse.update_current_trace(
        user_id=identity_key,
        session_id=session_id,
        tags=["chatbot", "user:authenticated" if is_authenticated else "user:anonymous"],
        metadata={
            "device_id": device_id,
            "customer_id": identity_key if is_authenticated else None,
            "model": model_name,
        },
    )

    # Get handler from context → nested under same trace as @observe
    try:
        langfuse_handler = langfuse.get_current_langchain_handler()
    except Exception:
        langfuse_handler = None
        logger.warning("⚠️ Could not get Langfuse langchain handler")

    streaming_callback = ProductIDStreamingCallback()

    exec_config = RunnableConfig(
        configurable={
            "user_id": identity_key,
            "transient_images": images or [],
            "run_id": run_id,
        },
        run_id=run_id,
        metadata={
            "session_id": session_id,
            "device_id": device_id,
        },
        callbacks=[langfuse_handler, streaming_callback] if langfuse_handler else [streaming_callback],
    )

    logger.info("🌊 Starting LLM streaming...")

    event_count = 0
    start_time = time.time()
    ai_text_response = ""
    final_product_ids = []
    all_accumulated_messages = []
    seen_message_ids = set()

    async def consume_events():
        nonlocal event_count, ai_text_response, final_product_ids, all_accumulated_messages, seen_message_ids

        async for event in graph.astream(initial_state, config=exec_config):
            event_count += 1
            elapsed = time.time() - start_time
            logger.info(f"📦 Event #{event_count} at t={elapsed:.2f}s | Keys: {list(event.keys())}")

            # Accumulate messages from ALL events (deduplicate by ID)
            for node_name, node_state in event.items():
                if isinstance(node_state, dict) and "messages" in node_state:
                    for msg in node_state["messages"]:
                        msg_id = id(msg)
                        if msg_id not in seen_message_ids:
                            seen_message_ids.add(msg_id)
                            all_accumulated_messages.append(msg)

    early_response = False

    stream_task = asyncio.create_task(consume_events())

    if return_user_insight:
        # Dev mode: Đợi stream xong để có đầy đủ user_insight
        logger.info("🔍 [DEV MODE] Waiting for full stream...")
        await stream_task
        ai_text_response = streaming_callback.ai_response_text
        final_product_ids = streaming_callback.product_skus
        logger.info("✅ Stream completed in dev mode")
    else:
        # Prod mode: Return ngay khi có product_ids
        wait_task = asyncio.create_task(streaming_callback.product_found_event.wait())

        done, pending = await asyncio.wait(
            [stream_task, wait_task],
            return_when=asyncio.FIRST_COMPLETED,
            timeout=30.0,
        )

        # bắn trả về ngay khi có product_ids hoặc timeout
        ai_text_response = streaming_callback.ai_response_text
        final_product_ids = streaming_callback.product_skus

        elapsed = time.time() - start_time
        if streaming_callback.product_ids_found:
            logger.info(f"⚡ Response ready at t={elapsed:.2f}s (early return)")
            early_response = True
        else:
            logger.info(f"✅ Stream completed at t={elapsed:.2f}s (no products)")

        # Cleanup (never cancel stream_task; we still need it)
        if wait_task in pending:
            wait_task.cancel()

    try:
        # Ensure stream completes to get all tool messages (skip if early response)
        if not early_response and not stream_task.done():
            logger.info("⏳ Waiting for stream to complete to extract tool products...")
            try:
                await stream_task
            except asyncio.CancelledError:
                logger.warning("⚠️ stream_task was cancelled; skip tool product extraction")

        # Early return path: try to extract ai_response directly from streaming content
        if early_response and not ai_text_response:
            raw_content = streaming_callback.accumulated_content
            if raw_content:
                # Try JSON parse first
                try:
                    ai_json = json.loads(raw_content)
                    if isinstance(ai_json, dict):
                        ai_text_response = ai_json.get("ai_response", ai_text_response)
                        if not final_product_ids and isinstance(ai_json.get("product_ids"), list):
                            final_product_ids = [str(s) for s in ai_json.get("product_ids", [])]
                        logger.info("ℹ️ Extracted ai_response from streaming JSON content")
                except json.JSONDecodeError:
                    # Fallback regex: extract ai_response regardless of field order
                    ai_match = re.search(r'"ai_response"\s*:\s*"((?:[^"\\]|\\.)*)"', raw_content)
                    if ai_match:
                        ai_text_response = ai_match.group(1).replace('\\"', '"').replace("\\n", "\n")
                        logger.info("ℹ️ Extracted ai_response via regex from streaming content")

        # Fallback: If ai_response is empty, take last AI message from accumulated events
        if not ai_text_response and all_accumulated_messages:
            for msg in reversed(all_accumulated_messages):
                if isinstance(msg, AIMessage) and msg.content:
                    ai_text_response = msg.content
                    logger.info("ℹ️ Filled ai_response from last AIMessage in accumulated events")
                    break

        # If ai_response looks like JSON, extract ai_response field directly
        if ai_text_response and ai_text_response.lstrip().startswith("{"):
            try:
                ai_json = json.loads(ai_text_response)
                if isinstance(ai_json, dict):
                    ai_text_response = ai_json.get("ai_response", ai_text_response)
                    if not final_product_ids and isinstance(ai_json.get("product_ids"), list):
                        final_product_ids = [str(s) for s in ai_json.get("product_ids", [])]
                    logger.info("ℹ️ Extracted ai_response from JSON content")
            except json.JSONDecodeError:
                logger.warning("⚠️ ai_response looked like JSON but failed to parse")

        # ✅ Extract products
        # PRIORITY: Use tool results from accumulated_messages (has correct color)
        # FALLBACK: Only use fetch_products_by_skus if no tool results available
        enriched_products = []  # Initialize before if block
        
        if all_accumulated_messages:
            all_extracted_products = extract_product_ids(all_accumulated_messages)
            logger.info(
                f"✅ Extracted {len(all_extracted_products)} products from {len(all_accumulated_messages)} accumulated messages"
            )

            # 🎯 FILTER: Only return products that AI mentioned in product_ids
            if final_product_ids:
                # AI mentioned specific SKUs → Filter to match only those
                mentioned_skus = set(final_product_ids)
                enriched_products = [p for p in all_extracted_products if p.get("sku") in mentioned_skus]
                logger.warning(
                    f"🎯 AI mentioned {len(mentioned_skus)} SKUs → Filtered {len(all_extracted_products)} down to {len(enriched_products)} products"
                )
            else:
                # AI didn't mention specific SKUs → Return all extracted products
                enriched_products = all_extracted_products
                logger.info("ℹ️ No specific SKUs mentioned by AI, returning all extracted products")
        else:
            logger.warning("⚠️ No accumulated messages found - returning empty product list")

        if not enriched_products:
            logger.warning("⚠️ No products extracted from tool messages - returning empty list")

    except Exception as e:
        logger.error(f"❌ Error extracting products from tool messages: {e}")
        enriched_products = []

    user_insight_dict = None

    if return_user_insight:
        callback_accumulated_content = streaming_callback.accumulated_content
        if callback_accumulated_content and identity_key:
            logger.info("🔍 [DEV] Extracting user_insight...")
            user_insight_dict = await extract_and_save_user_insight(callback_accumulated_content, identity_key)
    elif background_tasks:
        # Prod/Dev fast mode: để stream chạy xong rồi mới extract trong background
        logger.info("💾 [PROD] Scheduling background task for user_insight extraction (post-stream)")
        background_tasks.add_task(
            finalize_user_insight_after_stream,
            stream_task=stream_task,
            streaming_callback=streaming_callback,
            identity_key=identity_key,
            trace_id=langfuse.get_current_trace_id() if langfuse else run_id,
        )

    # ====================== BUILD RESPONSE ======================
    response_payload = {
        "ai_response": ai_text_response,
        "product_ids": enriched_products,  # ⚡ Full product objects from tool (name, price, sale_price, url, thumbnail)
    }

    if user_insight_dict is not None:
        response_payload["user_insight"] = user_insight_dict

    # ====================== SAVE TO CACHE ======================
    if REDIS_CACHE_TURN_ON:
        await redis_cache.set_response(user_id=identity_key, query=query, response_data=response_payload, ttl=300)
        logger.debug(f"💾 Cached response for identity_key={identity_key}")

    # Save to History (Background)
    background_tasks.add_task(
        handle_post_chat_async,
        memory=memory,
        identity_key=identity_key,
        human_query=query,
        ai_response=response_payload,
    )

    logger.info("chat_controller finished: cached=False")
    return {**response_payload, "cached": False}
