#!/usr/bin/env python
# vim:set ts=4 sw=4 et:
#
# Copyright 2014 Vladimir Ermakov.
#
# This file is part of the mavros package and subject to the license terms
# in the top-level LICENSE file of the mavros repository.
# https://github.com/mavlink/mavros/tree/master/LICENSE.md

from __future__ import print_function

import argparse
import threading

import rospy
import mavros
from mavros.utils import *
from mavros import mission as M

no_prettytable = False
try:
    from prettytable import PrettyTable
except ImportError:
    print("Waring: 'show' action disabled. install python-prettytable", file=sys.stderr)
    no_prettytable = True



def get_wp_file_io(args):
    return M.QGroundControlWP()


def _pull(args):
    try:
        ret = M.pull()
    except rospy.ServiceException as ex:
        fault(ex)

    if not ret.success:
        fault("Request failed. Check mavros logs")

    print_if(args.verbose, "Waypoints received:", ret.wp_received)
    return ret


def do_pull(args):
    _pull(args)


def do_show(args):
    str_bool = lambda x: 'Yes' if x else 'No'

    def str_frame(f):
        if M.FRAMES.has_key(f):
            return M.FRAMES[f] + ' ({})'.format(f)
        else:
            return 'UNK ({})'.format(f)

    def str_command(c):
        if M.NAV_CMDS.has_key(c):
            return M.NAV_CMDS[c] + ' ({})'.format(c)
        else:
            return 'UNK ({})'.format(c)

    done_evt = threading.Event()
    def _show_table(topic):
        pt = PrettyTable(('#', 'Curr', 'Auto',
                          'Frame', 'Command',
                          'P1', 'P2', 'P3', 'P4',
                          'X Lat', 'Y Long', 'Z Alt'))

        for seq, w in enumerate(topic.waypoints):
            pt.add_row((
                seq,
                str_bool(w.is_current),
                str_bool(w.autocontinue),
                str_frame(w.frame),
                str_command(w.command),
                w.param1,
                w.param2,
                w.param3,
                w.param4,
                w.x_lat,
                w.y_long,
                w.z_alt
            ))

        print(pt, file=sys.stdout)
        sys.stdout.flush()
        done_evt.set()

    if args.pull:
        _pull(args)

    # Waypoints topic is latched type, and it updates after pull
    sub = M.subscribe_waypoints(_show_table)
    if args.follow:
        rospy.spin()
    elif not done_evt.wait(30.0):
        fault("Something went wrong. Topic timed out.")


def do_load(args):
    wps = []
    wps_file = get_wp_file_io(args)
    with args.file:
        wps = [w for w in wps_file.read(args.file)]

    def _load_call(waypoint_list):
        try:
            ret = M.push(waypoints=waypoint_list)
        except rospy.ServiceException as ex:
            fault(ex)

        if not ret.success:
            fault("Request failed. Check mavros logs")

        print_if(args.verbose, "Waypoints transfered:", ret.wp_transfered)

    done_evt = threading.Event()
    def _fix_wp0(topic):
        if len(topic.waypoints) > 0:
            wps[0] = topic.waypoints[0]
            print_if(args.verbose, "HOME location: latitude:", wps[0].x_lat,
                     "longitude:", wps[0].y_long, "altitude:", wps[0].z_alt)
        else:
            print("Failed to get WP0! WP0 will be loaded from file.", file=sys.stderr)

        done_evt.set()

    if not args.preserve_home:
        _load_call(wps)
    else:
        # Note: _load_call() emit publish on this topic, so callback only changes
        # waypoint 0, and signal done event.
        sub = M.subscribe_waypoints(_fix_wp0)
        if not done_evt.wait(30.0):
            fault("Something went wrong. Topic timed out.")
        else:
            sub.unregister()
            _load_call(wps)


def do_dump(args):
    done_evt = threading.Event()
    def _write_file(topic):
        wps_file = get_wp_file_io(args)
        with args.file:
            wps_file.write(args.file, topic.waypoints)
        done_evt.set()

    # Waypoints topic is latched type, and it updates after pull
    _pull(args)
    sub = M.subscribe_waypoints(_write_file)
    if not done_evt.wait(30.0):
        fault("Something went wrong. Topic timed out.")


