# ============================================================
# autopilot_bridge.py
# Модуль 4: Мост к автопилоту (MAVLink / MSP / JSON)
#
# Конвертирует команды наведения в управляющие сигналы
# для полётного контроллера. Поддерживает:
#   - JSON export (для симулятора / внешнего потребителя)
#   - MAVLink (ArduPilot / PX4)
#   - MSP (Betaflight / INAV)
# ============================================================

import json
import socket
import struct
import time
import threading
from pathlib import Path

import numpy as np

from config import *
from config_intercept import *
from helpers import clamp


def _rc_value(normalized, scale, center=1500, rc_min=1000, rc_max=2000):
    """
    Конвертирует нормализованную команду [-1..+1] в RC PWM (мкс).

    Args:
        normalized: float в диапазоне [-1, +1]
        scale: масштаб (макс. отклонение от center)
        center: центр (обычно 1500)
        rc_min, rc_max: лимиты

    Returns:
        int — PWM в микросекундах
    """
    val = float(center) + float(normalized) * float(scale)
    return int(clamp(val, float(rc_min), float(rc_max)))


def select_proto_udp_target_state(confirmed, miss_streak):
    if not confirmed:
        return 0
    if int(miss_streak) > 0:
        return 3
    return 1


def build_proto_udp_payload(cmd):
    descriptor = int(clamp(getattr(cmd, "proto_descriptor", PROTO_UDP_DESCRIPTOR), 0, 255))
    target_state = int(clamp(getattr(cmd, "target_state", 0), 0, 255))

    if target_state == 0:
        object_id = 0
        off_y_px = 0
        off_x_px = 0
        off_y_pct = 0
        off_x_pct = 0
        bbox_area_pct = 0
    else:
        frame_w = max(1.0, float(getattr(cmd, "frame_w", 0) or 0))
        frame_h = max(1.0, float(getattr(cmd, "frame_h", 0) or 0))
        aim_x = float(getattr(cmd, "aim_x", 0.0) or 0.0)
        aim_y = float(getattr(cmd, "aim_y", 0.0) or 0.0)
        bbox_w = max(0.0, float(getattr(cmd, "bbox_w", 0.0) or 0.0))
        bbox_h = max(0.0, float(getattr(cmd, "bbox_h", 0.0) or 0.0))

        center_x = 0.5 * frame_w
        center_y = 0.5 * frame_h
        off_x_px = int(round(aim_x - center_x))
        off_y_px = int(round(center_y - aim_y))
        off_x_pct = int(round(clamp((off_x_px / max(1.0, center_x)) * 100.0, -100.0, 100.0)))
        off_y_pct = int(round(clamp((off_y_px / max(1.0, center_y)) * 100.0, -100.0, 100.0)))
        bbox_area_pct = int(round(clamp((bbox_w * bbox_h * 100.0) / max(1.0, frame_w * frame_h), 0.0, 100.0)))
        object_id = int(clamp(getattr(cmd, "target_id", 0) or 0, 0, 255))

    return struct.pack(
        "<BBBhhbbB",
        descriptor,
        object_id,
        target_state,
        int(clamp(off_y_px, -32768, 32767)),
        int(clamp(off_x_px, -32768, 32767)),
        int(clamp(off_y_pct, -128, 127)),
        int(clamp(off_x_pct, -128, 127)),
        int(clamp(bbox_area_pct, 0, 255)),
    )


class AutopilotCommand:
    """Унифицированная команда для любого бэкенда."""

    __slots__ = [
        "timestamp", "frame_id", "phase",
        "roll", "pitch", "yaw", "throttle",
        "steer_x", "steer_y", "cmd_x", "cmd_y",
        "target_id", "confirmed", "on_target",
        "target_state", "proto_descriptor",
        "frame_w", "frame_h", "aim_x", "aim_y", "bbox_w", "bbox_h",
        "range_m", "tau", "closing_vel", "confidence",
        "failsafe", "failsafe_action",
    ]

    def __init__(self):
        self.timestamp = 0.0
        self.frame_id = -1
        self.phase = "SEARCH"
        self.roll = 1500
        self.pitch = 1500
        self.yaw = 1500
        self.throttle = 1500
        self.steer_x = 0.0
        self.steer_y = 0.0
        self.cmd_x = 0.0
        self.cmd_y = 0.0
        self.target_id = None
        self.confirmed = False
        self.on_target = False
        self.target_state = 0
        self.proto_descriptor = int(PROTO_UDP_DESCRIPTOR)
        self.frame_w = None
        self.frame_h = None
        self.aim_x = None
        self.aim_y = None
        self.bbox_w = None
        self.bbox_h = None
        self.range_m = None
        self.tau = float('inf')
        self.closing_vel = 0.0
        self.confidence = 0.0
        self.failsafe = False
        self.failsafe_action = "hover"

    def to_dict(self):
        return {k: getattr(self, k) for k in self.__slots__}


