"""FR client for the Face Recognition server.

Sends a single JPEG frame to the FR WebSocket server and returns the list of
detections.  The server runs SCRFD + ViT + ByteTrack; per-connection tracking
state resets when the connection closes, which is fine for the demo pattern of
one ``IDENTIFY_PERSON`` call per scene.
"""

from __future__ import annotations

import json
import logging
import os
from dataclasses import dataclass, field
from typing import Any, Optional

logger = logging.getLogger(__name__)


@dataclass
class FRConfig:
    """Configuration for the FR client.

    Attributes:
        endpoint: WebSocket URL for the FR server.
        timeout_s: Connection + response timeout in seconds.
    """

    endpoint: str = field(
        default_factory=lambda: os.getenv(
            "FR_ENDPOINT", "ws://kluster.klass.dev:42067/"
        )
    )
    timeout_s: int = 10


class FRClient:
    """Async WebSocket client for the Face Recognition server.

    Usage::

        client = FRClient()
        detections = await client.identify(jpeg_bytes)
        # [{"identity": "alice", "confidence": 0.92, "bbox": [x1, y1, x2, y2]}, ...]

    The FR server expects:
    - **In:**  Raw binary JPEG/PNG bytes (one frame per message, any resolution)
    - **Out:** ``{"timestamp": ..., "detections": [{"identity": str|null,
               "confidence": float, "bbox": [x1, y1, x2, y2]}]}``

    A new WebSocket connection is opened per call.  This is adequate for the
    demo frequency of 1–2 calls per scene and avoids long-lived connection
    management complexity.
    """

    def __init__(self, config: Optional[FRConfig] = None) -> None:
        self.config = config or FRConfig()

    async def identify(self, jpeg_bytes: bytes) -> list[dict[str, Any]]:
        """Send a JPEG frame to the FR server and return detections.

        Args:
            jpeg_bytes: Raw JPEG bytes of the frame to analyse.

        Returns:
            List of detection dicts, each with keys:
            - ``identity``: str name or ``None`` if unrecognised
            - ``confidence``: float in [0, 1]
            - ``bbox``: ``[x1, y1, x2, y2]`` pixel coordinates

            Returns an empty list on any error (connection failure, timeout,
            malformed response).
        """
        if not jpeg_bytes:
            logger.warning("FRClient.identify called with empty bytes — skipping")
            return []

        try:
            import websockets  # lazy import — optional dependency
        except ImportError:
            logger.error(
                "websockets is not installed — cannot call FR server. "
                "Install it with: pip install websockets"
            )
            return []

        try:
            async with websockets.connect(
                self.config.endpoint,
                open_timeout=self.config.timeout_s,
                close_timeout=self.config.timeout_s,
            ) as ws:
                # Send raw binary frame
                await ws.send(jpeg_bytes)

                # Receive JSON response (with timeout)
                import asyncio

                raw = await asyncio.wait_for(ws.recv(), timeout=self.config.timeout_s)

            # Parse response
            data = json.loads(raw)
            detections: list[dict[str, Any]] = data.get("detections", [])
            logger.info(
                "FR identified %d detection(s) from %d-byte frame",
                len(detections),
                len(jpeg_bytes),
            )
            return detections

        except Exception as exc:
            logger.error("FR identify request failed: %s", exc)
            return []