def do_clear(args):
    try:
        ret = M.clear()
    except rospy.ServiceException as ex:
        fault(ex)

    if not ret.success:
        fault("Request failed, Check mavros logs")
    else:
        print_if(args.verbose, "Waypoints cleared")


def do_set_current(args):
    try:
        ret = M.set_current(wp_seq=args.seq)
    except rospy.ServiceException as ex:
        fault(ex)

    if not ret.success:
        fault("Request failed, Check mavros logs")
    else:
        print_if(args.verbose, "Set current done.")


def do_goto(args):
    rospy.init_node("mavwp", anonymous=True)

    wp = M.Waypoint(
        frame = args.frame,
        command = args.command,
        param1 = args.param1,
        param2 = args.param2,
        param3 = args.param3,
        param4 = args.param4,
        x_lat = args.x_lat,
        y_long = args.y_long,
        z_alt = args.z_alt
    )

    try:
        ret = M.goto(waypoint=wp)
    except rospy.ServiceException as ex:
        fault(ex)

    if not ret.success:
        fault("Request failed. Check mavros logs")

    print_if(args.verbose, "Goto done.")


def main():
    parser = argparse.ArgumentParser(description="Commad line tool for manipulating mission on MAVLink device.")
    parser.add_argument('-n', '--mavros-ns', help="ROS node namespace", default="/mavros")
    parser.add_argument('-v', '--verbose', action='store_true', help="verbose output")
    subarg = parser.add_subparsers()

    if not no_prettytable:
        show_args = subarg.add_parser('show', help="Show waypoints")
        show_args.add_argument('-f', '--follow', action='store_true', help="watch and show new data")
        show_args.add_argument('-p', '--pull', action='store_true', help="pull waypoints from FCU before show")
        show_args.set_defaults(func=do_show)

    load_args = subarg.add_parser('load', help="load waypoints from file")
    load_args.set_defaults(func=do_load)
    load_args.add_argument('-p', '--preserve-home', action='store_true', help="Preserve home location (WP 0, APM only)")
    load_args.add_argument('file', type=argparse.FileType('rb'), help="input file (QGC/MP format)")

    pull_args = subarg.add_parser('pull', help="pull waypoints from FCU")
    pull_args.set_defaults(func=do_pull)

    dump_args = subarg.add_parser('dump', help="dump waypoints to file")
    dump_args.set_defaults(func=do_dump)
    dump_args.add_argument('file', type=argparse.FileType('wb'), help="output file (QGC format)")

    clear_args = subarg.add_parser('clear', help="clear waypoints on device")
    clear_args.set_defaults(func=do_clear)

    setcur_args = subarg.add_parser('setcur', help="set current waypoints on device")
    setcur_args.add_argument('seq', type=int, help="waypoint seq id")
    setcur_args.set_defaults(func=do_set_current)

    goto_args = subarg.add_parser('goto', help="send goto waypoint (APM only)")
    goto_args.add_argument('frame', type=int, help="frame code")
    goto_args.add_argument('command', type=int, help="command code")
    goto_args.add_argument('param1', type=float, help="param1")
    goto_args.add_argument('param2', type=float, help="param2")
    goto_args.add_argument('param3', type=float, help="param3")
    goto_args.add_argument('param4', type=float, help="param4")
    goto_args.add_argument('x_lat', type=float, help="X latitude")
    goto_args.add_argument('y_long', type=float, help="Y longitude")
    goto_args.add_argument('z_alt', type=float, help="Z altitude")
    goto_args.set_defaults(func=do_goto)

    args = parser.parse_args(rospy.myargv(argv=sys.argv)[1:])

    rospy.init_node("mavwp", anonymous=True)
    mavros.set_namespace(args.mavros_ns)
    M.setup_services()

    args.func(args)


if __name__ == '__main__':
    main()
