
"""
Chatbot API Route
-----------------
`POST /chat` - endpoint chat cho CuCu Assistant (non-streaming).
`POST /chat/stream` - endpoint chat SSE streaming.
Logic xử lý nằm ở `agent.controller`.
"""

import logging
from dataclasses import dataclass
from typing import AsyncGenerator

from fastapi import APIRouter, BackgroundTasks, HTTPException, Request, Body
from fastapi.responses import StreamingResponse
from opentelemetry import trace

# Lazy imports - defer heavy AI modules to first use
from common.message_limit import message_limit_service
from config import DEFAULT_MODEL

logger = logging.getLogger(__name__)
tracer = trace.get_tracer(__name__)
router = APIRouter()

# Cache for lazy-loaded modules
_agent_modules = {}


def _get_agent_controller():
    """Lazy-load agent.controller to defer LangChain/LangGraph import."""
    if "controller" not in _agent_modules:
        from agent.controller import chat_controller, chat_controller_stream
        _agent_modules["controller"] = chat_controller
        _agent_modules["controller_stream"] = chat_controller_stream
    return _agent_modules["controller"], _agent_modules["controller_stream"]


def _get_query_request_model():
    """Lazy-load agent.models.QueryRequest."""
    if "QueryRequest" not in _agent_modules:
        from agent.models import QueryRequest
        _agent_modules["QueryRequest"] = QueryRequest
    return _agent_modules["QueryRequest"]


@dataclass
class Identity:
  """Định danh đơn giản dựa trên middleware (user_id / device_id)."""
  primary_id: str
  history_key: str
  rate_limit_key: str
  is_authenticated: bool


def _get_identity(request: Request) -> Identity:
  """Derive identity info từ request.state (do middleware set)."""
  user_id = getattr(request.state, "user_id", None)
  device_id = getattr(request.state, "device_id", "") or ""
  is_authenticated = bool(getattr(request.state, "is_authenticated", False))

  if is_authenticated and user_id:
    primary_id = str(user_id)
    history_key = primary_id
    rate_limit_key = primary_id
  else:
    primary_id = device_id or "anonymous"
    history_key = device_id or "anonymous"
    rate_limit_key = device_id or "anonymous"

  return Identity(
    primary_id=primary_id,
    history_key=history_key,
    rate_limit_key=rate_limit_key,
    is_authenticated=is_authenticated,
  )


@router.post("/chat", summary="Chat with CuCu Assistant")
async def cucu_chat(request: Request, background_tasks: BackgroundTasks):
  """Endpoint chat không stream - trả về response JSON đầy đủ một lần."""
  # Lazy-load AI modules on first call
  QueryRequest = _get_query_request_model()
  chat_controller, _ = _get_agent_controller()

  try:
    body = await request.json()
    if 'message' in body and 'user_query' not in body:
        body['user_query'] = body.pop('message')
    
    if not body.get('user_query'):
        body['user_query'] = 'xin chào'

    req = QueryRequest(**body)
  except Exception as e:
    logger.error(f"Validation error in chat: {e}")
    raise HTTPException(status_code=422, detail=f"Invalid request body: {str(e)}")

  identity = _get_identity(request)
  user_id = identity.primary_id

  logger.info("📥 [Incoming Chat] User=%s | Query=%s", user_id, req.user_query)

  span = trace.get_current_span()
  span.set_attribute("user.id", user_id)
  span.set_attribute("chat.user_query", req.user_query)

  try:
    result = await chat_controller(
      query=req.user_query,
      user_id=user_id,
      background_tasks=background_tasks,
      model_name=DEFAULT_MODEL,
      images=req.images,
      identity_key=identity.history_key,
    )

    usage_info = await message_limit_service.increment(
      identity_key=identity.rate_limit_key,
      is_authenticated=identity.is_authenticated,
    )

    logger.info("📤 [Outgoing Chat] User=%s", user_id)

    return {
      "status": "success",
      "ai_response": result["ai_response"],
      "product_ids": result.get("product_ids", []),
      "limit_info": {
        "limit": usage_info["limit"],
        "used": usage_info["used"],
        "remaining": usage_info["remaining"],
      },
    }
  except Exception as e:
    logger.error("Error in cucu_chat: %s", e, exc_info=True)
    raise HTTPException(status_code=500, detail=str(e)) from e


@router.post("/chat/stream", summary="Chat with CuCu Assistant (SSE Streaming)")
async def cucu_chat_stream(request: Request):
  """
  Endpoint chat SSE streaming — trả về token-by-token qua Server-Sent Events.
  """
  # Lazy-load AI modules on first call
  QueryRequest = _get_query_request_model()
  _, chat_controller_stream = _get_agent_controller()

  try:
    body = await request.json()
    if 'message' in body and 'user_query' not in body:
        body['user_query'] = body.pop('message')
    
    if not body.get('user_query'):
        body['user_query'] = 'xin chào'

    req = QueryRequest(**body)
  except Exception as e:
    raise HTTPException(status_code=422, detail=str(e))

  identity = _get_identity(request)
  user_id = identity.primary_id

  logger.info("📥 [Incoming Stream] User=%s | Query=%s", user_id, req.user_query)

  async def sse_generator() -> AsyncGenerator[str, None]:
    try:
      async for chunk_json in chat_controller_stream(
        query=req.user_query,
        user_id=user_id,
        model_name=DEFAULT_MODEL,
        images=req.images,
        identity_key=identity.history_key,
      ):
        yield f"data: {chunk_json}\n\n"

      # Increment rate limit after stream completes
      await message_limit_service.increment(
        identity_key=identity.rate_limit_key,
        is_authenticated=identity.is_authenticated,
      )
      logger.info("📤 [Stream Done] User=%s", user_id)
    except Exception as e:
      logger.error("Error in stream: %s", e, exc_info=True)
      yield f'data: {{"error": "{str(e)}"}}\n\n'

  return StreamingResponse(
    sse_generator(),
    media_type="text/event-stream",
    headers={
      "Cache-Control": "no-cache",
      "Connection": "keep-alive",
      "X-Accel-Buffering": "no",
    },
  )
