#!/usr/bin/env python

from argparse import ArgumentParser
import logging

import pymoku
import pymoku.version
from pymoku.finder import Finder
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor


log = logging.getLogger('moku')

parser = ArgumentParser()
subparsers = parser.add_subparsers(title="action", dest='action',
                                   description="Action to take")
subparsers.required = True

parser.add_argument('-v', action='store_true', help="Verbose logging.")


def listmokus(args):
    mokus = Finder(ipv6=False).find_all(timeout=2.0)
    sort_key = args.sort
    assert sort_key in ['serial', 'name'], f'Invalid sort key {sort_key}'

    print("{: <7} {: <20} {: >6} {: >6} {: >15}".format('HW', 'Name', 'Serial', 'FW', 'IP'))
    print("-" * (7 + 20 + 6 + 6 + 15 + 5))

    for m in sorted(mokus, key=lambda m: getattr(m, sort_key)):
        print(f"{m.hwver: <7} {m.name:<20} {m.serial:>06d} {m.fwver: 6d} {m.ip_addr:>15}")


parser_list = subparsers.add_parser(
    'list', help="List Moku:Labs on the network.")
parser_list.add_argument('--ipv6', '-6', help="IPv6", action='store_true')
parser_list.add_argument('--sort', '-s', help="Sort By: serial, name", default='serial')
parser_list.set_defaults(func=listmokus)


def update_firmware(args):
    m = pymoku.connect(args.ip_addr, attach=False)
    m.update_firmware(fw=args.fw)
    log.info(f"Upgrade started, please wait for your {type(m).__name__} to reboot")


parser_update = subparsers.add_parser('update', help="Upgrade Moku firmware")
parser_update.add_argument('ip_addr',
                           default=None,
                           help="IP Address of the Moku to connect upgrade")
parser_update.add_argument('fw',
                           type=str,
                           help="Upload specific fw file")
parser_update.set_defaults(func=update_firmware)


def upload_files(args):
    def _upload(files):
        m = pymoku.connect(args.ip_addr, attach=False)
        fs = m._fileserv
        for path in files:
            log.info(f'uploading {path}')
            mountpoint = 'p'
            if path.suffix == '.bar':
                mountpoint = 'b' if not args.persist else 'n'

            fs.send(mp=mountpoint, fname=path.name, data=path.read_bytes())

    threads = min(len(args.files), 3)
    if threads:
        with ThreadPoolExecutor(max_workers=threads) as pool:
            for i in range(threads):
                pool.submit(_upload, args.files[i::threads])

    if args.reload:
        m = pymoku.connect(args.ip_addr, attach=False)
        log.info('reloading bitstream cache')
        m._control.reload_bitstreams()


parser_upload = subparsers.add_parser('upload', help="Upload files to Moku")
parser_upload.add_argument('ip_addr',
                           default=None,
                           help="IP Address of the Moku to connect upgrade")
parser_upload.add_argument('-p', '--persist', action='store_true', help="Upload bitstream to persistent storage")
parser_upload.add_argument('-r', '--reload', action='store_true', help="Reload the bitstream cache after uploading")
parser_upload.add_argument('files', nargs='*', type=Path)
parser_upload.set_defaults(func=upload_files)


def clear_files(args):
    m = pymoku.connect(args.ip_addr, attach=False)
    fs = m._fileserv

    def _clear_partition(partition):
        files = fs.list(partition)
        for file in files:
            log.debug(f'removing {partition}:{file}')
            fs.remove(partition, file)

    if args.partition in ['persist', 'all']:
        _clear_partition('p')
        _clear_partition('n')

    if args.partition in ['ssd', 'all']:
        _clear_partition('s')

    if args.partition in ['tmp', 'all']:
        _clear_partition('b')


parser_clear = subparsers.add_parser('clear', help="Delete files from Moku")
parser_clear.add_argument('ip_addr',
                          default=None,
                          help="IP Address of the Moku to connect upgrade")
parser_clear.add_argument('partition', choices=['persist', 'media', 'ssd', 'tmp', 'all'])
parser_clear.set_defaults(func=clear_files)


def deploy(args):
    m = pymoku.connect(args.ip_addr, attach=False)

    if args.platform:
        log.debug(f'deploy {args.platform}')
        m.modify_hardware(platform=dict(ID=args.platform, force=args.force))

    if args.instr:
        log.debug(f'deploy {args.instr} to slot {args.slot:d}')
        m.modify_hardware({f'instr{args.slot:d}': dict(ID=args.instr, force=args.force)})


parser_deploy = subparsers.add_parser('deploy', help="Deploy bitstreams on the Moku")
parser_deploy.add_argument('ip_addr',
                           default=None,
                           help="IP Address of the Moku to connect upgrade")
parser_deploy.add_argument('-f', '--force', action='store_true', help="Deploy even if the current bitstream matches")
parser_deploy.add_argument('-p', '--platform', type=str, help="The id of the platform. ie. \"pymoku:platform_4\"")
parser_deploy.add_argument('-i', '--instr', type=str, help="The id of the instrument. ie. \"pymoku:oscilloscope\"")
parser_deploy.add_argument('-s', '--slot', type=int, default=1, help="The slot index to deploy an instrument to")
parser_deploy.set_defaults(func=deploy)


def main():
    args = parser.parse_args()
    lvl = [logging.INFO, logging.DEBUG][args.v]
    fmt = ['%(message)s', '%(asctime)8s.%(msecs)03d %(levelname)7s: %(message)s'][args.v]
    logging.basicConfig(format=fmt, datefmt='%H:%M:%S', level=lvl)
    log.debug("PyMoku %s" % pymoku.__version__)
    args.func(args)


if __name__ == '__main__':
    main()
