"""
zmq_client
PC 側の ZMQ 通信を担当するモジュール
画像の受信と操舵量の送信を行う
"""

import json

import cv2
import numpy as np
import zmq

from common import config


class PcZmqClient:
    """PC 側の ZMQ 通信クライアント

    画像受信（SUB）と操舵量送信（PUB）の2チャネルを管理する
    """

    def __init__(self) -> None:
        self._context = zmq.Context()
        self._image_socket: zmq.Socket | None = None
        self._control_socket: zmq.Socket | None = None

    def start(self) -> None:
        """通信ソケットを初期化してバインドする"""

        # 画像受信ソケット（SUB，Pi からの画像を受信）
        self._image_socket = self._context.socket(zmq.SUB)
        self._image_socket.setsockopt(zmq.CONFLATE, 1)
        self._image_socket.setsockopt_string(zmq.SUBSCRIBE, "")
        self._image_socket.bind(config.image_bind_address())

        # 操舵量送信ソケット（PUB，Pi へ操舵量を送信）
        self._control_socket = self._context.socket(zmq.PUB)
        self._control_socket.bind(config.control_bind_address())

    def receive_image(self) -> np.ndarray | None:
        """画像を非ブロッキングで受信する

        Returns:
            受信した画像の NumPy 配列，受信データがない場合は None
        """
        if self._image_socket is None:
            return None
        try:
            data = self._image_socket.recv(zmq.NOBLOCK)
            frame = cv2.imdecode(
                np.frombuffer(data, dtype=np.uint8),
                cv2.IMREAD_COLOR,
            )
            return frame
        except zmq.Again:
            return None

    def send_control(
        self, throttle: float, steer: float,
    ) -> None:
        """操舵量を送信する

        Args:
            throttle: 前後方向の出力 (-1.0 ~ +1.0)
            steer: 左右方向の出力 (-1.0 ~ +1.0)
        """
        if self._control_socket is None:
            return
        payload = json.dumps({
            "throttle": throttle,
            "steer": steer,
        }).encode("utf-8")
        self._control_socket.send(payload, zmq.NOBLOCK)

    def stop(self) -> None:
        """通信ソケットを閉じる"""
        if self._image_socket is not None:
            self._image_socket.close()
            self._image_socket = None
        if self._control_socket is not None:
            self._control_socket.close()
            self._control_socket = None
        self._context.term()