#!/usr/bin/env python3
"""Small MAVLink command-line debugger for a Kaptan flight controller."""

import argparse
import os
import sys
import time


def parse_arguments():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("serial_device", help="serial device connected to the flight controller")
    parser.add_argument("command", choices=("arm", "disarm", "read"))
    parser.add_argument("--baud", type=int, default=115200, help="serial baud rate (default: 115200)")
    parser.add_argument("--system-id", type=int, default=1, help="target MAVLink system ID (default: 1)")
    parser.add_argument(
        "--component-id", type=int, default=1, help="target MAVLink component ID (default: 1)"
    )
    return parser.parse_args()


def print_messages(connection, timeout):
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        message = connection.recv_match(blocking=True, timeout=deadline - time.monotonic())
        if message is None:
            return None
        print(message)
        yield message


def send_arm_command(connection, mavutil, arguments, armed):
    for _ in range(3):
        connection.mav.command_long_send(
            arguments.system_id,
            arguments.component_id,
            mavutil.mavlink.MAV_CMD_COMPONENT_ARM_DISARM,
            0,
            1.0 if armed else 0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
        )
        for message in print_messages(connection, timeout=1.0):
            if (
                message.get_type() == "COMMAND_ACK"
                and message.command == mavutil.mavlink.MAV_CMD_COMPONENT_ARM_DISARM
            ):
                return
    action = "arm" if armed else "disarm"
    raise SystemExit(f"timed out waiting for {action} acknowledgement")


def read_messages(connection):
    try:
        while True:
            message = connection.recv_match(blocking=True, timeout=1.0)
            if message is not None:
                print(message)
    except KeyboardInterrupt:
        pass


def main():
    arguments = parse_arguments()
    # Kaptan currently accepts MAVLink 2 frames only.
    os.environ["MAVLINK20"] = "1"
    from pymavlink import mavutil

    connection = mavutil.mavlink_connection(arguments.serial_device, baud=arguments.baud)
    if arguments.command == "read":
        read_messages(connection)
    else:
        send_arm_command(connection, mavutil, arguments, armed=arguments.command == "arm")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(130)