src/api/v2/chat.py

"""Chat endpoint — v2 pilot (pr/5 Phase 2).

`POST /api/v2/chat/stream` is the v2 of the only FE→Python call. It is identical to
`POST /api/v1/chat/stream` except:
  - the request carries an explicit **`analysis_id`** (replacing v1's `room_id`). The
    two are the same session id today (`analysis_id == room_id`), so the warm,
    process-shared `ChatHandler` and the v1 cache/history helpers are reused unchanged.
  - the `done` event carries the assistant **`message_id`**. It is always minted
    Python-side and is **never accepted from the caller** (server-authoritative — it keys
    the future `/observability` lookup; open-Q #1 resolved). The FE reads it off `done`.

Only chat moves to v2; the tools group + observability stay on `/api/v1` (contract:
API_ENDPOINTS_RESTRUCTURE.md §1).

Persistence (DEV_PLAN #25 — done). Python is **read-only** on the Go-owned
`analyses_messages` table: it *reads* turn history (so multi-turn context works) but no
longer writes the user/AI turns — Go is the sole writer. This removes the double-write
that appeared when both Go and Python's stream persisted the same turn. Aligns chat with
`/tools/help`, which was already generative-only.
"""

import json
import uuid
from typing import Any

from fastapi import APIRouter, Depends, HTTPException, Request
from pydantic import BaseModel
from sqlalchemy.ext.asyncio import AsyncSession
from sse_starlette.sse import EventSourceResponse

# Reuse the v1 chat machinery verbatim (warm ChatHandler + cache/history helpers) so
# v2 stays a thin field-rename over the same logic. Importing the module-private helpers
# is the established pattern here (handlers/help.py imports `_chat_handler` the same way).
from src.api.v1.chat import (
    _CACHEABLE_INTENTS,
    _chat_cache_key,
    _chat_handler,
    _fast_intent,
    cache_response,
    get_cached_response,
    load_history,
)
from src.db.postgres.connection import get_db
from src.db.redis.connection import get_redis
from src.middlewares.logging import get_logger, log_execution
from src.middlewares.rate_limit import limiter
from src.traceability import PostgresTraceabilityStore, TraceabilityScratchpad

logger = get_logger("chat_api_v2")

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

# Module-level store (mirrors the warm, process-shared `_chat_handler`): the greeting
# fast-path and cache-replay branches never enter `ChatHandler.handle`, so they write
# their (empty `chat`) traceability row directly here. KM-691.
_traceability_store = PostgresTraceabilityStore()


async def _save_empty_chat_trace(analysis_id: str, user_id: str, message_id: str) -> None:
    """Persist an empty `chat` traceability row for a turn that bypassed the handler
    (greeting / cache replay), so the FE's GET on `done` returns a payload, not a 404."""
    try:
        pad = TraceabilityScratchpad()
        pad.message_id = message_id
        pad.set_intent("chat")
        await _traceability_store.save(pad.build(analysis_id, user_id, message_id))
    except Exception as e:  # noqa: BLE001 — never break the reply on a trace slip
        logger.warning("traceability direct save failed", message_id=message_id, error=str(e))


def _mint_message_id() -> str:
    """Mint the assistant turn id. Server-authoritative — never accepted from the caller
    (it keys the GET /api/v1/traceability lookup). Returned on `done`; open-Q #1 resolved.

    A canonical UUID string, matching Go's `analyses_messages.id` shape, so the value stays
    format-compatible if we later swap to the real message-row id (still Python-minted now)."""
    return str(uuid.uuid4())


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


@router.post("/chat/stream")
# Rate limit per client IP. `slowapi` needs a Starlette `Request` param named
# `request`, so the JSON body moves to `body`. NOTE: if the FE reaches Python through
# the Go proxy, `get_remote_address` sees Go's IP (one bucket for everyone) — size the
# limit accordingly, or switch to a user-scoped key once identity is forwarded.
@limiter.limit("30/minute")
@log_execution(logger)
async def chat_stream(
    request: Request, body: ChatRequest, db: AsyncSession = Depends(get_db)
):
    """Chat endpoint with streaming response (v2 — keyed on `analysis_id`).

    SSE event sequence:
      1. sources  — always `[]` (KM-691): sources moved to GET /api/v1/traceability;
                    the stream stays text-only. Event kept for backward-compat.
      2. status   — slow-path progress pings (optional)
      3. chunk    — text fragments of the answer
      4. done     — {"message_id": "..."} for the traceability lookup
    """
    analysis_id = body.analysis_id
    message_id = _mint_message_id()
    redis = await get_redis()
    cache_key = _chat_cache_key(analysis_id, body.user_id, body.message)

    # v2 `done` always carries the turn id (v1 sent an empty `done`).
    done_event = {"event": "done", "data": json.dumps({"message_id": message_id})}

    # Redis cache hit (stateless `chat` intent only).
    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"]

        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 done_event

        # Write the row BEFORE the stream so the FE's GET on `done` can't race a 404.
        await _save_empty_chat_trace(analysis_id, body.user_id, message_id)
        return EventSourceResponse(stream_cached())

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

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

            # Write the row BEFORE the stream so the FE's GET on `done` can't race a 404.
            await _save_empty_chat_trace(analysis_id, body.user_id, message_id)
            return EventSourceResponse(stream_direct())

        history = await load_history(db, analysis_id, limit=10)
        handler = _chat_handler

        async def stream_response():
            logger.info("stream_response started", analysis_id=analysis_id, user_id=body.user_id)
            full_response = ""
            sources: list[dict[str, Any]] = []
            effective_intent: str | None = None
            async for event in handler.handle(
                body.message, body.user_id, history,
                analysis_id=analysis_id, message_id=message_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 (see _CACHEABLE_INTENTS).
                    if effective_intent in _CACHEABLE_INTENTS:
                        await cache_response(redis, cache_key, full_response, sources=sources)
                    # Persistence is Go's job now (DEV_PLAN #25): Python reads history but
                    # no longer writes turns, so Go stays the sole writer of analyses_messages.
                    yield done_event
                elif event["event"] == "status":
                    # slow-path progress: forward so the client shows activity.
                    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)}") from e