class AutopilotBridge:
    """
    Главный класс моста к автопилоту.

    Принимает guidance_state + intercept_params + range_state,
    формирует AutopilotCommand и отправляет в выбранный бэкенд.
    """

    def __init__(self):
        self.enabled = bool(AUTOPILOT_ENABLE)
        self.backend_name = str(AUTOPILOT_BACKEND).lower().strip()

        # Бэкенд
        self._backend = None
        if self.enabled:
            if self.backend_name == "mavlink":
                self._backend = _MAVLinkBackend()
            elif self.backend_name == "msp":
                self._backend = _MSPBackend()
            elif self.backend_name == "proto_udp":
                self._backend = _ProtoUDPBackend()
            else:
                self._backend = _JSONBackend()

        # Failsafe
        self._failsafe_miss_count = 0
        self._failsafe_active = False

        # Rate limiting
        self._last_send_time = 0.0
        self._send_interval = 0.0
        if self.backend_name == "mavlink":
            self._send_interval = 1.0 / max(1.0, float(MAVLINK_CMD_RATE_HZ))
        elif self.backend_name == "msp":
            self._send_interval = 1.0 / max(1.0, float(MSP_CMD_RATE_HZ))

        self.last_command = AutopilotCommand()

    def start(self):
        if self.enabled and self._backend is not None:
            self._backend.start()

    def stop(self):
        if self._backend is not None:
            self._backend.stop()

    def status_line(self):
        if not self.enabled:
            return "Autopilot bridge disabled"
        backend_status = self._backend.status_line() if self._backend else "no backend"
        return f"Autopilot bridge: {self.backend_name} — {backend_status}"

    def update(self, frame_id, timestamp, guidance_state, intercept_params,
               range_state, pn_state, confirmed, miss_streak):
        """
        Формирует и отправляет команду автопилоту.

        Args:
            frame_id:         int
            timestamp:        float
            guidance_state:   dict из guidance.py
            intercept_params: dict из intercept_fsm.py
            range_state:      dict из range_estimation.py
            pn_state:         dict из proportional_navigation.py
            confirmed:        bool
            miss_streak:      int

        Returns:
            AutopilotCommand
        """
        if not self.enabled:
            return self.last_command

        cmd = AutopilotCommand()
        cmd.timestamp = timestamp
        cmd.frame_id = frame_id
        cmd.confirmed = confirmed
        cmd.target_state = select_proto_udp_target_state(confirmed, miss_streak)

        # ─── Фаза и параметры ────────────────────────────────────
        phase = "SEARCH"
        throttle_bias = 0.5
        use_pn = False

        if intercept_params is not None:
            phase = intercept_params.get("phase", "SEARCH")
            throttle_bias = intercept_params.get("throttle_bias", 0.5)
            use_pn = intercept_params.get("use_pn", False)

            # Заморозка команд в TERMINAL
            if intercept_params.get("lock_commands") and intercept_params.get("frozen_cmd"):
                frozen = intercept_params["frozen_cmd"]
                cmd.steer_x = float(frozen.get("steer_x", 0.0))
                cmd.steer_y = float(frozen.get("steer_y", 0.0))
                cmd.cmd_x = float(frozen.get("cmd_x", 0.0))
                cmd.cmd_y = float(frozen.get("cmd_y", 0.0))
                use_pn = False  # используем замороженные

        cmd.phase = phase

        # ─── Наведение ───────────────────────────────────────────
        if use_pn and pn_state is not None and pn_state.get("active"):
            # Используем PN-команды
            cmd.steer_x = float(pn_state.get("accel_x", 0.0))
            cmd.steer_y = float(pn_state.get("accel_y", 0.0))
            cmd.cmd_x = cmd.steer_x
            cmd.cmd_y = cmd.steer_y
        elif not intercept_params or not intercept_params.get("lock_commands"):
            # Screen guidance
            if guidance_state is not None:
                cmd.steer_x = float(guidance_state.get("steer_x", 0.0))
                cmd.steer_y = float(guidance_state.get("steer_y", 0.0))
                cmd.cmd_x = float(guidance_state.get("cmd_x", 0.0))
                cmd.cmd_y = float(guidance_state.get("cmd_y", 0.0))
                cmd.on_target = bool(guidance_state.get("on_target", False))
                cmd.confidence = float(guidance_state.get("confidence", 0.0))
                cmd.target_id = guidance_state.get("target_id")
                cmd.frame_w = guidance_state.get("frame_w")
                cmd.frame_h = guidance_state.get("frame_h")
                cmd.aim_x = guidance_state.get("aim_x")
                cmd.aim_y = guidance_state.get("aim_y")
                cmd.bbox_w = guidance_state.get("box_w")
                cmd.bbox_h = guidance_state.get("box_h")

        # ─── Range info ─────────────────────────────────────────
        if range_state is not None:
            cmd.range_m = range_state.get("range_m")
            cmd.tau = range_state.get("tau", float('inf'))
            cmd.closing_vel = range_state.get("closing_vel", 0.0)

        # ─── RC channels ────────────────────────────────────────
        cmd.roll = _rc_value(cmd.steer_x, AUTOPILOT_ROLL_SCALE,
                             1500, AUTOPILOT_RC_MIN, AUTOPILOT_RC_MAX)
        cmd.pitch = _rc_value(-cmd.steer_y, AUTOPILOT_PITCH_SCALE,
                              1500, AUTOPILOT_RC_MIN, AUTOPILOT_RC_MAX)
        cmd.yaw = 1500  # yaw управляется отдельно или через steer_x
        cmd.throttle = _rc_value(
            throttle_bias * 2.0 - 1.0,  # [0..1] → [-1..+1]
            float(AUTOPILOT_RC_MAX - AUTOPILOT_THROTTLE_CENTER),
            AUTOPILOT_THROTTLE_CENTER,
            AUTOPILOT_RC_MIN,
            AUTOPILOT_RC_MAX,
        )

        # ─── Failsafe ───────────────────────────────────────────
        if AUTOPILOT_FAILSAFE_ENABLE:
            if not confirmed:
                self._failsafe_miss_count += 1
            else:
                self._failsafe_miss_count = 0

            if self._failsafe_miss_count >= int(AUTOPILOT_FAILSAFE_MISS_FRAMES):
                self._failsafe_active = True
                cmd.failsafe = True
                cmd.failsafe_action = str(AUTOPILOT_FAILSAFE_ACTION)
                # В failsafe: центрируем стики, throttle на hover
                cmd.roll = 1500
                cmd.pitch = 1500
                cmd.yaw = 1500
                cmd.throttle = 1500
                cmd.steer_x = 0.0
                cmd.steer_y = 0.0
            else:
                self._failsafe_active = False

        # ─── Отправка ───────────────────────────────────────────
        now = time.perf_counter()
        if (now - self._last_send_time) >= self._send_interval:
            if self._backend is not None:
                self._backend.send(cmd)
            self._last_send_time = now

        self.last_command = cmd
        return cmd


