"""State broadcast over UDP: the compact pose frame from the sync design.

Frame layout (little-endian, 45 bytes) — source of truth for docs and the
prototype bandwidth estimate:

    magic     2s   b"SH"
    version   B
    drone_id  8s   zero-padded ascii (PoC; production may switch to uint16 registry)
    ts_ns     q    epoch nanoseconds
    pos_mm    3i   position in the mission frame, millimeters (wire quantization only)
    att_cdeg  3h   roll/pitch/yaw, centi-degrees
    vel_cms   3h   velocity, cm/s
    frame_ref B
    flags     B
"""

from __future__ import annotations

import socket
import struct
from typing import Any

from .flight import Pose

FRAME = struct.Struct("<2sB8sq3i3h3hBB")
MAGIC = b"SH"
VERSION = 1


def encode_state(drone_id: str, pose: Pose, flags: int = 0) -> bytes:
    return FRAME.pack(
        MAGIC,
        VERSION,
        drone_id.encode()[:8].ljust(8, b"\0"),
        pose.ts_ns,
        int(pose.x * 1000), int(pose.y * 1000), int(pose.z * 1000),
        int(pose.roll * 100), int(pose.pitch * 100), int(pose.yaw * 100),
        int(pose.vx * 100), int(pose.vy * 100), int(pose.vz * 100),
        0,
        flags,
    )


def decode_state(payload: bytes) -> dict[str, Any] | None:
    if len(payload) != FRAME.size:
        return None
    magic, version, drone_id, ts_ns, px, py, pz, r, p, y, vx, vy, vz, frame_ref, flags = FRAME.unpack(payload)
    if magic != MAGIC or version != VERSION:
        return None
    return {
        "ts_ns": ts_ns,
        "peer_id": drone_id.rstrip(b"\0").decode(),
        "pos_x": px / 1000.0, "pos_y": py / 1000.0, "pos_z": pz / 1000.0,
        "roll": r / 100.0, "pitch": p / 100.0, "yaw": y / 100.0,
        "vel_x": vx / 100.0, "vel_y": vy / 100.0, "vel_z": vz / 100.0,
        "frame_ref": frame_ref,
        "flags": flags,
    }


def open_socket(port: int) -> socket.socket:
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    if hasattr(socket, "SO_REUSEPORT"):
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
    sock.bind(("", port))
    sock.setblocking(False)
    return sock