"""protocol.py —— 光电吊舱控制协议 G 版编解码（纯函数 + 数据类，无硬件依赖）。

规格：docs/光电吊舱控制协议G版-PB.docx。

四种定长帧，统一规则：
  * 多字节小端（低字节在前，高字节在后）；
  * 末尾校验和 = 从第 0 字节累加取低 8 位（sum(body) & 0xFF）；
  * 文档样例 eb 90 2b 00..00 a6 已锁定校验算法。

  表2-1 控制指令   同步字 EB 90  16B  控制器 → 吊舱  触发式
  表2-2 周期遥测   同步字 EE 16  32B  吊舱 → 控制器  60Hz
  表3   载机姿态   同步字 EB 91  32B  控制器 → 吊舱  1~60Hz
  表4   目标经纬度 同步字 EE 18  32B  吊舱 → 控制器  1~60Hz

控制器（本桥接）发送 表2-1/表3、接收 表2-2/表4；为便于回环测试与后续吊舱
模拟器，四种帧均提供对称的 encode/decode。
"""

from dataclasses import dataclass, field
from enum import IntEnum
import struct

# --- 同步字与帧长 -----------------------------------------------------------
SYNC_CONTROL = (0xEB, 0x90)    # 表2-1
SYNC_TELEMETRY = (0xEE, 0x16)  # 表2-2
SYNC_ATTITUDE = (0xEB, 0x91)   # 表3
SYNC_TARGET = (0xEE, 0x18)     # 表4

CONTROL_LEN = 16
TELEMETRY_LEN = 32
ATTITUDE_LEN = 32
TARGET_LEN = 32

# --- 量纲（每 LSB 物理量）---------------------------------------------------
ANGLE_SCALE = 0.01    # 角度 / 偏移角，°
RATE_SCALE = 0.01     # 角速度，°/s
RANGE_SCALE = 0.1     # 激光测距，m
ALT_SCALE = 0.2       # 高度，m
LATLON_SCALE = 1e-7   # 经纬度，°
SPEED_SCALE = 0.5     # 空速 / 地速，m/s

_S16_MIN, _S16_MAX = -32768, 32767
_U16_MAX = 65535
_S32_MIN, _S32_MAX = -(2 ** 31), 2 ** 31 - 1


def checksum(body: bytes) -> int:
    """从第 0 字节累加取低 8 位（body 不含校验和自身）。"""
    return sum(body) & 0xFF


def _to_s16(value: float, scale: float) -> int:
    return max(_S16_MIN, min(_S16_MAX, int(round(value / scale))))


def _to_u16(value: float, scale: float) -> int:
    return max(0, min(_U16_MAX, int(round(value / scale))))


def _to_s32(value: float, scale: float) -> int:
    return max(_S32_MIN, min(_S32_MAX, int(round(value / scale))))


