#!/usr/bin/env python3
"""Small user client for the local Kaptan host-SITL process."""

from __future__ import annotations

import argparse
import math
import secrets
import socket
import struct
import sys
import time


CONTROL = ("127.0.0.1", 14662)
TRUTH = ("127.0.0.1", 14663)
# The simulator lease identifies its one local user by UDP source address.
# Keep this source stable across separate CLI invocations so a completed move
# can be followed by the next command without a second, competing user.
USER_CONTROL = ("127.0.0.1", 14665)
USER_CLOCK = ("127.0.0.1", 14664)
VERSION = 1
STATUS_QUERY = struct.Struct("<4sB3x")
STATUS = struct.Struct("<4sB3xQH2x")
GOAL = struct.Struct("<4sB3xI3f")
GOAL_KEEPALIVE = struct.Struct("<4sB3xI")
GOAL_STATUS = struct.Struct("<4sB3xIB3x")
CLOCK_PACKET = struct.Struct("<4sB3xQ")
# KST2 carries pose/velocity, launch-relative attitude, and applied motor
# fractions. It is a test-only stream; the Rust controller never binds TRUTH.
TRUTH_PACKET = struct.Struct("<4sB3xQI14f")

ESTIMATOR_READY = 1 << 0
IMU_FRESH = 1 << 1
LOCAL_AIDING_FRESH = 1 << 2
OUTPUT_PERMITTED = 1 << 3
ACTUATOR_TRANSPORT = 1 << 4
READY = ESTIMATOR_READY | IMU_FRESH | LOCAL_AIDING_FRESH | ACTUATOR_TRANSPORT
GOAL_KEEPALIVE_MICROSECONDS = 100_000
GOAL_REJECTED = 0
GOAL_ACTIVE = 1
GOAL_SETTLED = 2
GOAL_LANDED = 3
# Independent truth-recorder assertion for the single, flat ground plane
# declared by the supplied world. This is not a physical contact detector and
# must not be carried into a flight target.
GROUND_PLANE_TOLERANCE_METRES = 0.10


def arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    subcommands = parser.add_subparsers(dest="command", required=True)
    ready = subcommands.add_parser("ready", help="wait for estimator and simulated sensors")
    ready.add_argument("--timeout", type=float, default=20.0)
    move = subcommands.add_parser("move", help="command absolute local north/east/up metres")
    move.add_argument("north", type=float)
    move.add_argument("east", type=float)
    move.add_argument("up", type=float)
    move.add_argument("--wait", action="store_true", help="wait for the settled result")
    move.add_argument(
        "--verify-truth",
        action="store_true",
        help="verify noise-free Gazebo truth on the test-only recorder port after settling",
    )
    move.add_argument("--truth-position-tolerance", type=float, default=0.75)
    move.add_argument("--truth-velocity-tolerance", type=float, default=0.50)
    move.add_argument("--truth-attitude-tolerance-rad", type=float, default=0.35)
    move.add_argument(
        "--truth-hold-seconds",
        type=float,
        default=0.75,
        help="require every independent truth sample to remain within tolerance for this interval",
    )
    # A distant, bounded goal is intentionally slow: the default must not turn
    # a valid long trajectory into a spurious client-side test failure.
    move.add_argument("--timeout", type=float, default=180.0)
    land = subcommands.add_parser(
        "land",
        help="descend to an explicit local landing point and revoke virtual motor authority",
    )
    land.add_argument("north", type=float, help="absolute local north landing coordinate in metres")
    land.add_argument("east", type=float, help="absolute local east landing coordinate in metres")
    land.add_argument("--wait", action="store_true", help="wait for the terminal landed result")
    land.add_argument(
        "--verify-truth",
        action="store_true",
        help="verify independent Gazebo ground truth and zero applied motor fractions after landing",
    )
    land.add_argument("--truth-position-tolerance", type=float, default=0.75)
    land.add_argument("--truth-velocity-tolerance", type=float, default=0.25)
    land.add_argument("--truth-attitude-tolerance-rad", type=float, default=0.35)
    land.add_argument("--truth-hold-seconds", type=float, default=0.75)
    land.add_argument("--timeout", type=float, default=180.0)
    return parser.parse_args()


