"""WebSocket endpoint streaming G1 motor temperatures to the 3D dashboard (N1).

Polls the arm controller's throttled rt/lowstate snapshot (arm.get_motor_temps
/ arm.get_current_q — NO second DDS subscriber, no second ChannelFactoryInitialize)
and pushes a Marcus-compatible 'motor_update' payload to each connected client.

Front-end: dashboard/static/temp3d/index.html (ported three.js view), which
opens this socket via a tiny shim in place of socket.io.
"""

from __future__ import annotations

import asyncio
import threading
import time

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from Project.Sanad.core.logger import get_logger
from Project.Sanad.dashboard.temp_motor_map import build_payload

log = get_logger("motor_temps_ws")

router = APIRouter()

MAX_WATCHERS = 20
PUSH_HZ = 8.0          # ~8 fps is plenty for a temperature heatmap

_count = 0
_count_lock = threading.Lock()


def _get_arm():
    """Lazy import — avoids a circular import on dashboard load."""
    try:
        from Project.Sanad.main import arm  # type: ignore
        return arm
    except Exception:
        return None


@router.websocket("/ws/motor-temps")
async def motor_temps_ws(ws: WebSocket):
    await ws.accept()

    global _count
    with _count_lock:
        if _count >= MAX_WATCHERS:
            await ws.close(code=1013, reason="Too many temperature watchers")
            return
        _count += 1

    period = 1.0 / PUSH_HZ
    try:
        while True:
            arm = _get_arm()
            temps: list = []
            positions: list = []
            if arm is not None:
                try:
                    temps = arm.get_motor_temps()
                except Exception:
                    temps = []
                try:
                    positions = arm.get_current_q()
                except Exception:
                    positions = []
            payload = build_payload(temps, positions, time.time())
            await ws.send_json(payload)
            await asyncio.sleep(period)
    except WebSocketDisconnect:
        pass
    except Exception:
        # Any other error (client gone mid-send, serialise issue) closes cleanly.
        try:
            await ws.close()
        except Exception:
            pass
    finally:
        with _count_lock:
            _count -= 1
