"""
Fashion Q&A Agent Controller (Refactored for 2-Agent Architecture)
Uses CANIFAGraph (Classifier -> Stylist) for high-performance deterministic routing.
"""

import asyncio
import logging
import uuid
import time
from typing import List, Dict, Any

from fastapi import BackgroundTasks
from langchain_core.messages import AIMessage, HumanMessage
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 async_flush_langfuse, get_callback_handler
from config import DEFAULT_MODEL, REDIS_CACHE_TURN_ON

from .controller_helpers import (
    check_response_cache,
    save_user_insight_to_redis,
    load_user_insight_from_redis,
)
from .graph import get_graph_manager
from .helper import handle_post_chat_async
from .models import get_config

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,
    device_id: str | None = None,
    skip_cache: bool = False,
) -> dict:
    """
    Controller logic using the new 2-Agent Architecture.
    """

    logger.info("🚀 [ChatController] Starting Turn | Identity: %s", identity_key)

    # 1. Cache Layer
    if REDIS_CACHE_TURN_ON and not skip_cache:
        cached_res = await check_response_cache(
            identity_key, query, background_tasks, model_name, is_authenticated, device_id
        )
        if cached_res:
            return cached_res

    # 2. Setup Config & Memory
    config = get_config()
    config.model_name = model_name
    
    manager = get_graph_manager(config)
    memory = await get_conversation_manager()

    # 3. Load History & Insight
    history_dicts = await memory.get_chat_history(identity_key, limit=10, include_product_ids=False)
    history = [
        HumanMessage(content=m["message"]) if m["is_human"] else AIMessage(content=m["message"]) 
        for m in history_dicts
    ][::-1]

    user_insight = await load_user_insight_from_redis(identity_key)

    # 4. Tracing Setup (Langfuse)
    trace_id = Langfuse.create_trace_id()
    session_id = f"{identity_key}-{str(uuid.uuid4())[:8]}"
    tags = ["stylist_pro", "user:authenticated" if is_authenticated else "user:anonymous"]

    langfuse_handler = get_callback_handler()
    exec_config = {
        "configurable": {
            "user_id": identity_key,
            "device_id": device_id,
            "model_name": model_name,
            "transient_images": images or [],
            "trace_id": trace_id,
        },
        "callbacks": [langfuse_handler] if langfuse_handler else [],
        "metadata": {
            "trace_id": trace_id,
            "session_id": session_id,
            "user_id": identity_key,
        },
        "tags": tags,
        "run_name": "CANIFAGraph",
    }

    # 5. Graph Invocation (Non-Streaming High-Performance)
    # We use .chat() method which includes enrichment and insight management
    try:
        start_time = time.time()
        
        from langfuse import propagate_attributes
        
        # We wrap the call to trace it manually in Langfuse if needed, 
        # but the LangChain callbacks in graph nodes handle most of it.
        with propagate_attributes(session_id=session_id, user_id=identity_key, tags=tags):
            result = await manager.chat(
                user_message=query,
                user_insight=user_insight,
                history=history,
                config=exec_config
            )
        
        elapsed = time.time() - start_time
        logger.info("✅ Graph Finished | Time: %.2fs | Products: %d", elapsed, len(result.get("products", [])))

    except Exception as e:
        logger.error("❌ Graph execution failed: %s", e, exc_info=True)
        raise e

    # 6. Response Preparation
    ai_response = result.get("response") or result.get("ai_response") or ""
    raw_products = result.get("products", [])
    updated_insight = result.get("updated_insight")

    # Trim massive JSON fields (similar_items, outfit_recommendations, description_text) to save network payload
    final_products = []
    for p in raw_products:
        final_products.append({
            "sku": p.get("magento_ref_code") or p.get("sku") or p.get("internal_ref_code"),
            "product_name": p.get("product_name") or p.get("name"),
            "sale_price": p.get("sale_price") or p.get("price"),
            "original_price": p.get("original_price") or p.get("price_vnd"),
            "product_image_url_thumbnail": p.get("product_image_url_thumbnail") or p.get("image"),
            "product_web_url": p.get("product_web_url") or p.get("url"),
            "master_color": p.get("master_color") or p.get("color_name") or p.get("product_color_name"),
            "sku_color": p.get("product_color_code") or p.get("color_code") or p.get("sku_color")
        })

    response_payload = {
        "ai_response": ai_response,
        "product_ids": final_products, # Stripped-down objects for UI
        "user_insight": updated_insight,
        "lead_stage": result.get("lead_stage"),
        "elapsed_ms": result.get("elapsed_ms")
    }

    # 7. Post-Chat Tasks (Async)
    # Save Insight to Redis
    if updated_insight:
        import json
        insight_str = json.dumps(updated_insight, ensure_ascii=False)
        background_tasks.add_task(save_user_insight_to_redis, identity_key, insight_str)

    # Save to Cache
    if REDIS_CACHE_TURN_ON and not skip_cache:
        background_tasks.add_task(
            redis_cache.set_response, 
            user_id=identity_key, 
            query=query, 
            response_data=response_payload, 
            ttl=300
        )

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

    # Flush Langfuse
    background_tasks.add_task(async_flush_langfuse)

    return {**response_payload, "cached": False}