def recv(socket_: socket.socket, deadline: float) -> bytes | None:
    remaining = deadline - time.monotonic()
    if remaining <= 0:
        return None
    socket_.settimeout(remaining)
    try:
        payload, sender = socket_.recvfrom(256)
    except TimeoutError:
        return None
    return payload if sender[0] == "127.0.0.1" else None


def latest_simulation_time(
    clock_socket: socket.socket, previous: int | None
) -> int | None:
    """Drains Gazebo's test-user clock without trusting host elapsed time."""
    latest = previous
    while True:
        try:
            payload, sender = clock_socket.recvfrom(64)
        except BlockingIOError:
            return latest
        if sender[0] != "127.0.0.1" or len(payload) != CLOCK_PACKET.size:
            continue
        magic, version, simulation_time = CLOCK_PACKET.unpack(payload)
        if magic != b"KSC1" or version != VERSION:
            continue
        if latest is not None and simulation_time < latest:
            raise RuntimeError("Gazebo simulation clock reset; restart the command in the new epoch")
        latest = simulation_time


def status(socket_: socket.socket, deadline: float) -> int | None:
    socket_.sendto(STATUS_QUERY.pack(b"KSQ1", VERSION), CONTROL)
    payload = recv(socket_, deadline)
    if payload is None or len(payload) != STATUS.size:
        return None
    magic, version, _, flags = STATUS.unpack(payload)
    if magic != b"KSS1" or version != VERSION:
        return None
    return flags


def wait_ready(socket_: socket.socket, timeout: float) -> None:
    deadline = time.monotonic() + timeout
    latest = None
    while time.monotonic() < deadline:
        latest = status(socket_, min(deadline, time.monotonic() + 0.5))
        if latest is not None and latest & READY == READY:
            return
    if latest is None:
        raise RuntimeError("host SITL is not running; run docker compose -f tools/sitl/compose.yaml up -d --build")
    names = (
        (ESTIMATOR_READY, "estimator initialization"),
        (IMU_FRESH, "fresh simulator IMU"),
        (LOCAL_AIDING_FRESH, "accepted simulator local-state aiding"),
        (ACTUATOR_TRANSPORT, "virtual actuator transport"),
    )
    missing = [name for bit, name in names if latest & bit != bit]
    raise RuntimeError(f"SITL is waiting for: {', '.join(missing)}")


def truth_sample_within_tolerance(
    position_error: float,
    velocity: float,
    attitude_error: float,
    motors_stopped: bool,
    ground_plane_error: float,
    position_tolerance: float,
    velocity_tolerance: float,
    attitude_tolerance: float,
    require_stopped_motors: bool,
    require_ground_plane: bool,
) -> bool:
    """Checks one independent truth sample against its declared contract."""
    return not (
        position_error > position_tolerance
        or velocity > velocity_tolerance
        or attitude_error > attitude_tolerance
        or (require_stopped_motors and not motors_stopped)
        or (
            require_ground_plane
            and ground_plane_error > GROUND_PLANE_TOLERANCE_METRES
        )
    )


