#!/usr/bin/env python3
"""Shared wire-format helpers for pipeview test tools.

Keep these functions in sync with:
  - crates/pipeview-core/src/frame/cobs.rs        (cobs_encode)
  - crates/pipeview-core/src/protocol/mixed.rs    (XP plot packet header)
"""

import math
import struct

PLOT_ESCAPE = 0x1E
PLOT_MARKER = ord("P")

DISCONNECT_WINERRORS = {
    10053,  # Software caused connection abort
    10054,  # Connection reset by peer
    10058,  # Socket shutdown race on Windows
}

# XP plot packet header field values (see protocol/mixed.rs)
PLOT_PACKET_MAGIC = b"XP"
PLOT_PACKET_VERSION = 1
PLOT_FORMAT_IDS = {
    "interleaved": 0,
    "block": 1,
    "xy": 2,
}
SAMPLE_TYPE_F32 = 8
ENDIAN_LITTLE = 0


def cobs_encode(payload: bytes) -> bytes:
    """Consistent Overhead Byte Stuffing encoder (matches Rust cobs_encode)."""
    if not payload:
        return b"\x01"

    out = bytearray([0])
    code_index = 0
    code = 1
    for byte in payload:
        if byte == 0:
            out[code_index] = code
            code_index = len(out)
            out.append(0)
            code = 1
        else:
            out.append(byte)
            code += 1
            if code == 0xFF:
                out[code_index] = code
                code_index = len(out)
                out.append(0)
                code = 1
    out[code_index] = code
    return bytes(out)


def build_plot_payload(
    sample_index: int,
    channels: int,
    plot_format: str,
    samples_per_channel: int,
    amplitude: float,
    frequency_hz: float,
    sample_rate_hz: float,
) -> bytes:
    """Build a raw little-endian f32 plot payload without any framing."""
    if channels <= 0:
        raise ValueError("channels must be >= 1")
    if plot_format not in PLOT_FORMAT_IDS:
        raise ValueError(f"unsupported plot format: {plot_format}")
    if plot_format == "xy" and channels != 2:
        raise ValueError("xy format requires exactly 2 channels")

    values = []
    if plot_format == "xy":
        for offset in range(samples_per_channel):
            t = (sample_index + offset) / sample_rate_hz
            values.extend(
                [
                    amplitude * math.cos(2 * math.pi * frequency_hz * t),
                    amplitude * math.sin(2 * math.pi * frequency_hz * t),
                ]
            )
    elif plot_format == "block":
        for ch in range(channels):
            phase = 2 * math.pi * ch / channels
            for offset in range(samples_per_channel):
                t = (sample_index + offset) / sample_rate_hz
                values.append(
                    amplitude * math.sin(2 * math.pi * frequency_hz * t + phase)
                )
    else:  # interleaved
        for offset in range(samples_per_channel):
            t = (sample_index + offset) / sample_rate_hz
            for ch in range(channels):
                phase = 2 * math.pi * ch / channels
                values.append(
                    amplitude * math.sin(2 * math.pi * frequency_hz * t + phase)
                )

    return struct.pack(f"<{len(values)}f", *values)


def build_plot_packet(
    payload: bytes,
    channels: int,
    samples_per_channel: int,
    plot_format: str,
) -> bytes:
    """Build an XP plot packet (the payload inside a MixedTextPlot COBS frame)."""
    if not 1 <= channels <= 255:
        raise ValueError("channels must be in 1..255")
    if not 0 <= samples_per_channel <= 0xFFFF:
        raise ValueError("samples_per_channel must fit in u16")
    if plot_format not in PLOT_FORMAT_IDS:
        raise ValueError(f"unsupported plot format: {plot_format}")

    header = bytearray()
    header.extend(PLOT_PACKET_MAGIC)
    header.append(PLOT_PACKET_VERSION)
    header.append(PLOT_FORMAT_IDS[plot_format])
    header.append(SAMPLE_TYPE_F32)
    header.append(ENDIAN_LITTLE)
    header.append(channels)
    header.extend(struct.pack("<H", samples_per_channel))
    header.extend(struct.pack("<I", len(payload)))
    header.extend(payload)
    return bytes(header)


def build_mixed_plot_frame(
    payload: bytes,
    channels: int,
    samples_per_channel: int,
    plot_format: str,
) -> bytes:
    """Build a complete MixedTextPlot plot frame (escape + marker + COBS packet + 0x00)."""
    packet = build_plot_packet(
        payload=payload,
        channels=channels,
        samples_per_channel=samples_per_channel,
        plot_format=plot_format,
    )
    return bytes([PLOT_ESCAPE, PLOT_MARKER]) + cobs_encode(packet) + b"\x00"


def is_disconnect_error(err: OSError) -> bool:
    """Return True when the socket error is an expected client-disconnect error."""
    return isinstance(
        err,
        (
            BrokenPipeError,
            ConnectionAbortedError,
            ConnectionResetError,
        ),
    ) or getattr(err, "winerror", None) in DISCONNECT_WINERRORS