"""
将 JPEG 帧通过 OrangePi 的 FrameHeader(16B) 分片协议转发出去（UDP）。
用于 Python -> Unity 图像显示链路，便于 Unity 端复用同一套分片重组逻辑。

Header（小端）:
  uint32 magic     = 0x57494649 ("WIFI")
  uint32 frame_id
  uint16 seq
  uint16 total
  uint16 data_len
  uint16 reserved
"""

from __future__ import annotations

import socket
import struct
import threading
from dataclasses import dataclass


MAGIC = 0x57494649
HEADER_FMT = "<I I H H H H"
HEADER_SIZE = struct.calcsize(HEADER_FMT)
MAX_UDP_SIZE = 1400
MAX_PAYLOAD = MAX_UDP_SIZE - HEADER_SIZE  # 1384


@dataclass
class ForwardConfig:
    target_ip: str = "127.0.0.1"
    target_port: int = 12366


class UDPFrameForwarder:
    def __init__(self, target_ip: str = "127.0.0.1", target_port: int = 12366) -> None:
        self.cfg = ForwardConfig(target_ip=target_ip, target_port=target_port)
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self._frame_id = 0
        self._lock = threading.Lock()

    def set_target(self, ip: str, port: int) -> None:
        with self._lock:
            self.cfg.target_ip = ip
            self.cfg.target_port = int(port)

    def send_jpeg(self, jpeg_bytes: bytes) -> None:
        if not jpeg_bytes:
            return
        with self._lock:
            target = (self.cfg.target_ip, self.cfg.target_port)
        frame_id = self._frame_id
        self._frame_id = (self._frame_id + 1) & 0xFFFFFFFF

        total = (len(jpeg_bytes) + MAX_PAYLOAD - 1) // MAX_PAYLOAD
        if total <= 0:
            return
        if total > 65535:
            # 单帧过大，不发送
            return

        for seq in range(total):
            start = seq * MAX_PAYLOAD
            end = min(start + MAX_PAYLOAD, len(jpeg_bytes))
            payload = jpeg_bytes[start:end]
            header = struct.pack(HEADER_FMT, MAGIC, frame_id, seq, total, len(payload), 0)
            self.sock.sendto(header + payload, target)

    def close(self) -> None:
        try:
            self.sock.close()
        except Exception:
            pass