"""
CuCu Note Agent Controller
Langfuse will auto-trace via LangChain integration (no code changes needed).
"""

import json
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 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

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")

    # Graph is a singleton — built once and reused across requests.
    # LLM instances are cached by LLMFactory, tools are static.
    graph = build_graph()

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

    # Load History
    history_dicts = await memory.get_chat_history(effective_identity_key, limit=5)
    history_messages: list[BaseMessage] = []
    for m in history_dicts:
        if m["is_human"]:
            history_messages.append(HumanMessage(content=m["message"]))
        else:
            # AI responses may be saved as JSON — extract readable text
            ai_content = m["message"]
            try:
                parsed = json.loads(ai_content)
                if isinstance(parsed, dict) and "ai_response" in parsed:
                    ai_content = parsed["ai_response"]
            except (json.JSONDecodeError, TypeError):
                pass
            history_messages.append(AIMessage(content=ai_content))

    # Prepare State
    # - history: previous conversation messages (for context)
    # - messages: starts with current query only (tool calls will be appended by LangGraph)
    # - user_query: current user message
    user_query_message: BaseMessage = HumanMessage(content=query)

    initial_state: AgentState = {
        "user_query": user_query_message,
        "messages": [user_query_message],
        "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 ""
    logger.info("RAW LLM output (%d chars): %s", len(ai_raw_content), ai_raw_content[:500])
    ai_text_response, final_product_ids = parse_ai_response(ai_raw_content, all_product_ids)
    logger.info("PARSED ai_response (%d chars): %s", len(ai_text_response), ai_text_response[:300])

    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}


async def chat_controller_stream(
    query: str,
    user_id: str,
    model_name: str = DEFAULT_MODEL,
    images: list[str] | None = None,
    identity_key: str | None = None,
):
    """
    Streaming controller — yields SSE events with token chunks.
    
    Each yield is a JSON string:
      {"token": "partial text"}       — during streaming
      {"done": true, "ai_response": "full text"}  — final event
    
    History is saved after stream completes.
    """
    effective_identity_key = identity_key or user_id
    
    logger.info(
        "chat_controller_stream start: model=%s, user_id=%s",
        model_name, user_id,
    )

    # ====================== 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 (stream) for identity_key={effective_identity_key}")
            # Stream cached response as one chunk
            ai_text = cached_response.get("ai_response", "")
            yield json.dumps({"token": ai_text}, ensure_ascii=False)
            yield json.dumps({"done": True, "ai_response": ai_text}, ensure_ascii=False)
            # Save history
            memory = await get_conversation_manager()
            await handle_post_chat_async(
                memory=memory,
                identity_key=effective_identity_key,
                human_query=query,
                ai_response=cached_response,
            )
            return

    # ====================== STREAM LLM FLOW ======================
    graph = build_graph()
    memory = await get_conversation_manager()

    # Load History
    history_dicts = await memory.get_chat_history(effective_identity_key, limit=5)
    history_messages: list[BaseMessage] = []
    for m in history_dicts:
        if m["is_human"]:
            history_messages.append(HumanMessage(content=m["message"]))
        else:
            ai_content = m["message"]
            try:
                parsed = json.loads(ai_content)
                if isinstance(parsed, dict) and "ai_response" in parsed:
                    ai_content = parsed["ai_response"]
            except (json.JSONDecodeError, TypeError):
                pass
            history_messages.append(AIMessage(content=ai_content))

    user_query_message: BaseMessage = HumanMessage(content=query)

    initial_state: AgentState = {
        "user_query": user_query_message,
        "messages": [user_query_message],
        "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,stream"},
        callbacks=[langfuse_handler] if langfuse_handler else [],
    )

    # Stream using astream_events
    start_time = time.time()
    full_response = ""
    is_final_response = False

    session_id = f"{user_id}-{run_id_str[:8]}"
    
    with propagate_attributes(user_id=user_id, session_id=session_id):
        async for event in graph.astream_events(
            initial_state, config=exec_config, version="v2"
        ):
            kind = event.get("event", "")
            
            # Only stream tokens from the chat model (not tool calls)
            if kind == "on_chat_model_stream":
                chunk = event.get("data", {}).get("chunk")
                if chunk and hasattr(chunk, "content") and chunk.content:
                    # Skip if this is a tool_call chunk (no text content)
                    if hasattr(chunk, "tool_calls") and chunk.tool_calls:
                        continue
                    # Only stream the FINAL response (after tool execution)
                    # We detect this by tracking: if tools were called, 
                    # the final agent response comes after tool results
                    token = chunk.content
                    full_response += token
                    is_final_response = True
                    yield json.dumps({"token": token}, ensure_ascii=False)

            # When the agent finishes and we got tokens, prepare to end
            elif kind == "on_chain_end" and event.get("name") == "agent":
                if is_final_response:
                    # Reset for potential next agent iteration
                    pass

    duration = time.time() - start_time
    logger.info("chat_controller_stream finished in %.2fs (%d chars)", duration, len(full_response))

    # Parse and yield final event
    ai_text_response, _ = parse_ai_response(full_response, [])
    yield json.dumps({"done": True, "ai_response": ai_text_response}, ensure_ascii=False)

    # Build response payload for caching and history
    response_payload = {
        "ai_response": ai_text_response,
        "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,
        )

    # Save to history
    await handle_post_chat_async(
        memory=memory,
        identity_key=effective_identity_key,
        human_query=query,
        ai_response=response_payload,
    )
