"""WebSocket handler for streaming chat."""

from __future__ import annotations

import json
import logging

from fastapi import WebSocket, WebSocketDisconnect

log = logging.getLogger(__name__)


async def ws_chat_handler(websocket: WebSocket, conversation_id: str) -> None:
    """Handle a WebSocket chat connection.

    Protocol:
    - Client sends: {"message": "user text"}
    - Server sends: {"type": "chunk", "delta": "text", "done": false}
    - Server sends: {"type": "chunk", "delta": "", "done": true} on completion
    - Server sends: {"type": "title", "title": "..."} after auto-titling
    """
    await websocket.accept()
    agent = websocket.app.state.agent
    memory = websocket.app.state.memory

    # Conversation row is created lazily on first add_message — opening a WS
    # without sending anything must not leave an empty row behind.

    log.info("WebSocket connected: conversation=%s", conversation_id)

    try:
        while True:
            raw = await websocket.receive_text()
            try:
                data = json.loads(raw)
            except json.JSONDecodeError:
                await websocket.send_json({"type": "error", "message": "Invalid JSON"})
                continue

            # Keepalive ping from client
            if data.get("type") == "ping":
                await websocket.send_json({"type": "pong"})
                continue

            # Regenerate path — the client just hit Try-again on the last
            # assistant reply, OR it just edited a user message via PATCH.
            # The REST route already truncated everything after the target
            # user message; we only need to produce a fresh reply against
            # the truncated history without inserting a new user row.
            #
            # Resolution rules:
            #   1. If the client passes ``message`` (the replay text from
            #      the REST response), use it. Cheaper — no extra DB read.
            #   2. Otherwise look up the last user message ourselves so
            #      raw ``{regenerate: true}`` works as a one-shot signal.
            #   3. If no user message exists at all, surface an error so
            #      the UI can recover (it's a bug, not a normal path).
            is_regenerate = bool(data.get("regenerate"))
            if is_regenerate:
                user_message = data.get("message", "").strip()
                if not user_message:
                    last = await memory.get_last_user_message(conversation_id)
                    if not last:
                        await websocket.send_json({
                            "type": "error",
                            "message": "No user message to regenerate from",
                        })
                        continue
                    user_message = last["content"]
            else:
                user_message = data.get("message", "").strip()
                if not user_message:
                    await websocket.send_json({"type": "error", "message": "Empty message"})
                    continue

            try:
                _provenance = {
                    "kind": "chat_stream_regenerate" if is_regenerate else "chat_stream",
                    "source_channel": "websocket",
                }
                async for chunk in agent.chat_stream(
                    conversation_id, user_message,
                    provenance=_provenance,
                    skip_user_add=is_regenerate,
                ):
                    # Send tool status events for visibility
                    if hasattr(chunk, "tool_name") and chunk.tool_name:
                        await websocket.send_json({
                            "type": "tool_status",
                            "tool": chunk.tool_name,
                            "status": chunk.tool_status or "running",
                        })
                    await websocket.send_json({
                        "type": "chunk",
                        "delta": chunk.delta,
                        "done": chunk.done,
                    })
                # After streaming completes, auto-title in the background
                # so it never blocks the handler from processing the next
                # user message (was causing "infinite thinking" on the
                # second message in a conversation).
                import asyncio

                async def _auto_title() -> None:
                    try:
                        new_title = await memory.auto_title(
                            conversation_id, llm=agent._llm,
                        )
                        if new_title:
                            await websocket.send_json({
                                "type": "title",
                                "title": new_title,
                            })
                    except Exception:
                        log.debug("Auto-title failed", exc_info=True)

                asyncio.create_task(_auto_title())
            except Exception as exc:
                from agntspace.infra.cost_tracker import BudgetExceededError
                from agntspace.llm.base import LLMError

                if isinstance(exc, BudgetExceededError):
                    log.warning("Budget exceeded during chat: %s", exc)
                    detail = f"⚠ {exc}"
                elif isinstance(exc, LLMError):
                    log.error("LLM error during chat: %s", exc.user_message)
                    detail = exc.user_message
                else:
                    log.exception("Error during streaming response")
                    detail = str(exc)
                    if len(detail) > 300:
                        detail = detail[:300] + "…"
                await websocket.send_json({
                    "type": "error",
                    "message": detail or "Internal error during response generation",
                })
    except WebSocketDisconnect:
        log.info("WebSocket disconnected: conversation=%s", conversation_id)