def verify_truth(
    north: float,
    east: float,
    up: float,
    timeout: float,
    position_tolerance: float,
    velocity_tolerance: float,
    attitude_tolerance: float,
    hold_seconds: float,
    require_stopped_motors: bool,
    require_ground_plane: bool,
    control_socket: socket.socket | None = None,
    clock_socket: socket.socket | None = None,
    keepalive: bytes | None = None,
) -> None:
    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as truth_socket:
        truth_socket.bind(TRUTH)
        acquisition_deadline = time.monotonic() + timeout
        hold_deadline: int | None = None
        samples = 0
        worst_position_error = 0.0
        worst_velocity = 0.0
        worst_attitude_error = 0.0
        worst_ground_plane_error = 0.0
        worst_motor_fraction = 0.0
        last_out_of_tolerance: str | None = None
        simulation_time: int | None = None
        next_keepalive_time: int | None = None
        last_truth_time: int | None = None
        while time.monotonic() < acquisition_deadline:
            if clock_socket is not None:
                simulation_time = latest_simulation_time(clock_socket, simulation_time)
            if (
                keepalive is not None
                and control_socket is not None
                and simulation_time is not None
                and (next_keepalive_time is None or simulation_time >= next_keepalive_time)
            ):
                control_socket.sendto(keepalive, CONTROL)
                next_keepalive_time = simulation_time + GOAL_KEEPALIVE_MICROSECONDS
            payload = recv(truth_socket, min(acquisition_deadline, time.monotonic() + 0.01))
            if payload is None or len(payload) != TRUTH_PACKET.size:
                continue
            (
                magic,
                version,
                truth_time,
                _,
                north_m,
                east_m,
                down_m,
                north_v,
                east_v,
                down_v,
                attitude_w,
                attitude_x,
                attitude_y,
                attitude_z,
                motor_1,
                motor_2,
                motor_3,
                motor_4,
            ) = TRUTH_PACKET.unpack(payload)
            if magic != b"KST2" or version != VERSION:
                continue
            if last_truth_time is not None and truth_time < last_truth_time:
                raise RuntimeError("Gazebo truth time regressed")
            last_truth_time = truth_time
            position_error = math.sqrt(
                (north_m - north) ** 2 + (east_m - east) ** 2 + (down_m + up) ** 2
            )
            # `up` is converted to local-NED down above. A terminal land
            # request explicitly sets it to zero, so this distinguishes the
            # declared ground plane from a low-velocity armed hover.
            ground_plane_error = abs(down_m + up)
            velocity = math.sqrt(north_v**2 + east_v**2 + down_v**2)
            quaternion_norm = math.sqrt(
                attitude_w**2 + attitude_x**2 + attitude_y**2 + attitude_z**2
            )
            if not math.isfinite(quaternion_norm) or abs(quaternion_norm - 1.0) > 1e-3:
                raise RuntimeError("Gazebo truth contained a non-unit attitude quaternion")
            attitude_error = 2.0 * math.acos(min(1.0, abs(attitude_w) / quaternion_norm))
            motors = (motor_1, motor_2, motor_3, motor_4)
            maximum_motor_fraction = max(abs(motor) for motor in motors)
            motors_stopped = all(abs(motor) <= 0.01 for motor in motors)
            if not truth_sample_within_tolerance(
                position_error,
                velocity,
                attitude_error,
                motors_stopped,
                ground_plane_error,
                position_tolerance,
                velocity_tolerance,
                attitude_tolerance,
                require_stopped_motors,
                require_ground_plane,
            ):
                # A lifecycle result says the controller's own settle rule was
                # met. The independent plant may still be damping. Require one
                # uninterrupted truth hold, rather than rejecting the first
                # transient sample after that result.
                hold_deadline = None
                samples = 0
                last_out_of_tolerance = (
                    "position_error="
                    f"{position_error:.3f} m, velocity={velocity:.3f} m/s, "
                    f"attitude_error={attitude_error:.3f} rad, "
                    f"ground_plane_error={ground_plane_error:.3f} m, motors={motors}"
                )
                continue
            if hold_deadline is None:
                # Use the truth record's own simulated timestamp. A separately
                # drained clock record can be ahead of buffered truth packets,
                # so it cannot prove an uninterrupted physical hold.
                hold_deadline = truth_time + int(hold_seconds * 1_000_000)
                worst_position_error = position_error
                worst_velocity = velocity
                worst_attitude_error = attitude_error
                worst_ground_plane_error = ground_plane_error
                worst_motor_fraction = maximum_motor_fraction
            else:
                worst_position_error = max(worst_position_error, position_error)
                worst_velocity = max(worst_velocity, velocity)
                worst_attitude_error = max(worst_attitude_error, attitude_error)
                worst_ground_plane_error = max(
                    worst_ground_plane_error, ground_plane_error
                )
                worst_motor_fraction = max(worst_motor_fraction, maximum_motor_fraction)
            samples += 1
            if truth_time >= hold_deadline:
                summary = (
                    "truth verified: "
                    f"samples={samples}, position_error<={worst_position_error:.3f} m, "
                    f"velocity<={worst_velocity:.3f} m/s, "
                    f"attitude_error<={worst_attitude_error:.3f} rad"
                )
                if require_ground_plane:
                    summary += f", ground_plane_error<={worst_ground_plane_error:.3f} m"
                if require_stopped_motors:
                    summary += f", motor_fraction<={worst_motor_fraction:.3f}"
                print(summary)
                return
    if last_out_of_tolerance is not None:
        raise RuntimeError(
            "Gazebo truth never held within tolerance before timeout: "
            + last_out_of_tolerance
        )
    raise RuntimeError("did not receive a Gazebo truth record")


