src/api/v1/chat.py

"""Chat endpoint with streaming support."""

import uuid
import json
from typing import List, Dict, Any, Optional

from fastapi import APIRouter, Depends, HTTPException
from langchain_core.messages import HumanMessage, AIMessage
from pydantic import BaseModel
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sse_starlette.sse import EventSourceResponse

from src.agents.chat_handler import ChatHandler
from src.config.settings import settings
from src.db.postgres.connection import get_db
from src.db.postgres.models import ChatMessage, MessageSource
from src.db.redis.connection import get_redis
from src.middlewares.logging import get_logger, log_execution

logger = get_logger("chat_api")

router = APIRouter(prefix="/api/v1", tags=["Chat"])

# One shared ChatHandler for the process. It holds no per-request state (user_id
# is passed into handle()), and lazily builds + caches the Orchestrator/Chatbot
# chains — so reusing it keeps the Azure OpenAI clients (and their httpx/TLS pools)
# warm across requests instead of re-handshaking on the first call of every request.
# enable_slow_path is env-gated (ENABLE_SLOW_PATH): when on, structured intents route
# Orchestrator -> Planner -> TaskRunner -> Assembler so the team can test e2e here.
_chat_handler = ChatHandler(
    enable_tracing=True,
    enable_slow_path=settings.enable_slow_path,
    enable_gate=settings.enable_gate,
)

_GREETINGS = frozenset(["hi", "hello", "hey", "halo", "hai", "hei"])
_GOODBYES = frozenset(["bye", "goodbye", "thanks", "thank you", "terima kasih", "sampai jumpa"])


def _fast_intent(message: str) -> Optional[str]:
    """Return a direct response for obvious greetings/farewells, else None."""
    lower = message.lower().strip().rstrip("!.,?")
    if lower in _GREETINGS:
        return "Hello! How can I assist you today?"
    if lower in _GOODBYES:
        return "Goodbye! Have a great day!"
    return None


class ChatRequest(BaseModel):
    user_id: str
    room_id: str
    message: str


async def get_cached_response(redis, cache_key: str) -> Optional[dict]:
    cached = await redis.get(cache_key)
    if cached:
        data = json.loads(cached)
        if isinstance(data, dict) and "response" in data:
            return data
        # legacy: plain string cached before this change
        return {"response": data, "sources": []}
    return None


# 1h TTL per the 2026-06-11 checkpoint decision (Redis = retrieval/query caching
# only, short-lived). Was 24h, which served stale answers after re-ingestion.
_CHAT_CACHE_TTL_SECONDS = 3600

# Only stateless replies are safe to cache. The cache key is (room, user, message)
# with no analysis-state/data version, so caching a state- or data-dependent answer
# (help / problem_statement / check / structured_flow / unstructured_flow) would
# replay a stale answer after the state or data changes — and, since the read check
# runs before the gate, could even bypass the gate when the same message repeats.
# So we cache ONLY the `chat` intent. Caching analysis answers needs proper
# invalidation on data/state change — deferred. The write is gated by the intent the
# handler already emits; the read stays as-is (safe because only `chat` is ever
# stored).
_CACHEABLE_INTENTS = frozenset({"chat"})


def _chat_cache_key(room_id: str, user_id: str, message: str) -> str:
    # user_id is part of the key so one user's cached answer can never be
    # replayed to another (R5); room_id stays first so the room-wide clear
    # endpoint can keep matching on a `chat:{room_id}:*` prefix.
    # LIMITATION (T-G): the key omits conversation history, so a repeated message
    # replays its cached answer even if the conversation has since moved on. Only
    # the stateless `chat` intent is cached, so the blast radius is small — but a
    # history-aware key (hash of last-N turns) would close it. Flagged to Harry.
    return f"{settings.redis_prefix}chat:{room_id}:{user_id}:{message}"


async def cache_response(redis, cache_key: str, response: str, sources: list):
    await redis.setex(
        cache_key,
        _CHAT_CACHE_TTL_SECONDS,
        json.dumps({"response": response, "sources": sources}),
    )


async def load_history(db: AsyncSession, room_id: str, limit: int = 10) -> list:
    """Load recent chat messages for a room as LangChain message objects (oldest-first)."""
    result = await db.execute(
        select(ChatMessage)
        .where(ChatMessage.room_id == room_id)
        .order_by(ChatMessage.created_at.asc())
        .limit(limit)
    )
    rows = result.scalars().all()
    return [
        HumanMessage(content=row.content) if row.role == "user" else AIMessage(content=row.content)
        for row in rows
    ]


async def save_messages(
    db: AsyncSession,
    room_id: str,
    user_content: str,
    assistant_content: str,
    sources: Optional[List[Dict[str, Any]]] = None,
):
    """Persist user and assistant messages, and attach sources to the assistant message."""
    db.add(ChatMessage(id=str(uuid.uuid4()), room_id=room_id, role="user", content=user_content))
    assistant_id = str(uuid.uuid4())
    db.add(ChatMessage(id=assistant_id, room_id=room_id, role="assistant", content=assistant_content))
    for src in (sources or []):
        page = src.get("page_label")
        db.add(MessageSource(
            id=str(uuid.uuid4()),
            message_id=assistant_id,
            document_id=src.get("document_id"),
            filename=src.get("filename"),
            page_label=str(page) if page is not None else None,
        ))
    await db.commit()


