"""
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 Langfuse, get_client as get_langfuse

from common.cache import redis_cache
from common.conversation_manager import get_conversation_manager
from common.langfuse_client import get_callback_handler
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__)


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)
    skip_cache: bool = False,  # Skip response cache (for dev endpoint)
) -> 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, key=%s", model_name, identity_key)

    # ====================== CACHE LAYER ======================
    if REDIS_CACHE_TURN_ON and not skip_cache:
        cached_response = await redis_cache.get_response(user_id=identity_key, query=query)
        if cached_response:
            logger.info("CACHE HIT key=%s", 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}

            # --- Langfuse: log cache hit trace ---
            try:
                langfuse = get_langfuse()
                if langfuse:
                    trace_id = Langfuse.create_trace_id()
                    session_id = f"{identity_key}-cache"
                    tags = ["chatbot", "cache-hit", "user:authenticated" if is_authenticated else "user:anonymous"]
                    langfuse.trace(
                        id=trace_id,
                        name="chat-request-cache-hit",
                        user_id=identity_key,
                        session_id=session_id,
                        tags=tags,
                        input={"query": query, "user_id": identity_key},
                        output={
                            "ai_response": (cached_response.get("ai_response", "") or "")[:500],
                            "product_count": len(cached_response.get("product_ids", [])),
                            "source": "cache",
                        },
                        metadata={
                            "device_id": device_id,
                            "customer_id": identity_key if is_authenticated else None,
                            "model": model_name,
                            "cache_hit": True,
                        },
                    )
            except Exception:
                logger.debug("Failed to log cache hit trace", exc_info=True)

            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())
    trace_id = Langfuse.create_trace_id()
    session_id = f"{identity_key}-{run_id[:8]}"
    tags = ["chatbot", "user:authenticated" if is_authenticated else "user:anonymous"]

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

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

        async for event in graph.astream(initial_state, config=exec_config):
            event_count += 1
            for _node, 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

    # --- Langfuse: wrap everything in a named trace with custom trace_id ---
    langfuse = get_langfuse()
    observation_ctx = (
        langfuse.start_as_current_observation(
            as_type="span",
            name="chat-request",
            trace_context={"trace_id": trace_id},
        )
        if langfuse
        else None
    )

    if observation_ctx:
        span = observation_ctx.__enter__()
        span.update_trace(
            name="chat-request",
            user_id=identity_key,
            session_id=session_id,
            tags=tags,
            input={"query": query, "user_id": identity_key},
            metadata={
                "device_id": device_id,
                "customer_id": identity_key if is_authenticated else None,
                "model": model_name,
            },
        )

    # CallbackHandler MUST be created INSIDE the observation context
    # so it inherits the trace_id and nests under "chat-request"
    langfuse_handler = get_callback_handler()

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

    try:
        stream_task = asyncio.create_task(consume_events())

        if return_user_insight:
            await stream_task
            ai_text_response = streaming_callback.ai_response_text
            final_product_ids = streaming_callback.product_skus
        else:
            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,
            )
            ai_text_response = streaming_callback.ai_response_text
            final_product_ids = streaming_callback.product_skus

            if streaming_callback.product_ids_found:
                early_response = True

            if wait_task in pending:
                wait_task.cancel()
    finally:
        # Update trace output before closing
        if observation_ctx:
            try:
                elapsed = time.time() - start_time
                span.update_trace(
                    output={
                        "ai_response": (ai_text_response or "")[:500],
                        "product_count": len(final_product_ids) if final_product_ids else 0,
                        "elapsed_s": round(elapsed, 2),
                        "early_return": early_response,
                    },
                )
            except Exception:
                pass
            try:
                observation_ctx.__exit__(None, None, None)
            except Exception:
                pass

    try:
        # Ensure stream completes for tool messages
        if not early_response and not stream_task.done():
            try:
                await stream_task
            except asyncio.CancelledError:
                logger.warning("stream_task cancelled")

        # Extract ai_response from streaming content (fallbacks)
        if early_response and not ai_text_response:
            raw_content = streaming_callback.accumulated_content
            if raw_content:
                try:
                    raw_normalized = raw_content.replace("{{", "{").replace("}}", "}")
                    ai_json = json.loads(raw_normalized)
                    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["product_ids"]]
                except json.JSONDecodeError:
                    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")

        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
                    break

        # Parse JSON-wrapped ai_response
        if ai_text_response and ai_text_response.lstrip().startswith("{"):
            try:
                ai_normalized = ai_text_response.replace("{{", "{").replace("}}", "}")
                ai_json = json.loads(ai_normalized)
                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["product_ids"]]
            except json.JSONDecodeError:
                pass

        # Extract & filter products
        enriched_products = []
        if all_accumulated_messages:
            all_extracted = extract_product_ids(all_accumulated_messages)
            if final_product_ids:
                mentioned = set(final_product_ids)
                enriched_products = [p for p in all_extracted if p.get("sku") in mentioned]
            else:
                enriched_products = all_extracted

    except Exception as e:
        logger.error("Error extracting products: %s", e)
        enriched_products = []

    user_insight_dict = None

    if return_user_insight:
        content = streaming_callback.accumulated_content
        if content and identity_key:
            user_insight_dict = await extract_and_save_user_insight(content, identity_key)
    elif background_tasks:
        background_tasks.add_task(
            finalize_user_insight_after_stream,
            stream_task=stream_task,
            streaming_callback=streaming_callback,
            identity_key=identity_key,
            trace_id=trace_id,
        )

    response_payload = {
        "ai_response": ai_text_response,
        "product_ids": enriched_products,
    }
    if user_insight_dict is not None:
        response_payload["user_insight"] = user_insight_dict

    if REDIS_CACHE_TURN_ON and not skip_cache:
        await redis_cache.set_response(user_id=identity_key, query=query, response_data=response_payload, ttl=300)

    background_tasks.add_task(
        handle_post_chat_async,
        memory=memory,
        identity_key=identity_key,
        human_query=query,
        ai_response=response_payload,
    )
    return {**response_payload, "cached": False}