def move(
    socket_: socket.socket,
    clock_socket: socket.socket,
    north: float,
    east: float,
    up: float,
    timeout: float,
    wait: bool,
    verify: bool,
    position_tolerance: float,
    velocity_tolerance: float,
    attitude_tolerance: float,
    hold_seconds: float,
) -> None:
    if not all(math.isfinite(value) for value in (north, east, up)):
        raise ValueError("goal coordinates must be finite")
    wait_ready(socket_, timeout)
    goal_id = secrets.randbelow(0xFFFF_FFFF) + 1
    goal = GOAL.pack(b"KSG1", VERSION, goal_id, north, east, up)
    command_goal(
        socket_,
        clock_socket,
        goal,
        goal_id,
        north,
        east,
        up,
        timeout,
        wait,
        verify,
        position_tolerance,
        velocity_tolerance,
        attitude_tolerance,
        hold_seconds,
        GOAL_SETTLED,
        False,
        "move",
    )


def land(
    socket_: socket.socket,
    clock_socket: socket.socket,
    north: float,
    east: float,
    timeout: float,
    wait: bool,
    verify: bool,
    position_tolerance: float,
    velocity_tolerance: float,
    attitude_tolerance: float,
    hold_seconds: float,
) -> None:
    if not all(math.isfinite(value) for value in (north, east)):
        raise ValueError("landing coordinates must be finite")
    wait_ready(socket_, timeout)
    goal_id = secrets.randbelow(0xFFFF_FFFF) + 1
    landing = GOAL.pack(b"KSD1", VERSION, goal_id, north, east, 0.0)
    command_goal(
        socket_,
        clock_socket,
        landing,
        goal_id,
        north,
        east,
        0.0,
        timeout,
        wait,
        verify,
        position_tolerance,
        velocity_tolerance,
        attitude_tolerance,
        hold_seconds,
        GOAL_LANDED,
        True,
        "land",
    )


