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

import logging
import time
import uuid

from fastapi import BackgroundTasks
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
from langchain_core.runnables import RunnableConfig

from common.cache import redis_cache
from common.conversation_manager import get_conversation_manager
from common.langfuse_client import get_callback_handler
from common.llm_factory import create_llm
from config import DEFAULT_MODEL, REDIS_CACHE_TURN_ON
from langfuse import propagate_attributes

from .graph import build_graph
from .helper import extract_product_ids, handle_post_chat_async, parse_ai_response
from .models import AgentState, get_config
from .tools.get_tools import get_all_tools

logger = logging.getLogger(__name__)


async def chat_controller(
    query: str,
    user_id: str,
    background_tasks: BackgroundTasks,
    model_name: str = DEFAULT_MODEL,
    images: list[str] | None = None,
    identity_key: str | None = None,
) -> 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
    
    Args:
        identity_key: Key for saving/loading history (identity.history_key)
                      Guest: device_id, User: user_id
    """
    effective_identity_key = identity_key or user_id
    
    logger.info(
        "chat_controller start: model=%s, user_id=%s, identity_key=%s", 
        model_name, user_id, effective_identity_key
    )

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

    # ====================== NORMAL LLM FLOW ======================
    logger.info("chat_controller: proceed with live LLM call")

    config = get_config()
    config.model_name = model_name

    llm = create_llm(model_name=model_name, streaming=False, json_mode=True)
    tools = get_all_tools()
    graph = build_graph(config, llm=llm, tools=tools)

    # Init ConversationManager (Singleton)
    memory = await get_conversation_manager()

    # Load History
    history_dicts = await memory.get_chat_history(effective_identity_key, limit=20)
    history_messages: list[BaseMessage] = [
        HumanMessage(content=m["message"]) if m["is_human"] else AIMessage(content=m["message"])
        for m in history_dicts
    ]

    # Prepare State
    user_query_message: BaseMessage = HumanMessage(content=query)
    messages_with_query: list[BaseMessage] = [*history_messages, user_query_message]

    initial_state: AgentState = {
        "user_query": user_query_message,
        "messages": messages_with_query,
        "history": history_messages,
        "user_id": user_id,
        "images_embedding": [],
        "ai_response": None,
    }
    
    run_uuid = uuid.uuid4()
    run_id_str = str(run_uuid)
    langfuse_handler = get_callback_handler()
    
    exec_config = RunnableConfig(
        configurable={
            "user_id": user_id,
            "transient_images": images or [],
            "run_id": run_id_str,
        },
        run_id=run_uuid,
        metadata={"run_id": run_id_str, "tags": "chatbot,production"},
        callbacks=[langfuse_handler] if langfuse_handler else [],
    )

    # Execute Graph
    start_time = time.time()
    session_id = f"{user_id}-{run_id_str[:8]}"
    
    with propagate_attributes(user_id=user_id, session_id=session_id):
        result = await graph.ainvoke(initial_state, config=exec_config)
    
    duration = time.time() - start_time

    # Parse Response
    all_product_ids = extract_product_ids(result.get("messages", []))
    ai_raw_content = result.get("ai_response").content if result.get("ai_response") else ""
    ai_text_response, final_product_ids = parse_ai_response(ai_raw_content, all_product_ids)

    response_payload = {
        "ai_response": ai_text_response,
        "product_ids": final_product_ids,
    }

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

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

    logger.info("chat_controller finished in %.2fs", duration)
    return {**response_payload, "cached": False}