class ControlCode(IntEnum):
    """表2-1 字节 2「控制信息」取值。"""

    NONE = 0x00                 # 空指令
    VISIBLE_1 = 0x01            # 可见光1
    VISIBLE_2 = 0x02            # 可见光2
    IR_1 = 0x03                 # 红外1
    IR_2 = 0x04                 # 红外2
    TARGET_DETECT_ON = 0x07     # 目标识别开（参数X：0=人/车，1=无人机）
    TARGET_DETECT_OFF = 0x08    # 目标识别关（默认）
    VIDEO_RECORD_ON = 0x09      # 视频录制开
    VIDEO_RECORD_OFF = 0x0A     # 视频录制关（默认）
    BOX_TRACK = 0x0C            # 框选跟踪（左上 X/Y + 右下 X2/Y2）
    POINT_TRACK = 0x0D          # 指点跟踪（X/Y）
    STOP_TRACK = 0x0E           # 停止跟踪
    AUTO_LOCK_ON_FOUND = 0x0F   # 发现目标自动锁定
    DETECT_TO_TRACK = 0x10      # 识别转跟踪
    IR_WHITE_HOT = 0x11         # 红外白热
    IR_BLACK_HOT = 0x12         # 红外黑热
    ANTI_OCCLUSION_OFF = 0x17   # 关抗遮挡
    ANTI_OCCLUSION_ON = 0x18    # 开抗遮挡
    GIMBAL_SEARCH = 0x24        # 云台搜索（X=方位转速0.1°/s，Y=俯仰转速0.1°/s）
    OPTICAL_ZOOM = 0x25         # 光学变焦（X=倍数0.1倍）
    SET_FRAME_ANGLE = 0x26      # 指定框架角度（X=方位0.01°，Y=俯仰0.01°）
    MOTOR_ON = 0x27             # 开电机
    MOTOR_OFF = 0x28            # 关电机
    FOLLOW_OFF = 0x29           # 关跟随
    FOLLOW_ON = 0x2A            # 跟随
    CENTER = 0x2B               # 归中
    SUPPRESS_GYRO_DRIFT = 0x2C  # 抑制陀螺仪漂移
    ELECTRIC_LOCK = 0x30        # 电锁模式
    ELECTRIC_UNLOCK = 0x31      # 解除电锁
    PIP_ON = 0x40               # 开画中画
    PIP_OFF = 0x41              # 关画中画
    OSD_ON = 0x51               # OSD 显示开
    OSD_OFF = 0x52              # OSD 显示关
    SET_ZOOM = 0x5A             # 指定变焦倍数（X=倍数0.1倍）
    TAKE_PHOTO = 0x5B           # 拍照指令
    EIS_ON = 0x5C               # 开电子稳像
    EIS_OFF = 0x5D              # 关电子稳像
    SERVO_STAB_ON = 0x60        # 开伺服稳像
    SERVO_STAB_OFF = 0x61       # 关伺服稳像
    VERSION_QUERY = 0x62        # 版本查询
    SET_BITRATE = 0x64          # 设置码率（X=kbps）


# --- 表2-1 控制指令 ---------------------------------------------------------
@dataclass
class ControlCommand:
    """控制器 → 吊舱 控制帧。param* 语义随 code 变化，见 ControlCode 注释。"""

    code: int = ControlCode.NONE
    param_x: int = 0          # S16
    param_y: int = 0          # S16
    param3: int = 0           # U8（画中画小画面等）
    file_save_mode: int = 0   # S8（0=时间戳，1=前缀+序号）
    param_x2: int = 0         # S16（框选跟踪右下角横向）
    param_y2: int = 0         # S16（框选跟踪右下角纵向）


def encode_control(cmd: ControlCommand) -> bytes:
    body = struct.pack(
        "<BBBhhBbhh2s",
        SYNC_CONTROL[0], SYNC_CONTROL[1],
        cmd.code & 0xFF,
        _to_s16(cmd.param_x, 1), _to_s16(cmd.param_y, 1),
        cmd.param3 & 0xFF,
        max(-128, min(127, int(cmd.file_save_mode))),
        _to_s16(cmd.param_x2, 1), _to_s16(cmd.param_y2, 1),
        b"\x00\x00",
    )
    return body + bytes([checksum(body)])


def decode_control(data: bytes) -> ControlCommand:
    _check(data, CONTROL_LEN, SYNC_CONTROL)
    code, x, y, param3, save, x2, y2 = struct.unpack("<BhhBbhh", data[2:13])
    return ControlCommand(code=code, param_x=x, param_y=y, param3=param3,
                          file_save_mode=save, param_x2=x2, param_y2=y2)