# ─────────────────────────────────────────────────────────────
# BACKENDS
# ─────────────────────────────────────────────────────────────

class _JSONBackend:
    """
    Экспорт команд в JSON-файл.
    Расширенная версия текущего guidance_state.json.
    """

    def __init__(self):
        self.path = Path(AUTOPILOT_JSON_PATH)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.every = int(max(1, AUTOPILOT_JSON_EVERY))
        self._frame_count = 0

    def start(self):
        pass

    def stop(self):
        pass

    def status_line(self):
        return f"JSON export → {self.path}"

    def send(self, cmd):
        self._frame_count += 1
        if self._frame_count % self.every != 0:
            return

        payload = cmd.to_dict()
        # Сериализация: inf → null для JSON
        for k, v in payload.items():
            if isinstance(v, float) and (v == float('inf') or v == float('-inf')):
                payload[k] = None

        tmp = self.path.with_suffix(self.path.suffix + ".tmp")
        try:
            with tmp.open("w", encoding="utf-8") as f:
                json.dump(payload, f, ensure_ascii=True, indent=2)
            tmp.replace(self.path)
        except OSError:
            pass


class _MAVLinkBackend:
    """
    MAVLink backend для ArduPilot / PX4.

    Отправляет RC_CHANNELS_OVERRIDE или MANUAL_CONTROL.
    Требует: pip install pymavlink
    """

    def __init__(self):
        self._conn = None
        self._connected = False

    def start(self):
        try:
            from pymavlink import mavutil
            self._conn = mavutil.mavlink_connection(
                str(MAVLINK_CONNECTION),
                source_system=int(MAVLINK_SYSTEM_ID),
                source_component=int(MAVLINK_COMPONENT_ID),
            )
            self._connected = True
        except Exception as e:
            self._connected = False
            print(f"[autopilot] MAVLink connection failed: {e}")

    def stop(self):
        if self._conn is not None:
            try:
                self._conn.close()
            except Exception:
                pass

    def status_line(self):
        if self._connected:
            return f"MAVLink connected: {MAVLINK_CONNECTION}"
        return "MAVLink not connected"

    def send(self, cmd):
        if not self._connected or self._conn is None:
            return

        try:
            if cmd.failsafe:
                # Failsafe: отправляем hover/RTL
                if cmd.failsafe_action == "rtl":
                    self._conn.set_mode_rtl()
                elif cmd.failsafe_action == "land":
                    self._conn.set_mode("LAND")
                else:
                    # hover — отпускаем стики
                    self._send_rc_override(cmd)
            else:
                self._send_rc_override(cmd)
        except Exception:
            pass

    def _send_rc_override(self, cmd):
        """Отправка RC_CHANNELS_OVERRIDE."""
        self._conn.mav.rc_channels_override_send(
            self._conn.target_system,
            self._conn.target_component,
            cmd.roll,      # chan1 — roll
            cmd.pitch,     # chan2 — pitch
            cmd.throttle,  # chan3 — throttle
            cmd.yaw,       # chan4 — yaw
            0, 0, 0, 0,   # chan5-8
        )


