"""Shared helpers for CV subscribers."""

from __future__ import annotations

import json

from pathlib import Path

from typing import Iterator

import cv2

import numpy as np

import yaml

import zmq

ROOT = Path(__file__).resolve().parents[1]

CONFIG_PATH = ROOT / "configs" / "runtime.yaml"

def load_config() -> dict:

    with open(CONFIG_PATH, "r", encoding="utf-8") as f:

        return yaml.safe_load(f)

def make_frame_subscriber(endpoint: str, ctx: zmq.Context | None = None) -> zmq.Socket:

    ctx = ctx or zmq.Context.instance()

    sock = ctx.socket(zmq.SUB)

    sock.setsockopt(zmq.RCVHWM, 1)

    sock.setsockopt(zmq.SUBSCRIBE, b"")

    sock.connect(endpoint)

    return sock

def make_pub(endpoint: str, ctx: zmq.Context | None = None) -> zmq.Socket:

    ctx = ctx or zmq.Context.instance()

    sock = ctx.socket(zmq.PUB)

    sock.setsockopt(zmq.SNDHWM, 8)

    sock.bind(endpoint)

    return sock

def iter_latest_frames(sock: zmq.Socket, timeout_ms: int = 2000) -> Iterator[tuple[dict, np.ndarray]]:

    """Yield (meta_dict, decoded_bgr_frame) for each arriving message, dropping stale ones.

    Blocks up to timeout_ms for the next message; raises TimeoutError if idle too long.

    """

    poller = zmq.Poller()

    poller.register(sock, zmq.POLLIN)

    while True:

        events = dict(poller.poll(timeout_ms))

        if sock not in events:

            raise TimeoutError(f"No frame within {timeout_ms}ms")

        latest: list[bytes] | None = None

        while True:

            try:

                latest = sock.recv_multipart(flags=zmq.NOBLOCK)

            except zmq.Again:

                break

        if latest is None or len(latest) != 2:

            continue

        meta = json.loads(latest[0].decode("utf-8"))

        arr = np.frombuffer(latest[1], dtype=np.uint8)

        frame = cv2.imdecode(arr, cv2.IMREAD_COLOR)

        if frame is None:

            continue

        yield meta, frame

def send_event(sock: zmq.Socket, event: dict) -> None:

    sock.send(json.dumps(event).encode("utf-8"))

def event_subscriber(endpoint: str, ctx: zmq.Context | None = None) -> zmq.Socket:

    ctx = ctx or zmq.Context.instance()

    sock = ctx.socket(zmq.SUB)

    sock.setsockopt(zmq.RCVHWM, 64)

    sock.setsockopt(zmq.SUBSCRIBE, b"")

    sock.connect(endpoint)

    return sock