# --- 表2-2 周期遥测 ---------------------------------------------------------
@dataclass
class Telemetry:
    """吊舱 → 控制器 周期遥测帧（60Hz）。角度单位 °，角速度 °/s，测距 m。"""

    status1: int = 0
    status2: int = 0
    status3: int = 0
    off_target_x_deg: float = 0.0   # 目标脱靶量 X 轴偏移角
    off_target_y_deg: float = 0.0   # 目标脱靶量 Y 轴偏移角
    frame_roll_deg: float = 0.0     # 横滚框架角
    frame_pitch_deg: float = 0.0    # 俯仰框架角
    frame_yaw_deg: float = 0.0      # 方位框架角
    att_roll_deg: float = 0.0       # 横滚姿态角
    att_pitch_deg: float = 0.0      # 俯仰姿态角
    roll_rate_dps: float = 0.0
    pitch_rate_dps: float = 0.0
    yaw_rate_dps: float = 0.0
    laser_range_m: float | None = None  # None 表示无效（原始 0）
    self_check: int = 0

    @property
    def track_source(self) -> int:
        """Bit7-6 跟踪视频源：0可见1 1可见2 2红外1 3红外2。"""
        return (self.status1 >> 6) & 0x3

    @property
    def is_target_locked(self) -> bool:
        """Bit2 目标跟踪状态：True=锁定，False=搜索。"""
        return bool((self.status1 >> 2) & 0x1)

    @property
    def video_storing(self) -> bool:
        return bool((self.status2 >> 5) & 0x1)

    @property
    def motor_on(self) -> bool:
        return bool((self.status2 >> 3) & 0x1)

    @property
    def follow_on(self) -> bool:
        return bool((self.status2 >> 2) & 0x1)

    @property
    def electric_lock_on(self) -> bool:
        return bool((self.status2 >> 1) & 0x1)

    @property
    def big_picture(self) -> int:
        return (self.status3 >> 6) & 0x3

    @property
    def small_picture(self) -> int:
        return (self.status3 >> 4) & 0x3


def encode_telemetry(t: Telemetry) -> bytes:
    laser = 0 if t.laser_range_m is None else _to_u16(t.laser_range_m, RANGE_SCALE)
    body = struct.pack(
        "<BBBBBBhhhhhhhhhhHB2s",
        SYNC_TELEMETRY[0], SYNC_TELEMETRY[1],
        t.status1 & 0xFF, t.status2 & 0xFF, 0, t.status3 & 0xFF,
        _to_s16(t.off_target_x_deg, ANGLE_SCALE),
        _to_s16(t.off_target_y_deg, ANGLE_SCALE),
        _to_s16(t.frame_roll_deg, ANGLE_SCALE),
        _to_s16(t.frame_pitch_deg, ANGLE_SCALE),
        _to_s16(t.frame_yaw_deg, ANGLE_SCALE),
        _to_s16(t.att_roll_deg, ANGLE_SCALE),
        _to_s16(t.att_pitch_deg, ANGLE_SCALE),
        _to_s16(t.roll_rate_dps, RATE_SCALE),
        _to_s16(t.pitch_rate_dps, RATE_SCALE),
        _to_s16(t.yaw_rate_dps, RATE_SCALE),
        laser,
        t.self_check & 0xFF,
        b"\x00\x00",
    )
    return body + bytes([checksum(body)])


def decode_telemetry(data: bytes) -> Telemetry:
    _check(data, TELEMETRY_LEN, SYNC_TELEMETRY)
    (s1, s2, _rsv, s3, ox, oy, froll, fpitch, fyaw, aroll, apitch,
     rrate, prate, yrate, laser, selfcheck, _rsv2) = struct.unpack(
        "<BBBBhhhhhhhhhhHB2s", data[2:31])
    return Telemetry(
        status1=s1, status2=s2, status3=s3,
        off_target_x_deg=ox * ANGLE_SCALE,
        off_target_y_deg=oy * ANGLE_SCALE,
        frame_roll_deg=froll * ANGLE_SCALE,
        frame_pitch_deg=fpitch * ANGLE_SCALE,
        frame_yaw_deg=fyaw * ANGLE_SCALE,
        att_roll_deg=aroll * ANGLE_SCALE,
        att_pitch_deg=apitch * ANGLE_SCALE,
        roll_rate_dps=rrate * RATE_SCALE,
        pitch_rate_dps=prate * RATE_SCALE,
        yaw_rate_dps=yrate * RATE_SCALE,
        laser_range_m=(laser * RANGE_SCALE) if laser != 0 else None,
        self_check=selfcheck,
    )


