"""Binary TCP protocol shared by the K209 remote server and Python client."""

from __future__ import annotations

import socket
import struct
from typing import BinaryIO

import numpy as np

COMMAND_IDENTITY = b"I"
COMMAND_LIMITS = b"L"
COMMAND_CONFIGURE = b"C"
COMMAND_ACQUIRE = b"A"

STATUS_OK = b"O"
STATUS_ERROR = b"E"

DEFAULT_REMOTE_HOST = "127.0.0.1"
DEFAULT_REMOTE_PORT = 50209

CONFIG_STRUCT = struct.Struct("!ddIdd")
LIMITS_STRUCT = struct.Struct("!ddddIdd")
U32_STRUCT = struct.Struct("!I")


def recv_exact(stream: socket.socket | BinaryIO, size: int) -> bytes:
    """Read exactly `size` bytes from a socket-like object."""
    chunks: list[bytes] = []
    remaining = int(size)
    while remaining > 0:
        chunk = stream.recv(remaining) if isinstance(stream, socket.socket) else stream.read(remaining)
        if not chunk:
            raise ConnectionError(f"K209 remote connection closed with {remaining} bytes pending")
        chunks.append(chunk)
        remaining -= len(chunk)
    return b"".join(chunks)


def send_all(stream: socket.socket | BinaryIO, payload: bytes) -> None:
    """Write all payload bytes to a socket-like object."""
    if isinstance(stream, socket.socket):
        stream.sendall(payload)
        return
    stream.write(payload)


def send_u32(stream: socket.socket | BinaryIO, value: int) -> None:
    """Send one network-order uint32."""
    send_all(stream, U32_STRUCT.pack(int(value)))


def recv_u32(stream: socket.socket | BinaryIO) -> int:
    """Read one network-order uint32."""
    return int(U32_STRUCT.unpack(recv_exact(stream, U32_STRUCT.size))[0])


def send_error(stream: socket.socket | BinaryIO, message: str) -> None:
    """Send protocol error response."""
    payload = str(message).encode("utf-8", errors="replace")
    send_all(stream, STATUS_ERROR)
    send_u32(stream, len(payload))
    send_all(stream, payload)


def read_status(stream: socket.socket | BinaryIO) -> None:
    """Read response status and raise remote error when needed."""
    status = recv_exact(stream, 1)
    if status == STATUS_OK:
        return
    if status == STATUS_ERROR:
        message = recv_exact(stream, recv_u32(stream)).decode("utf-8", errors="replace")
        raise RuntimeError(f"K209 remote server error: {message}")
    raise RuntimeError(f"K209 remote server returned invalid status byte: {status!r}")


def send_float32_array(stream: socket.socket | BinaryIO, values: np.ndarray) -> None:
    """Send a float32 array as a length-prefixed little-endian payload."""
    payload = np.asarray(values, dtype="<f4").tobytes(order="C")
    send_u32(stream, len(payload))
    send_all(stream, payload)


def recv_float32_array(stream: socket.socket | BinaryIO, expected_values: int) -> np.ndarray:
    """Read a length-prefixed little-endian float32 array."""
    payload_size = recv_u32(stream)
    expected_size = int(expected_values) * np.dtype(np.float32).itemsize
    if payload_size != expected_size:
        raise RuntimeError(f"K209 remote payload has {payload_size} bytes, expected {expected_size}")
    return np.frombuffer(recv_exact(stream, payload_size), dtype="<f4").copy()