"""WebSocket endpoint for streaming chat responses with user confirmation."""

from __future__ import annotations

import asyncio
import json
import logging

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from orchestrator.orchestrator import Orchestrator

logger = logging.getLogger(__name__)

router = APIRouter(tags=["websocket"])


@router.websocket("/ws/chat")
async def websocket_chat(websocket: WebSocket):
    """WebSocket endpoint for streaming chat responses.

    Client sends:
      {"question": "...", "conversation_id": "..." (optional)}
      {"type": "confirm"}   — approve pending SQL/API action
      {"type": "reject"}    — reject pending SQL/API action
      {"type": "ping"}      — heartbeat

    Server streams:
      {"type": "stream_start|stream_chunk|status|sql|results|stream_end|error", ...}
      {"type": "confirm_api", "description": str, "metadata_preview": str}
      {"type": "confirm_sql", "query": str, "tables": [...]}
    """
    await websocket.accept()
    orchestrator: Orchestrator | None = None

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

            msg_type = message.get("type", "")

            # Handle ping/pong heartbeat
            if msg_type == "ping":
                await websocket.send_json({"type": "pong"})
                continue

            # Handle confirmation responses
            if msg_type == "confirm" and orchestrator:
                orchestrator.confirm()
                continue

            if msg_type == "reject" and orchestrator:
                orchestrator.reject()
                continue

            # Handle new questions
            question = message.get("question", "").strip()
            if not question:
                await websocket.send_json({"type": "error", "message": "Empty question"})
                continue

            conversation_id = message.get("conversation_id")

            # Create orchestrator for this question
            settings = websocket.app.state.settings
            db = websocket.app.state.db
            orchestrator = Orchestrator(settings, db)

            # Run the pipeline in a task so we can receive confirmation messages
            async def _run_pipeline():
                try:
                    async for event in orchestrator.process_question(question, conversation_id):
                        await websocket.send_json(event)
                except Exception as e:
                    logger.error(f"Orchestrator error: {e}", exc_info=True)
                    await websocket.send_json({
                        "type": "error",
                        "message": "An error occurred while processing your question.",
                    })

            # Start pipeline as background task
            pipeline_task = asyncio.create_task(_run_pipeline())

            # While pipeline runs, listen for confirmation messages
            while not pipeline_task.done():
                try:
                    # Wait for client message with short timeout
                    reply_data = await asyncio.wait_for(
                        websocket.receive_text(), timeout=0.5
                    )
                    try:
                        reply = json.loads(reply_data)
                        reply_type = reply.get("type", "")

                        if reply_type == "confirm":
                            orchestrator.confirm()
                        elif reply_type == "reject":
                            orchestrator.reject()
                        elif reply_type == "ping":
                            await websocket.send_json({"type": "pong"})
                    except json.JSONDecodeError:
                        pass
                except asyncio.TimeoutError:
                    # No message received, check if pipeline finished
                    continue
                except WebSocketDisconnect:
                    pipeline_task.cancel()
                    raise

            # Wait for pipeline to complete
            await pipeline_task

    except WebSocketDisconnect:
        logger.info("WebSocket client disconnected")
    except Exception as e:
        logger.error(f"WebSocket error: {e}", exc_info=True)
