import asyncio
import logging

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from src.backend.api.routes import _frame_to_dict
from src.backend.services.image_saver import ImageSaver
from src.backend.services.phase_timing import timer

try:
    from src.backend.services.detector import Detector
except Exception:
    Detector = None  # onnxruntime is optional; YOLO features stay disabled

logger = logging.getLogger(__name__)

ws_router = APIRouter()

_camera = None
_saver: ImageSaver | None = None
_detector = None  # Detector | None when onnxruntime is installed


def set_camera(camera):
    global _camera
    _camera = camera


def set_image_saver(saver: ImageSaver) -> None:
    global _saver
    _saver = saver


def set_detector(detector) -> None:
    global _detector
    _detector = detector


@ws_router.websocket("/ws")
async def websocket_stream(ws: WebSocket):
    """WebSocket stream of RGB+depth frames.

    Client → server text messages:
      'start'       — begin streaming
      'stop'        — pause streaming
      'close'       — disconnect
      'predict_on'  — enable YOLO inference per frame (adds 'detections' to JSON)
      'predict_off' — disable YOLO inference
      'conf:0.55'   — set YOLO confidence threshold (0.0–1.0)
    """
    await ws.accept()
    logger.info("WebSocket client connected")
    streaming = False
    predicting = False

    try:
        while True:
            if streaming and _camera and _camera.is_open():
                # Waiting for the frame, encoding and saving run in threads: the
                # event loop keeps serving REST and the other sockets meanwhile.
                frame = await asyncio.to_thread(_camera.capture, 500)
                if frame:
                    data = await asyncio.to_thread(_frame_to_dict, frame)
                    if predicting and _detector is not None:
                        data["detections"] = await asyncio.to_thread(_detector.detect, frame.rgb)
                    with timer.phase("ws_send"):
                        await ws.send_json(data)
                    if _saver and _saver.is_recording:
                        await asyncio.to_thread(_saver.save_rec_frame, frame.rgb)
                await asyncio.sleep(0.033)

            try:
                msg = await asyncio.wait_for(ws.receive_text(), timeout=0.01)
                if msg == "start":
                    streaming = True
                elif msg == "stop":
                    streaming = False
                elif msg == "close":
                    break
                elif msg == "predict_on":
                    predicting = _detector is not None
                    if not predicting:
                        logger.warning("predict_on requested but detector not loaded")
                elif msg == "predict_off":
                    predicting = False
                elif msg.startswith("conf:") and _detector is not None:
                    try:
                        _detector.set_conf(float(msg.split(":", 1)[1]))
                    except ValueError:
                        logger.warning(f"Invalid conf value: {msg!r}")
            except asyncio.TimeoutError:
                pass

            if not streaming:
                await asyncio.sleep(0.1)

    except WebSocketDisconnect:
        logger.info("WebSocket client disconnected")