from __future__ import annotations

import importlib
import math
import threading
import time
from contextlib import suppress
from dataclasses import dataclass
from typing import Any, Final

from k1link.simulation.contracts import AckermannControlSetpoint

OFFBOARD_HEARTBEAT_SECONDS: Final = 0.1
OFFBOARD_WARMUP_HEARTBEATS: Final = 12
ROS_PUBLISHER_STARTUP_TIMEOUT_SECONDS: Final = 15.0
OFFBOARD_ADMISSION_TIMEOUT_SECONDS: Final = 60.0
STOCK_ROVER_MAX_THROTTLE_SPEED_MPS: Final = 3.1
VEHICLE_COMMAND_DO_SET_MODE: Final = 176
VEHICLE_COMMAND_COMPONENT_ARM_DISARM: Final = 400


class Px4RoverControlError(RuntimeError):
    """The virtual-only PX4 rover command boundary failed closed."""


@dataclass(frozen=True, slots=True)
class Px4RoverControlSnapshot:
    run_id: str
    armed: bool
    offboard: bool
    ttl_expired_count: int


class Ros2Px4AckermannControl:
    """Publish canonical speed+steering through the PX4 ROS 2 offboard boundary."""

    def __init__(
        self,
        *,
        heartbeat_seconds: float = OFFBOARD_HEARTBEAT_SECONDS,
        admission_timeout_seconds: float = OFFBOARD_ADMISSION_TIMEOUT_SECONDS,
    ) -> None:
        if not 0.05 <= heartbeat_seconds <= 0.25:
            raise ValueError("PX4 heartbeat period must remain within 4..20 Hz")
        if admission_timeout_seconds <= 0:
            raise ValueError("PX4 admission timeout must be positive")
        self.heartbeat_seconds = heartbeat_seconds
        self.admission_timeout_seconds = admission_timeout_seconds
        self._lock = threading.RLock()
        self._started = threading.Event()
        self._control_ready = threading.Event()
        self._disarmed = threading.Event()
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None
        self._run_id: str | None = None
        self._latest_speed_mps = 0.0
        self._latest_steering_normalized = 0.0
        self._valid_until_monotonic_ns = 0
        self._ttl_expired_count = 0
        self._ttl_expired_for_latest = False
        self._armed = False
        self._offboard = False
        self._disarm_requested = False
        self._failure: BaseException | None = None

    def start(self, run_id: str) -> Px4RoverControlSnapshot:
        with self._lock:
            if self._thread is not None:
                raise Px4RoverControlError("PX4 command controller already owns a run")
            self._run_id = run_id
            self._latest_speed_mps = 0.0
            self._latest_steering_normalized = 0.0
            self._valid_until_monotonic_ns = 0
            self._ttl_expired_for_latest = False
            self._failure = None
            self._stop.clear()
            self._started.clear()
            self._control_ready.clear()
            self._disarmed.clear()
            self._disarm_requested = False
            self._thread = threading.Thread(
                target=self._run,
                name=f"px4-rover-control:{run_id}",
                daemon=True,
            )
            self._thread.start()
        if not self._started.wait(timeout=ROS_PUBLISHER_STARTUP_TIMEOUT_SECONDS):
            self._abort_thread()
            raise Px4RoverControlError("PX4 ROS 2 command publisher did not start")
        self._raise_failure()
        if not self._control_ready.wait(timeout=self.admission_timeout_seconds):
            self._abort_thread()
            self._raise_failure()
            raise Px4RoverControlError("PX4 did not confirm armed Offboard rover control")
        self._raise_failure()
        return self.snapshot(run_id)

    def submit(self, command: AckermannControlSetpoint) -> Px4RoverControlSnapshot:
        with self._lock:
            if self._run_id != command.run_id or self._thread is None:
                raise Px4RoverControlError("PX4 command belongs to another or inactive run")
            self._raise_failure_locked()
            ttl_ns = command.valid_until_sim_ns - command.issued_at_sim_ns
            self._latest_speed_mps = command.speed_mps
            self._latest_steering_normalized = command.steering_normalized
            self._valid_until_monotonic_ns = time.monotonic_ns() + ttl_ns
            self._ttl_expired_for_latest = False
        return self.snapshot(command.run_id)

    def stop(self, run_id: str) -> Px4RoverControlSnapshot:
        with self._lock:
            if self._run_id != run_id or self._thread is None:
                raise Px4RoverControlError("PX4 stop belongs to another or inactive run")
            self._latest_speed_mps = 0.0
            self._latest_steering_normalized = 0.0
            self._valid_until_monotonic_ns = 0
            self._disarm_requested = True
            self._disarmed.clear()
        self._disarmed.wait(timeout=2)
        snapshot = self.snapshot(run_id)
        self._stop.set()
        thread = self._thread
        thread.join(timeout=3)
        if thread.is_alive():
            raise Px4RoverControlError("PX4 command publisher did not stop")
        self._raise_failure()
        with self._lock:
            self._thread = None
            self._run_id = None
        return snapshot

    def snapshot(self, run_id: str) -> Px4RoverControlSnapshot:
        with self._lock:
            if self._run_id != run_id:
                raise Px4RoverControlError("PX4 snapshot belongs to another run")
            return Px4RoverControlSnapshot(
                run_id=run_id,
                armed=self._armed,
                offboard=self._offboard,
                ttl_expired_count=self._ttl_expired_count,
            )

    def _abort_thread(self) -> None:
        self._stop.set()
        thread = self._thread
        if thread is not None:
            thread.join(timeout=3)
        with self._lock:
            self._thread = None
            self._run_id = None

    def _raise_failure(self) -> None:
        with self._lock:
            self._raise_failure_locked()

    def _raise_failure_locked(self) -> None:
        if self._failure is not None:
            raise Px4RoverControlError(
                f"PX4 ROS 2 command publisher failed: {type(self._failure).__name__}"
            ) from self._failure

    def _run(self) -> None:
        rclpy: Any | None = None
        node: Any | None = None
        try:
            rclpy = importlib.import_module("rclpy")
            qos_module = importlib.import_module("rclpy.qos")
            messages = importlib.import_module("px4_msgs.msg")
            rclpy.init(args=None)
            node = rclpy.create_node("missioncore_px4_ackermann_control")
            qos = qos_module.QoSProfile(
                reliability=qos_module.ReliabilityPolicy.BEST_EFFORT,
                durability=qos_module.DurabilityPolicy.TRANSIENT_LOCAL,
                history=qos_module.HistoryPolicy.KEEP_LAST,
                depth=1,
            )
            offboard_publisher = node.create_publisher(
                messages.OffboardControlMode,
                "/fmu/in/offboard_control_mode",
                qos,
            )
            throttle_publisher = node.create_publisher(
                messages.RoverThrottleSetpoint,
                "/fmu/in/rover_throttle_setpoint",
                qos,
            )
            steering_publisher = node.create_publisher(
                messages.RoverSteeringSetpoint,
                "/fmu/in/rover_steering_setpoint",
                qos,
            )
            command_publisher = node.create_publisher(
                messages.VehicleCommand,
                "/fmu/in/vehicle_command",
                qos,
            )

            def receive_status(status: Any) -> None:
                with self._lock:
                    self._armed = status.arming_state == messages.VehicleStatus.ARMING_STATE_ARMED
                    self._offboard = (
                        status.nav_state == messages.VehicleStatus.NAVIGATION_STATE_OFFBOARD
                    )
                    if self._armed and self._offboard:
                        self._control_ready.set()

            node.create_subscription(
                messages.VehicleStatus,
                "/fmu/out/vehicle_status_v1",
                receive_status,
                qos,
            )
            self._started.set()
            heartbeat_count = 0
            last_mode_request_monotonic = 0.0
            while not self._stop.is_set():
                now_monotonic = time.monotonic()
                now_monotonic_ns = time.monotonic_ns()
                rclpy.spin_once(node, timeout_sec=0)
                with self._lock:
                    expired = (
                        self._valid_until_monotonic_ns > 0
                        and now_monotonic_ns >= self._valid_until_monotonic_ns
                    )
                    if expired:
                        self._latest_speed_mps = 0.0
                        self._latest_steering_normalized = 0.0
                        self._valid_until_monotonic_ns = 0
                        if not self._ttl_expired_for_latest:
                            self._ttl_expired_count += 1
                            self._ttl_expired_for_latest = True
                    speed_mps = self._latest_speed_mps
                    steering_normalized = self._latest_steering_normalized
                    disarm_requested = self._disarm_requested
                    control_confirmed = self._armed and self._offboard
                timestamp_us = int(node.get_clock().now().nanoseconds // 1_000)
                self._publish_setpoint(
                    messages,
                    offboard_publisher,
                    throttle_publisher,
                    steering_publisher,
                    timestamp_us,
                    speed_mps,
                    steering_normalized,
                )
                heartbeat_count += 1
                if (
                    heartbeat_count >= OFFBOARD_WARMUP_HEARTBEATS
                    and not control_confirmed
                    and now_monotonic - last_mode_request_monotonic >= 1.0
                ):
                    self._publish_vehicle_command(
                        messages,
                        command_publisher,
                        timestamp_us,
                        VEHICLE_COMMAND_DO_SET_MODE,
                        param1=1.0,
                        param2=6.0,
                    )
                    self._publish_vehicle_command(
                        messages,
                        command_publisher,
                        timestamp_us,
                        VEHICLE_COMMAND_COMPONENT_ARM_DISARM,
                        param1=1.0,
                    )
                    last_mode_request_monotonic = now_monotonic
                if disarm_requested:
                    self._publish_vehicle_command(
                        messages,
                        command_publisher,
                        timestamp_us,
                        VEHICLE_COMMAND_COMPONENT_ARM_DISARM,
                        param1=0.0,
                    )
                    with self._lock:
                        self._armed = False
                        self._offboard = False
                        self._disarm_requested = False
                    self._disarmed.set()
                self._stop.wait(self.heartbeat_seconds)
        except BaseException as exc:
            with self._lock:
                self._failure = exc
            self._started.set()
            self._control_ready.set()
            self._disarmed.set()
        finally:
            if node is not None:
                node.destroy_node()
            if rclpy is not None:
                with suppress(Exception):
                    rclpy.shutdown()

    @staticmethod
    def _publish_setpoint(
        messages: Any,
        offboard_publisher: Any,
        throttle_publisher: Any,
        steering_publisher: Any,
        timestamp_us: int,
        speed_mps: float,
        steering_normalized: float,
    ) -> None:
        offboard = messages.OffboardControlMode()
        offboard.timestamp = timestamp_us
        offboard.position = False
        offboard.velocity = False
        offboard.acceleration = False
        offboard.attitude = False
        offboard.body_rate = False
        offboard.thrust_and_torque = True
        offboard.direct_actuator = False
        offboard_publisher.publish(offboard)

        throttle = messages.RoverThrottleSetpoint()
        throttle.timestamp = timestamp_us
        throttle.throttle_body_x = float(speed_mps / STOCK_ROVER_MAX_THROTTLE_SPEED_MPS)
        throttle.throttle_body_y = math.nan
        throttle_publisher.publish(throttle)

        steering = messages.RoverSteeringSetpoint()
        steering.timestamp = timestamp_us
        steering.normalized_steering_setpoint = float(steering_normalized)
        steering_publisher.publish(steering)

    @staticmethod
    def _publish_vehicle_command(
        messages: Any,
        publisher: Any,
        timestamp_us: int,
        command: int,
        *,
        param1: float,
        param2: float = 0.0,
    ) -> None:
        message = messages.VehicleCommand()
        message.timestamp = timestamp_us
        message.param1 = param1
        message.param2 = param2
        message.command = command
        message.target_system = 1
        message.target_component = 1
        message.source_system = 1
        message.source_component = 1
        message.from_external = True
        publisher.publish(message)
