import time
import pymavlink.mavutil as utility
import pymavlink.dialects.v20.all as dialect
import argparse
import pika


def main():


    ## KOLEJKA

    connection_parameters = pika.ConnectionParameters('localhost')

    connection = pika.BlockingConnection(connection_parameters)

    channel = connection.channel()

    channel.queue_declare(queue="letterbox")
    channel.basic_publish(exchange="", routing_key="letterbox", body="@Connected to query".encode())

    parser = argparse.ArgumentParser()

    parser.add_argument('--comm', dest='comm',
                        help='np. --comm udpin:127.0.0.1:14560 1 (0 -> GUIDED, 1 -> STABILIZE, 2 -> LAND, 3 -> RTL, 4-> MANUAL)',
                        type=str, nargs=2)
    args = parser.parse_args()
    comm_str = args.comm[0]
    print(comm_str)
    MODE_FLAG = args.comm[1]
    # connect to vehicle
    print("@Connecting")
    channel.basic_publish(exchange="", routing_key="letterbox", body="@Connecting".encode())
    vehicle = utility.mavlink_connection(device=comm_str)
    # wait for a heartbeat
    vehicle.wait_heartbeat()

    # inform user
    msg="Connected to system:", vehicle.target_system, ", component:", vehicle.target_component
    print(msg)
    channel.basic_publish(exchange="", routing_key="letterbox", body=str(msg).encode())

    FLIGHT_MODE = ""
    if MODE_FLAG == "0":
        FLIGHT_MODE = "GUIDED"
    elif MODE_FLAG == "1":
        FLIGHT_MODE = "STABILIZE"
    elif MODE_FLAG == "2":
        FLIGHT_MODE = "LAND"
    elif MODE_FLAG == "3":
        FLIGHT_MODE = "RTL"
    elif MODE_FLAG == "4":
        FLIGHT_MODE = "MANUAL"
    # get supported flight modes
    flight_modes = vehicle.mode_mapping()

    # check the desired flight mode is supported
    if FLIGHT_MODE not in flight_modes.keys():
        # inform user that desired flight mode is not supported by the vehicle
        print(FLIGHT_MODE, "is not supported")

        # create change mode message
    set_mode_message = dialect.MAVLink_command_long_message(
            target_system=vehicle.target_system,
            target_component=vehicle.target_component,
            command=dialect.MAV_CMD_DO_SET_MODE,
            confirmation=0,
            param1=dialect.MAV_MODE_FLAG_CUSTOM_MODE_ENABLED,
            param2=flight_modes[FLIGHT_MODE],
            param3=0,
            param4=0,
            param5=0,
            param6=0,
            param7=0
        )

        # catch HEARTBEAT message
    message = vehicle.recv_match(type=dialect.MAVLink_heartbeat_message.msgname, blocking=True)

        # convert this message to dictionary
    message = message.to_dict()

        # get mode id
    mode_id = message["custom_mode"]

        # get mode name
    flight_mode_names = list(flight_modes.keys())
    flight_mode_ids = list(flight_modes.values())
    flight_mode_index = flight_mode_ids.index(mode_id)
    flight_mode_name = flight_mode_names[flight_mode_index]

        # print mode name
    print("Mode name before:", flight_mode_name)

        # change flight mode
    vehicle.mav.send(set_mode_message)
        # do below always
    while True:

            # catch COMMAND_ACK message
            message = vehicle.recv_match(type=dialect.MAVLink_command_ack_message.msgname, blocking=True)

            # convert this message to dictionary
            message = message.to_dict()

            # check is the COMMAND_ACK is for DO_SET_MODE
            if message["command"] == dialect.MAV_CMD_DO_SET_MODE:

                # check the command is accepted or not
                if message["result"] == dialect.MAV_RESULT_ACCEPTED:

                    # inform the user
                    msg="Changing mode to", FLIGHT_MODE, "accepted from the vehicle"
                    print(msg)
                    channel.basic_publish(exchange="", routing_key="letterbox", body=str(msg).encode())


                # not accepted
                else:

                    # inform the user
                    msg="Changing mode to", FLIGHT_MODE, "failed"
                    print(msg)
                    channel.basic_publish(exchange="", routing_key="letterbox", body=str(msg).encode())

                # break the loop
                    break

        # catch HEARTBEAT message
            message = vehicle.recv_match(type=dialect.MAVLink_heartbeat_message.msgname, blocking=True)

        # convert this message to dictionary
            message = message.to_dict()

        # get mode id
            mode_id = message["custom_mode"]

        # get mode name
            flight_mode_names = list(flight_modes.keys())
            flight_mode_ids = list(flight_modes.values())
            flight_mode_index = flight_mode_ids.index(mode_id)
            flight_mode_name = flight_mode_names[flight_mode_index]

        # print mode name
            msg="Mode name after:", flight_mode_name
            print(msg)
            channel.basic_publish(exchange="", routing_key="letterbox", body=str(msg).encode())
            connection.close()

            # exit the code
            exit(1)



if __name__ == '__main__':
    main()