class _MSPBackend:
    """
    MSP (MultiWii Serial Protocol) backend для Betaflight / INAV.

    Отправляет MSP_SET_RAW_RC (200) через serial.
    Требует: pip install pyserial
    """

    def __init__(self):
        self._ser = None
        self._connected = False

    def start(self):
        try:
            import serial
            self._ser = serial.Serial(
                str(MSP_PORT),
                int(MSP_BAUD),
                timeout=0.05,
            )
            self._connected = True
        except Exception as e:
            self._connected = False
            print(f"[autopilot] MSP serial open failed: {e}")

    def stop(self):
        if self._ser is not None:
            try:
                self._ser.close()
            except Exception:
                pass

    def status_line(self):
        if self._connected:
            return f"MSP connected: {MSP_PORT}@{MSP_BAUD}"
        return "MSP not connected"

    def send(self, cmd):
        if not self._connected or self._ser is None:
            return

        try:
            channels = [
                cmd.roll,      # roll
                cmd.pitch,     # pitch
                cmd.throttle,  # throttle
                cmd.yaw,       # yaw
                1500, 1500, 1500, 1500,  # aux 1-4
            ]
            self._send_msp_set_raw_rc(channels)
        except Exception:
            pass

    def _send_msp_set_raw_rc(self, channels):
        """
        Формирует и отправляет MSP пакет SET_RAW_RC (cmd=200).

        Формат MSP v1:
        $M< [size] [cmd] [data...] [checksum]
        """
        data = bytearray()
        for ch in channels[:8]:
            val = int(clamp(ch, 1000, 2000))
            data.append(val & 0xFF)
            data.append((val >> 8) & 0xFF)

        size = len(data)
        cmd_id = 200  # MSP_SET_RAW_RC

        # Checksum = XOR of size, cmd, data
        checksum = size ^ cmd_id
        for b in data:
            checksum ^= b

        packet = bytearray([
            ord('$'), ord('M'), ord('<'),
            size,
            cmd_id,
        ])
        packet.extend(data)
        packet.append(checksum & 0xFF)

        self._ser.write(packet)


class _ProtoUDPBackend:
    def __init__(self):
        self._sock = None
        self._enabled = bool(PROTO_UDP_ENABLE)
        self._host = str(PROTO_UDP_HOST)
        self._port = int(PROTO_UDP_PORT)

    def start(self):
        if not self._enabled:
            return
        try:
            self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        except OSError as e:
            self._sock = None
            print(f"[autopilot] proto_udp socket failed: {e}")

    def stop(self):
        if self._sock is not None:
            try:
                self._sock.close()
            except OSError:
                pass
            self._sock = None

    def status_line(self):
        if not self._enabled:
            return "PROTO_UDP disabled"
        return f"PROTO_UDP ready: {self._host}:{self._port}"

    def send(self, cmd):
        if (not self._enabled) or (self._sock is None):
            return
        try:
            payload = build_proto_udp_payload(cmd)
            self._sock.sendto(payload, (self._host, self._port))
        except OSError:
            pass