def command_goal(
    socket_: socket.socket,
    clock_socket: socket.socket,
    goal: bytes,
    goal_id: int,
    north: float,
    east: float,
    up: float,
    timeout: float,
    wait: bool,
    verify: bool,
    position_tolerance: float,
    velocity_tolerance: float,
    attitude_tolerance: float,
    hold_seconds: float,
    terminal_lifecycle: int,
    require_stopped_motors: bool,
    action: str,
) -> None:
    keepalive = GOAL_KEEPALIVE.pack(b"KSK1", VERSION, goal_id)
    deadline = time.monotonic() + timeout
    simulation_time: int | None = None
    next_send_time: int | None = None
    active = False
    while time.monotonic() < deadline:
        simulation_time = latest_simulation_time(clock_socket, simulation_time)
        if next_send_time is None or (
            simulation_time is not None and simulation_time >= next_send_time
        ):
            socket_.sendto(keepalive if active else goal, CONTROL)
            if simulation_time is not None:
                next_send_time = simulation_time + GOAL_KEEPALIVE_MICROSECONDS
        payload = recv(socket_, min(deadline, time.monotonic() + 0.01))
        if payload is None or len(payload) != GOAL_STATUS.size:
            continue
        magic, version, received_id, lifecycle = GOAL_STATUS.unpack(payload)
        if magic != b"KSL1" or version != VERSION or received_id != goal_id:
            continue
        if lifecycle == GOAL_REJECTED:
            raise RuntimeError(f"{action} was rejected because SITL was not estimator-ready")
        if lifecycle == GOAL_ACTIVE:
            active = True
            print(f"{action} active: north={north:g} m, east={east:g} m, up={up:g} m")
            break
    if not active:
        raise RuntimeError("goal was not activated")
    if not wait:
        return
    while time.monotonic() < deadline:
        simulation_time = latest_simulation_time(clock_socket, simulation_time)
        if simulation_time is not None and (
            next_send_time is None or simulation_time >= next_send_time
        ):
            socket_.sendto(keepalive, CONTROL)
            next_send_time = simulation_time + GOAL_KEEPALIVE_MICROSECONDS
        payload = recv(socket_, min(deadline, time.monotonic() + 0.01))
        if payload is None or len(payload) != GOAL_STATUS.size:
            continue
        magic, version, received_id, lifecycle = GOAL_STATUS.unpack(payload)
        if (
            magic == b"KSL1"
            and version == VERSION
            and received_id == goal_id
            and lifecycle == terminal_lifecycle
        ):
            print("landed" if terminal_lifecycle == GOAL_LANDED else "settled")
            if verify:
                verify_truth(
                    north,
                    east,
                    up,
                    min(5.0, timeout),
                    position_tolerance,
                    velocity_tolerance,
                    attitude_tolerance,
                    hold_seconds,
                    require_stopped_motors,
                    terminal_lifecycle == GOAL_LANDED,
                    None if require_stopped_motors else socket_,
                    clock_socket,
                    None if require_stopped_motors else keepalive,
                )
            return
    raise RuntimeError(f"{action} did not reach its terminal state before timeout")


def main() -> int:
    values = arguments()
    if values.timeout <= 0:
        raise ValueError("timeout must be positive")
    if values.command in ("move", "land"):
        tolerances = (
            values.truth_position_tolerance <= 0
            or values.truth_velocity_tolerance <= 0
            or values.truth_attitude_tolerance_rad <= 0
            or values.truth_hold_seconds <= 0
            or not math.isfinite(values.truth_position_tolerance)
            or not math.isfinite(values.truth_velocity_tolerance)
            or not math.isfinite(values.truth_attitude_tolerance_rad)
            or not math.isfinite(values.truth_hold_seconds)
        )
        if tolerances:
            raise ValueError("truth tolerances and hold duration must be finite and positive")
    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as client:
        client.bind(USER_CONTROL)
        if values.command == "ready":
            wait_ready(client, values.timeout)
            print("ready")
        elif values.command == "move":
            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as clock:
                clock.bind(USER_CLOCK)
                clock.setblocking(False)
                move(
                    client,
                    clock,
                    values.north,
                    values.east,
                    values.up,
                    values.timeout,
                    values.wait,
                    values.verify_truth,
                    values.truth_position_tolerance,
                    values.truth_velocity_tolerance,
                    values.truth_attitude_tolerance_rad,
                    values.truth_hold_seconds,
                )
        else:
            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as clock:
                clock.bind(USER_CLOCK)
                clock.setblocking(False)
                land(
                    client,
                    clock,
                    values.north,
                    values.east,
                    values.timeout,
                    values.wait,
                    values.verify_truth,
                    values.truth_position_tolerance,
                    values.truth_velocity_tolerance,
                    values.truth_attitude_tolerance_rad,
                    values.truth_hold_seconds,
                )
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (OSError, RuntimeError, ValueError, struct.error) as error:
        print(f"kaptan-sim: {error}", file=sys.stderr)
        raise SystemExit(2) from error