"""
Lead Stage Agent Controller — Entry point with Langfuse tracing.

Wraps LeadStageGraph (Classifier ⇄ Tools → Stylist) with:
- Langfuse trace + span context
- Conversation history loading
- User insight injection
- Background task for history persistence

Every call creates a named Langfuse trace "lead-stage-chat" so the full
Classifier → Tool → Stylist 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
from langfuse import get_client as get_langfuse

from agent.controller_helpers import load_user_insight_from_redis, save_user_insight_to_redis
from agent.helper import handle_post_chat_async
from common.conversation_manager import get_conversation_manager
from common.langfuse_client import async_flush_langfuse, get_callback_handler

from .graph import get_lead_stage_agent

logger = logging.getLogger(__name__)


async def lead_stage_chat_controller(
    *,
    query: str,
    identity_key: str,
    background_tasks: BackgroundTasks,
    model_name: str | None = None,
    conversation_id: str | None = None,
    is_authenticated: bool = False,
    device_id: str | None = None,
) -> dict:
    """
    Controller cho Lead Stage Agent — full Langfuse tracing.

    Flow:
    1. Load user_insight từ Redis
    2. Load chat history từ ConversationManager
    3. Build Langfuse trace context
    4. Run LeadStageGraph (Classifier ⇄ Tools → Stylist)
    5. Save conversation background
    6. Return response + pipeline + lead_stage metadata

    Returns:
        dict with: status, ai_response, products, lead_stage, pipeline, timing, trace_id
    """
    start_time = time.time()
    run_id = str(uuid.uuid4())
    session_id = conversation_id or f"{identity_key}-lead-{run_id[:8]}"

    logger.info(f"📥 [Lead Controller] identity={identity_key} | query={query[:80]}")

    # ═══ 1. LOAD USER INSIGHT ═══
    user_insight = None
    if identity_key:
        user_insight = await load_user_insight_from_redis(str(identity_key))

    # ═══ 2. LOAD CHAT HISTORY ═══
    memory = await get_conversation_manager()
    history_dicts = await memory.get_chat_history(
        str(identity_key),
        limit=10,
        include_product_ids=False,
        conversation_id=conversation_id,
    )
    history_msgs = [
        HumanMessage(content=m["message"]) if m["is_human"] else AIMessage(content=m["message"])
        for m in history_dicts
    ][::-1]

    # ═══ 3. BUILD LANGFUSE TRACE ═══
    trace_id = Langfuse.create_trace_id()
    tags = ["lead_stage", "experiment"]
    if is_authenticated:
        tags.append("user:authenticated")
    else:
        tags.append("user:anonymous")

    langfuse = get_langfuse()
    observation_ctx = None

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

    # Enter observation context
    span = None
    if observation_ctx:
        try:
            span = observation_ctx.__enter__()
            span.update_trace(
                name="lead-stage-chat",
                user_id=str(identity_key),
                session_id=session_id,
                tags=tags,
                input={"query": query, "identity_key": identity_key},
                metadata={
                    "device_id": device_id,
                    "customer_id": identity_key if is_authenticated else None,
                    "model": model_name,
                    "conversation_id": conversation_id,
                    "history_turns": len(history_msgs),
                    "has_user_insight": user_insight is not None,
                },
            )
        except Exception as e:
            logger.warning(f"⚠️ Langfuse trace update failed: {e}")

    # Create CallbackHandler INSIDE observation context for proper nesting
    langfuse_handler = get_callback_handler()

    exec_config = RunnableConfig(
        callbacks=[langfuse_handler] if langfuse_handler else [],
        run_name="lead-stage-graph",
        run_id=run_id,
        metadata={
            "langfuse_session_id": session_id,
            "langfuse_user_id": str(identity_key),
            "langfuse_tags": tags,
            "trace_id": trace_id,
            "conversation_id": session_id,
            "customer_id": identity_key if is_authenticated else None,
            "device_id": device_id,
        },
    )

    # ═══ 4. RUN LEAD STAGE GRAPH ═══
    try:
        agent = get_lead_stage_agent(model_name)
        chat_result = await agent.chat(
            user_message=query,
            user_insight=user_insight,
            history=history_msgs,
            config=exec_config,
        )

        ai_response = chat_result.get("response", "")
        lead_stage = chat_result.get("lead_stage") or {}
        pipeline = chat_result.get("pipeline", [])
        products = chat_result.get("products", [])
        total_elapsed = chat_result.get("elapsed_ms", 0)

        # Timing breakdown from pipeline
        classifier_ms = 0
        stylist_ms = 0
        tool_ms = 0
        for p in pipeline:
            step = p.get("step", "")
            ms = p.get("elapsed_ms", 0)
            if "classifier" in step:
                classifier_ms += ms
            elif step in ("stylist", "responder"):
                stylist_ms += ms
            elif step == "tool_result":
                tool_ms += ms

        # Update Langfuse trace output
        if span:
            try:
                span.update_trace(
                    output={
                        "ai_response": (ai_response or "")[:500],
                        "stage": lead_stage.get("stage_name", "UNKNOWN"),
                        "stage_confidence": lead_stage.get("confidence"),
                        "product_count": len(products),
                        "elapsed_ms": round(total_elapsed),
                        "classifier_ms": classifier_ms,
                        "stylist_ms": stylist_ms,
                    },
                )
            except Exception:
                pass

        logger.info(
            f"✅ [Lead Controller] Stage: {lead_stage.get('stage_name', 'N/A')} | "
            f"Products: {len(products)} | "
            f"Time: {total_elapsed:.0f}ms (C:{classifier_ms}ms S:{stylist_ms}ms)"
        )

        # ═══ 5. SAVE CONVERSATION (background) ═══
        # Chỉ lưu text của Stylist + compact product list (sku + name)
        compact_products = [
            {"sku": p.get("sku_code") or p.get("sku") or p.get("internal_ref_code", ""),
             "name": p.get("name") or p.get("product_name", "")}
            for p in products[:6]
        ]
        response_payload = {
            "ai_response": ai_response,        # Chỉ text của Stylist
            "product_ids": compact_products,   # Compact: SKU + tên (không full object)
            "lead_stage": lead_stage,
        }

        background_tasks.add_task(
            handle_post_chat_async,
            memory=memory,
            identity_key=str(identity_key),
            human_query=query,
            ai_response=response_payload,
            conversation_id=conversation_id,
        )

        # ── Persist Updated Insight về Redis ──
        updated_insight = chat_result.get("updated_insight")
        if updated_insight and identity_key:
            import json as _json
            insight_str = _json.dumps(updated_insight, ensure_ascii=False)
            background_tasks.add_task(save_user_insight_to_redis, str(identity_key), insight_str)
            logger.info(f"💾 [Lead Controller] Queued insight persist | stage={updated_insight.get('STAGE', '?')}")


        # Lưu Postgres (persistent, không bị TTL như Redis)
        from common.lead_flow_postgres import save_lead_turn
        background_tasks.add_task(
            save_lead_turn,
            device_id=str(device_id or identity_key),
            conv_id=str(conversation_id or session_id),
            human_message=query,
            ai_response_text=ai_response,
            products=products,
            lead_stage=lead_stage,
        )

        # ═══ 6. RETURN ═══
        return {
            "status": "success",
            "ai_response": ai_response,
            "products": products,
            "trace_id": trace_id,
            "lead_stage": lead_stage,
            "pipeline": pipeline,
            "timing": {
                "classifier_ms": classifier_ms,
                "stylist_ms": stylist_ms,
                "tool_ms": tool_ms,
                "total_ms": round(total_elapsed),
            },
        }

    except Exception as e:
        elapsed = round((time.time() - start_time) * 1000)
        logger.error(f"❌ [Lead Controller] Error after {elapsed}ms: {e}", exc_info=True)

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

        return {
            "status": "error",
            "error_code": "LEAD_STAGE_ERROR",
            "message": f"Lead Stage Error: {str(e)[:200]}",
            "trace_id": trace_id,
        }

    finally:
        # Close Langfuse observation context
        if observation_ctx:
            try:
                observation_ctx.__exit__(None, None, None)
            except Exception:
                pass

        # Async flush Langfuse in background
        background_tasks.add_task(async_flush_langfuse)