# --- 表3 载机姿态和经纬度 ---------------------------------------------------
@dataclass
class CarrierAttitude:
    """控制器 → 吊舱 载机姿态/位置/时间帧（供吊舱解算目标经纬度）。"""

    roll_deg: float = 0.0
    pitch_deg: float = 0.0
    yaw_deg: float = 0.0
    latitude_deg: float = 0.0
    longitude_deg: float = 0.0
    abs_alt_m: float = 0.0
    rel_alt_m: float = 0.0
    year: int = 2000        # 完整年份（编码时存 year-2000）
    month: int = 1
    day: int = 1
    hour: int = 0
    minute: int = 0
    second: int = 0
    centisecond: int = 0    # 百分秒，10ms 单位，0~99
    airspeed_mps: float = 0.0
    ground_speed_mps: float = 0.0


def encode_attitude(a: CarrierAttitude) -> bytes:
    body = struct.pack(
        "<BBhhhiihhBBBBBBBHH",
        SYNC_ATTITUDE[0], SYNC_ATTITUDE[1],
        _to_s16(a.roll_deg, ANGLE_SCALE),
        _to_s16(a.pitch_deg, ANGLE_SCALE),
        _to_s16(a.yaw_deg, ANGLE_SCALE),
        _to_s32(a.latitude_deg, LATLON_SCALE),
        _to_s32(a.longitude_deg, LATLON_SCALE),
        _to_s16(a.abs_alt_m, ALT_SCALE),
        _to_s16(a.rel_alt_m, ALT_SCALE),
        (a.year - 2000) & 0xFF, a.month & 0xFF, a.day & 0xFF,
        a.hour & 0xFF, a.minute & 0xFF, a.second & 0xFF, a.centisecond & 0xFF,
        _to_u16(a.airspeed_mps, SPEED_SCALE),
        _to_u16(a.ground_speed_mps, SPEED_SCALE),
    )
    return body + bytes([checksum(body)])


def decode_attitude(data: bytes) -> CarrierAttitude:
    _check(data, ATTITUDE_LEN, SYNC_ATTITUDE)
    (roll, pitch, yaw, lat, lon, absalt, relalt, yr, mo, day, hour, mi, sec,
     cs, air, gs) = struct.unpack("<hhhiihhBBBBBBBHH", data[2:31])
    return CarrierAttitude(
        roll_deg=roll * ANGLE_SCALE,
        pitch_deg=pitch * ANGLE_SCALE,
        yaw_deg=yaw * ANGLE_SCALE,
        latitude_deg=lat * LATLON_SCALE,
        longitude_deg=lon * LATLON_SCALE,
        abs_alt_m=absalt * ALT_SCALE,
        rel_alt_m=relalt * ALT_SCALE,
        year=2000 + yr, month=mo, day=day, hour=hour, minute=mi,
        second=sec, centisecond=cs,
        airspeed_mps=air * SPEED_SCALE,
        ground_speed_mps=gs * SPEED_SCALE,
    )


# --- 表4 目标经纬度 ---------------------------------------------------------
@dataclass
class TargetGeo:
    """吊舱 → 控制器 目标经纬度帧（吊舱解算结果）。"""

    latitude_deg: float = 0.0
    longitude_deg: float = 0.0
    abs_alt_m: float = 0.0
    rel_alt_m: float = 0.0
    year: int = 2000
    month: int = 1
    day: int = 1
    hour: int = 0
    minute: int = 0
    second: int = 0
    centisecond: int = 0