@router.delete("/chat/cache")
async def clear_chat_cache(room_id: str, user_id: str, message: str):
    """Delete the Redis cache entry for a specific room + user + message pair."""
    redis = await get_redis()
    cache_key = _chat_cache_key(room_id, user_id, message)
    deleted = await redis.delete(cache_key)
    return {"deleted": deleted > 0, "cache_key": cache_key}


@router.delete("/chat/cache/room/{room_id}")
async def clear_room_cache(room_id: str):
    """Delete all Redis cache entries for a room."""
    redis = await get_redis()
    pattern = f"{settings.redis_prefix}chat:{room_id}:*"
    keys = await redis.keys(pattern)
    if keys:
        await redis.delete(*keys)
    return {"deleted_count": len(keys), "room_id": room_id}


@router.delete("/retrieval/cache/{user_id}")
async def clear_retrieval_cache(user_id: str):
    """Delete all cached retrieval results for a user. Call this after uploading/processing new documents."""
    from src.retrieval.router import retrieval_router
    deleted = await retrieval_router.invalidate_cache(user_id)
    return {"deleted_count": deleted, "user_id": user_id}


@router.post("/chat/stream")
@log_execution(logger)
async def chat_stream(request: ChatRequest, db: AsyncSession = Depends(get_db)):
    """Chat endpoint with streaming response.

    SSE event sequence:
      1. sources  — JSON array of source refs from ChatHandler (table for
                    structured; deduped document_id/page_label for unstructured)
      2. chunk    — text fragments of the answer
      3. done     — signals end of stream
    """
    redis = await get_redis()
    cache_key = _chat_cache_key(request.room_id, request.user_id, request.message)

    # Redis cache hit
    cached = await get_cached_response(redis, cache_key)
    logger.info("cache check", cache_key=cache_key, cache_hit=cached is not None)
    if cached:
        logger.info("Returning cached response")
        cached_text = cached["response"]
        cached_sources = cached["sources"]
        await save_messages(db, request.room_id, request.message, cached_text, sources=cached_sources)

        async def stream_cached():
            yield {"event": "sources", "data": json.dumps(cached_sources)}
            for i in range(0, len(cached_text), 50):
                yield {"event": "chunk", "data": cached_text[i:i + 50]}
            yield {"event": "done", "data": ""}

        return EventSourceResponse(stream_cached())

    try:
        # Fast intent: greetings/farewells bypass LLM entirely
        direct = _fast_intent(request.message)
        if direct:
            await cache_response(redis, cache_key, direct, sources=[])
            await save_messages(db, request.room_id, request.message, direct, sources=[])

            async def stream_direct():
                yield {"event": "sources", "data": json.dumps([])}
                yield {"event": "chunk", "data": direct}
                yield {"event": "done", "data": ""}

            return EventSourceResponse(stream_direct())

        history = await load_history(db, request.room_id, limit=10)
        handler = _chat_handler

        async def stream_response():
            logger.info("stream_response started", room_id=request.room_id, user_id=request.user_id)
            full_response = ""
            sources: List[Dict[str, Any]] = []
            effective_intent: Optional[str] = None
            async for event in handler.handle(
                request.message, request.user_id, history, analysis_id=request.room_id
            ):
                if event["event"] == "intent":
                    # consumed internally (not forwarded); gates caching below.
                    try:
                        effective_intent = json.loads(event["data"]).get("intent")
                    except (TypeError, ValueError, AttributeError):
                        effective_intent = None
                elif event["event"] == "sources":
                    try:
                        sources = json.loads(event["data"]) or []
                    except (TypeError, ValueError):
                        sources = []
                    yield event
                elif event["event"] == "chunk":
                    full_response += event["data"]
                    yield event
                elif event["event"] == "done":
                    # Only cache stateless `chat` replies — caching a state/data-
                    # dependent answer would replay it stale (see _CACHEABLE_INTENTS).
                    if effective_intent in _CACHEABLE_INTENTS:
                        await cache_response(redis, cache_key, full_response, sources=sources)
                    logger.info("saving messages", sources_count=len(sources), sources=sources)
                    try:
                        await save_messages(db, request.room_id, request.message, full_response, sources=sources)
                    except Exception as e:
                        logger.error("save_messages failed", room_id=request.room_id, error=str(e))
                    yield event
                elif event["event"] == "status":
                    # slow-path progress ("Planning…", "Running N steps…"): forward
                    # so the client shows activity and the SSE connection stays alive.
                    yield event
                elif event["event"] == "error":
                    yield event
                    return

        return EventSourceResponse(stream_response())

    except Exception as e:
        logger.error("Chat failed", error=str(e))
        raise HTTPException(status_code=500, detail=f"Chat failed: {str(e)}")