#!/usr/bin/env python2
# -*- coding: utf-8 -*-
# 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 os
import rospy
import mavros
from mavros.utils import *
from mavros.nuttx_crc32 import *
from mavros import ftp

no_progressbar = False
try:
    import progressbar as pbar
except ImportError:
    print("Prigressbar disabled. install python-progressbar", file=sys.stderr)
    no_progressbar = True

# optimized transfer size for FTP message payload
# XXX: bug in ftp.cpp cause a doubling request of last package.
# -1 fixes that.
FTP_PAGE_SIZE = 239 * 18 - 1
FTP_PWD_FILE = '/tmp/.mavftp_pwd'


def _resolve_path(path = None):
    """
    Resolve FTP path using PWD file
    """
    if os.path.exists(FTP_PWD_FILE):
        with open(FTP_PWD_FILE, 'r') as fd:
            pwd = fd.readline()
    else:
        # default home location is root directory
        pwd = os.environ.get('MAVFTP_HOME', '/')

    if not path:
        return os.path.normpath(pwd)    # no path - PWD location
    elif path.startswith('/'):
        return os.path.normpath(path)   # absolute path
    else:
        return os.path.normpath(os.path.join(pwd, path))


class ProgressBar:
    """
    Wrapper class for hiding file transfer brogressbar construction
    """
    def __init__(self, quiet, operation, maxval):
        if no_progressbar or quiet or maxval == 0:
            print_if(maxval == 0, "Can't show progressbar for unknown file size", file=sys.stderr)
            self.pbar = None
            return

        self.pbar = pbar.ProgressBar(
            widgets=[operation, pbar.Percentage(), ' ', pbar.Bar(), ' ', pbar.ETA(), ' ', pbar.FileTransferSpeed()],
            maxval=maxval).start()

    def update(self, value):
        if self.pbar:
            self.pbar.update(value)

    def __enter__(self):
        if self.pbar:
            self.pbar.start()

        return self

    def __exit__(self, type, value, traceback):
        if self.pbar:
            self.pbar.finish()


def do_change_directory(args):
    # TODO: check that path is exist
    path = args.path or os.environ.get('MAVFTP_HOME')
    if args.path and not path.startswith('/'):
        path = _resolve_path(path)
    if path and not path.startswith('/'):
        fault("Path is not absolute:", path)

    if path:
        with open(FTP_PWD_FILE, 'w') as fd:
            fd.write(os.path.normpath(path))
    else:
        if os.path.exists(FTP_PWD_FILE):
            os.unlink(FTP_PWD_FILE)


def do_list(args):
    args.path = _resolve_path(args.path)
    for ent in ftp.listdir(args.path):
        n = ent.name
        if ent.type == ftp.FileEntry.TYPE_DIRECTORY:
            n += '/'
        else:
            n += '\t{}'.format(ent.size)

        print(n)


def do_cat(args):
    args.no_progressbar = True
    args.no_verify = False
    args.file = sys.stdout
    do_download(args)


def do_remove(args):
    args.path = _resolve_path(args.path)
    ftp.unlink(args.path)


def do_reset(args):
    ftp.reset_server()


def do_mkdir(args):
    args.path = _resolve_path(args.path)
    ftp.mkdir(args.path)


def do_rmdir(args):
    args.path = _resolve_path(args.path)
    ftp.rmdir(args.path)


def do_download(args):
    local_crc = 0
    args.path = _resolve_path(args.path)

    if not args.file:
        # if file argument is not set, use $PWD/basename
        args.file = open(os.path.basename(args.path), 'wb')

    print_if(args.verbose, "Downloading from", args.path, "to", args.file.name, file=sys.stderr)

    with args.file as to_fd, \
            ftp.open(args.path, 'r') as from_fd, \
            ProgressBar(args.no_progressbar, "Downloading: ", from_fd.size) as bar:
        while True:
            buf = from_fd.read(FTP_PAGE_SIZE)
            if len(buf) == 0:
                break

            local_crc = nuttx_crc32(buf, local_crc)
            to_fd.write(buf)
            bar.update(from_fd.tell())

    if not args.no_verify:
        print_if(args.verbose, "Verifying...", file=sys.stderr)
        remote_crc = ftp.checksum(args.path)
        if local_crc != remote_crc:
            fault("Verification failed: 0x{local_crc:08x} != 0x{remote_crc:08x}".format(**locals()))