def encode_target(g: TargetGeo) -> bytes:
    body = struct.pack(
        "<BBiihhBBBBBBB10s",
        SYNC_TARGET[0], SYNC_TARGET[1],
        _to_s32(g.latitude_deg, LATLON_SCALE),
        _to_s32(g.longitude_deg, LATLON_SCALE),
        _to_s16(g.abs_alt_m, ALT_SCALE),
        _to_s16(g.rel_alt_m, ALT_SCALE),
        (g.year - 2000) & 0xFF, g.month & 0xFF, g.day & 0xFF,
        g.hour & 0xFF, g.minute & 0xFF, g.second & 0xFF, g.centisecond & 0xFF,
        b"\x00" * 10,
    )
    return body + bytes([checksum(body)])


def decode_target(data: bytes) -> TargetGeo:
    _check(data, TARGET_LEN, SYNC_TARGET)
    (lat, lon, absalt, relalt, yr, mo, day, hour, mi, sec, cs,
     _rsv) = struct.unpack("<iihhBBBBBBB10s", data[2:31])
    return TargetGeo(
        latitude_deg=lat * LATLON_SCALE,
        longitude_deg=lon * LATLON_SCALE,
        abs_alt_m=absalt * ALT_SCALE,
        rel_alt_m=relalt * ALT_SCALE,
        year=2000 + yr, month=mo, day=day, hour=hour, minute=mi,
        second=sec, centisecond=cs,
    )


def _check(data: bytes, length: int, sync: tuple) -> None:
    if len(data) != length:
        raise ValueError(f"帧长错误: 期望 {length} 字节, 实际 {len(data)}")
    if data[0] != sync[0] or data[1] != sync[1]:
        raise ValueError(
            f"同步字错误: 期望 {sync[0]:02X} {sync[1]:02X}, "
            f"实际 {data[0]:02X} {data[1]:02X}")
    if checksum(data[:length - 1]) != data[length - 1]:
        raise ValueError("校验和错误")


# --- 字节流分帧器 -----------------------------------------------------------
_FRAME_SPECS = {
    SYNC_CONTROL: (CONTROL_LEN, decode_control),
    SYNC_TELEMETRY: (TELEMETRY_LEN, decode_telemetry),
    SYNC_ATTITUDE: (ATTITUDE_LEN, decode_attitude),
    SYNC_TARGET: (TARGET_LEN, decode_target),
}


class FrameParser:
    """串口字节流分帧：按同步字滑窗取定长帧，校验通过则解码输出。

    默认只接收吊舱 → 控制器的两种帧（遥测 EE16 / 目标 EE18）；可通过 kinds
    指定要识别的同步字集合（如吊舱模拟器侧识别 EB90/EB91）。粘包、半帧、坏校验
    均安全处理：坏校验丢弃该同步字首字节后重扫，半帧留待下次 feed。
    """

    def __init__(self, kinds: tuple = (SYNC_TELEMETRY, SYNC_TARGET)) -> None:
        for k in kinds:
            if k not in _FRAME_SPECS:
                raise ValueError(f"未知帧类型: {k}")
        self._specs = {k: _FRAME_SPECS[k] for k in kinds}
        self._buf = bytearray()

    def feed(self, data: bytes):
        """喂入任意长度字节，返回本次可解析出的帧列表（按到达顺序）。"""
        self._buf.extend(data)
        frames = []
        while True:
            idx, sync = self._find_sync()
            if idx is None:
                # 无同步字：仅保留最后 1 字节（可能是下一帧同步字首字节）
                if len(self._buf) > 1:
                    del self._buf[:-1]
                break
            if idx > 0:
                del self._buf[:idx]
            length, decoder = self._specs[sync]
            if len(self._buf) < length:
                break  # 半帧，等待更多数据
            frame = bytes(self._buf[:length])
            if checksum(frame[:length - 1]) == frame[length - 1]:
                frames.append(decoder(frame))
                del self._buf[:length]
            else:
                del self._buf[:1]  # 坏校验：跳过该同步字首字节重扫
        return frames

    def _find_sync(self):
        earliest_idx, earliest_sync = None, None
        for sync in self._specs:
            i = self._buf.find(bytes(sync))
            if i >= 0 and (earliest_idx is None or i < earliest_idx):
                earliest_idx, earliest_sync = i, sync
        return earliest_idx, earliest_sync