"""
Tag Search Agent Controller — Entry point with Langfuse tracing.

Wraps TagSearchGraph (Planner ⇄ Tools → Responder) with:
- Langfuse trace + span context (timing: planner_ms, responder_ms, tool_ms)
- Session-based Redis history (via ai_tag_search Redis helpers)
- Background flush for Langfuse

Every call creates a named Langfuse trace "tag-search-chat" so the full
Planner → Tool → Responder pipeline is visible in the Langfuse dashboard.
"""

import logging
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.langfuse_client import async_flush_langfuse, get_callback_handler

from .tag_search_graph import get_tag_search_agent

logger = logging.getLogger(__name__)


async def tag_search_chat_controller(
    *,
    query: str,
    session_id: str | None = None,
    background_tasks: BackgroundTasks,
    model_name: str | None = None,
    history: list | None = None,  # pre-loaded LangChain messages (optional)
    user_id: str | None = None,
) -> dict:
    """
    Controller cho Tag Search Agent — full Langfuse tracing.

    Flow:
    1. Build Langfuse trace context
    2. Run TagSearchGraph (Planner ⇄ Tools → Responder)
    3. Extract timing breakdown từ pipeline diagnostics
    4. Update Langfuse trace output
    5. Flush Langfuse in background

    Returns:
        dict with: status, response, products, product_ids, pipeline, timing, trace_id
    """
    start_time = time.time()
    run_id = str(uuid.uuid4())
    trace_id = Langfuse.create_trace_id()

    effective_session = session_id or f"ts-{run_id[:8]}"
    effective_user = user_id or effective_session

    logger.info(
        "📥 [TagSearch Controller] session=%s | query=%.80s",
        effective_session[:12], query,
    )

    # ═══ 1. BUILD LANGFUSE TRACE ═══
    tags = ["tag_search", "experiment"]
    langfuse = get_langfuse()
    observation_ctx = None

    if langfuse:
        try:
            observation_ctx = langfuse.start_as_current_observation(
                as_type="span",
                name="tag-search-chat",
                trace_context={"trace_id": trace_id},
            )
        except Exception as e:
            logger.warning("⚠️ Langfuse span init failed: %s", e)

    span = None
    if observation_ctx:
        try:
            span = observation_ctx.__enter__()
            span.update_trace(
                name="tag-search-chat",
                user_id=effective_user,
                session_id=effective_session,
                tags=tags,
                input={"query": query},
                metadata={
                    "model": model_name,
                    "session_id": effective_session,
                    "history_turns": len(history) if history else 0,
                },
            )
        except Exception as e:
            logger.warning("⚠️ Langfuse trace update failed: %s", e)

    # Langfuse callback handler — tracks each LLM call inside the graph
    langfuse_handler = get_callback_handler()

    exec_config = RunnableConfig(
        callbacks=[langfuse_handler] if langfuse_handler else [],
        run_name="tag-search-graph",
        run_id=run_id,
        metadata={
            "langfuse_session_id": effective_session,
            "langfuse_user_id": effective_user,
            "langfuse_tags": tags,
            "trace_id": trace_id,
        },
    )

    # ═══ 2. RUN TAG SEARCH GRAPH ═══
    try:
        agent = get_tag_search_agent(model_name)
        chat_result = await agent.chat(
            user_message=query,
            history=history or [],
            config=exec_config,
        )

        ai_response = chat_result.get("response", "")
        products = chat_result.get("products", [])
        pipeline = chat_result.get("pipeline", [])
        agent_path = chat_result.get("agent_path", "unknown")
        total_elapsed = chat_result.get("elapsed_ms", 0)

        # ─── Extract product_ids (SKU array) ───
        product_ids = [p["sku"] for p in products if p.get("sku")]

        # ─── Timing breakdown from pipeline diagnostics ───
        planner_ms = 0
        responder_ms = 0
        tool_ms = 0
        for step in pipeline:
            step_type = step.get("step", "")
            ms = step.get("elapsed_ms", 0) or 0
            if "planner" in step_type:
                planner_ms += ms
            elif step_type == "responder":
                responder_ms += ms
            # tool_ms không có elapsed_ms in pipeline (tool timing tracked by tool itself)

        wall_ms = round((time.time() - start_time) * 1000)

        # ─── Log timing breakdown ───
        logger.info(
            "✅ [TagSearch Controller] path=%s | products=%d | "
            "planner=%dms | responder=%dms | total=%dms | wall=%dms",
            agent_path, len(products),
            planner_ms, responder_ms, round(total_elapsed), wall_ms,
        )

        # ═══ 3. UPDATE LANGFUSE TRACE OUTPUT ═══
        if span:
            try:
                span.update_trace(
                    output={
                        "ai_response": (ai_response or "")[:500],
                        "agent_path": agent_path,
                        "product_count": len(products),
                        "product_ids": product_ids[:10],  # top 10 for trace
                        "elapsed_ms": round(total_elapsed),
                        "planner_ms": planner_ms,
                        "responder_ms": responder_ms,
                    },
                )
            except Exception:
                pass

        # ═══ 4. RETURN ═══
        return {
            "status": "success",
            "response": ai_response,
            "products": products,
            "product_ids": product_ids,
            "product_count": len(products),
            "pipeline": pipeline,
            "agent_path": agent_path,
            "tool_calls": chat_result.get("tool_calls", []),
            "trace_id": trace_id,
            "session_id": effective_session,
            "timing": {
                "planner_ms": planner_ms,
                "responder_ms": responder_ms,
                "tool_ms": tool_ms,
                "total_ms": round(total_elapsed),
                "wall_ms": wall_ms,
            },
        }

    except Exception as e:
        elapsed = round((time.time() - start_time) * 1000)
        logger.error(
            "❌ [TagSearch Controller] Error after %dms: %s", elapsed, e,
            exc_info=True,
        )

        if span:
            try:
                span.update_trace(
                    output={"error": str(e)[:500], "elapsed_ms": elapsed},
                    level="ERROR",
                )
            except Exception:
                pass

        return {
            "status": "error",
            "error_code": "TAG_SEARCH_ERROR",
            "message": f"Tag Search Error: {str(e)[:200]}",
            "trace_id": trace_id,
        }

    finally:
        if observation_ctx:
            try:
                observation_ctx.__exit__(None, None, None)
            except Exception:
                pass

        # Async flush Langfuse in background — non-blocking
        if background_tasks:
            background_tasks.add_task(async_flush_langfuse)