def do_upload(args):
    mode = 'cw' if args.no_overwrite else 'w'
    local_crc = 0

    if args.path:
        args.path = _resolve_path(args.path)
    else:
        args.path = _resolve_path(os.path.basename(args.file.name))

    print_if(args.verbose, "Uploading from", args.file.name, "to", args.path, file=sys.stderr)

    # for stdin it is 0
    from_size = os.fstat(args.file.fileno()).st_size

    with args.file as from_fd, \
            ftp.open(args.path, mode) as to_fd, \
            ProgressBar(args.no_progressbar, "Uploading: ", from_size) as bar:
        while True:
            buf = from_fd.read(FTP_PAGE_SIZE)
            if len(buf) == 0:
                break

            local_crc = nuttx_crc32(buf, local_crc)
            to_fd.write(buf)
            bar.update(to_fd.tell())

    if not args.no_verify:
        print_if(args.verbose, "Verifying...", file=sys.stderr)
        remote_crc = ftp.checksum(args.path)
        if local_crc != remote_crc:
            fault("Verification failed: 0x{local_crc:08x} != 0x{remote_crc:08x}".format(**locals()))


def do_verify(args):
    local_crc = 0

    if args.path:
        args.path = _resolve_path(args.path)
    else:
        args.path = _resolve_path(os.path.basename(args.file.name))

    print_if(args.verbose, "Verifying", args.file.name, "and", args.path, file=sys.stderr)

    with args.file as fd:
        while True:
            buf = fd.read(4096 * 32)   # use 128k block for CRC32 calculation
            if len(buf) == 0:
                break

            local_crc = nuttx_crc32(buf, local_crc)

    remote_crc = ftp.checksum(args.path)

    print_if(args.verbose, "CRC32 for local and remote files:")
    print_if(args.verbose, "0x{local_crc:08x}  {args.file.name}".format(**locals()))
    print_if(args.verbose, "0x{remote_crc:08x}  {args.path}".format(**locals()))

    if local_crc != remote_crc:
        print("{args.file.name}: FAULT".format(**locals()))
        sys.exit(1)
    else:
        print("{args.file.name}: OK".format(**locals()))


def main():
    parser = argparse.ArgumentParser(description="File manipulation tool for MAVLink-FTP.")
    parser.add_argument('-n', '--mavros-ns', help="ROS node namespace", default=mavros.DEFAULT_NAMESPACE)
    parser.add_argument('-v', '--verbose', action='store_true', help="verbose output")
    subarg = parser.add_subparsers()

    # argparse from python2 don't support subparser aliases
    # there exist a hack, but i don't want it now.

    cd_args = subarg.add_parser('cd', help="change directory")
    cd_args.set_defaults(func=do_change_directory)
    cd_args.add_argument('path', type=str, nargs='?', help="directory path")

    list_args = subarg.add_parser('list', help="list files and dirs")
    list_args.set_defaults(func=do_list)
    list_args.add_argument('path', type=str, nargs='?', help="directory path")

    cat_args = subarg.add_parser('cat', help="cat file")
    cat_args.set_defaults(func=do_cat)
    cat_args.add_argument('path', type=str, help="file path")

    remove_args = subarg.add_parser('remove', help="remove file")
    remove_args.set_defaults(func=do_remove)
    remove_args.add_argument('path', type=str, help="file path")

    mkdir_args = subarg.add_parser('mkdir', help="create direcotory")
    mkdir_args.set_defaults(func=do_mkdir)
    # mkdir_args.add_argument('-p', action='store_true', help="dir path")
    mkdir_args.add_argument('path', type=str, help="dir path")

    rmdir_args = subarg.add_parser('rmdir', help="remove directory")
    rmdir_args.set_defaults(func=do_rmdir)
    rmdir_args.add_argument('path', type=str, help="dir path")

    download_args = subarg.add_parser('download', help="download file")
    download_args.set_defaults(func=do_download)
    download_args.add_argument('path', type=str, help="file to send")
    download_args.add_argument('file', type=argparse.FileType('wb'), nargs='?', help="save path")
    download_args.add_argument('-q', '--no-progressbar', action="store_true", help="do not show progressbar")
    download_args.add_argument('--no-verify', action="store_true", help="do not perform verify step")

    upload_args = subarg.add_parser('upload', help="upload file")
    upload_args.set_defaults(func=do_upload)
    upload_args.add_argument('file', type=argparse.FileType('rb'), help="file to send")
    upload_args.add_argument('path', type=str, nargs='?', help="save path")
    upload_args.add_argument('-n', '--no-overwrite', action="store_true", help="do not overwrite existing file")
    upload_args.add_argument('-q', '--no-progressbar', action="store_true", help="do not show progressbar")
    upload_args.add_argument('--no-verify', action="store_true", help="do not perform verify step")

    verify_args = subarg.add_parser('verify', help="verify files")
    verify_args.set_defaults(func=do_verify)
    verify_args.add_argument('file', type=argparse.FileType('rb'), help="local file")
    verify_args.add_argument('path', type=str, nargs='?', help="remote file")

    reset_args = subarg.add_parser('reset', help="reset")
    reset_args.set_defaults(func=do_reset)

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

    rospy.init_node("mavftp", anonymous=True)
    mavros.set_namespace(args.mavros_ns)
    args.func(args)


if __name__ == '__main__':
    main()
