#!/usr/bin/env python3

import argparse
import logging
import os
import re
import shutil
import subprocess
import sys
import tempfile

try:
    import coloredlogs
    have_colored_logs = True
except ImportError:
    have_colored_logs = False


BOARD_TYPE = 'librem5r3'
IMAGE = '{}.img'
UBOOT = 'u-boot-{}.imx'
UUU_SCRIPT = 'flash_{}.lst'
UUU_SCRIPT_TMPL = '''uuu_version 1.0.1

{mods}
SDP: boot -f {uboot}
# This command will be run when use SPL
SDPU: delay 1000
SDPU: write -f {uboot} -offset 0x57c00
SDPU: jump
SDPV: delay 1000
SDPV: write -f {uboot} -skipspl
SDPV: jump
# This command will be run when ROM support stream mode
SDPS: boot -f {uboot}
SDPU: delay 1000
FB: ucmd setenv fastboot_dev mmc
FB: ucmd setenv mmcdev 0
FB: flash -raw2sparse all {image}
FB: Done
'''


def write_uuu_script(target, image, uboot, mods):
    with open(target, 'w+') as f:
        f.write(UUU_SCRIPT_TMPL.format(image=os.path.basename(image),
                                       uboot=os.path.basename(uboot),
                                       mods=mods))


def flash_image(uuu_target, debug):
    if debug:
        subprocess.check_call(['uuu', '-v', uuu_target])
    else:
        subprocess.check_call(['uuu', uuu_target])


def main():
    uuu_mods = ''
    parser = argparse.ArgumentParser(description='Flash a librem5 phone')
    parser.add_argument('--dir', type=str, default=None,
                        help='Look in dir for image and u-boot')
    parser.add_argument('--board', type=str, default=BOARD_TYPE,
                        help='Name of the image to flash ( devkit, librem5r2, librem5r3 )')

    group = parser.add_argument_group(title='Testing and debugging options')
    group.add_argument('--debug', action="store_true", default=False,
                       help='Enable debug output')
    group.add_argument('--uboot', type=str,
                       help='u-boot version to flash')
    args = parser.parse_args()

    level = logging.DEBUG if args.debug else logging.INFO
    if have_colored_logs:
        coloredlogs.install(level=level, fmt='%(asctime)s %(levelname)s %(message)s')
    else:
        logging.basicConfig(level=level, format='%(asctime)s %(levelname)s %(message)s')

    # uboot builds don't carry board revisions (yet?)
    if args.uboot:
        uboot_board = args.uboot
    else:
        if re.match('librem5r[0-9]$', args.board):
            uboot_board = args.board[:-2]

    imagedir = args.dir if args.dir is not None else '.'
    outdir = tempfile.mkdtemp(prefix='flash_', dir='.')

    try:
        logging.info("loading from {}".format(imagedir))
        if args.dir == outdir:
            os.makedirs(args.dir, exist_ok=True)

        image_target = os.path.join(imagedir, IMAGE.format(args.board))
        uboot_target = os.path.join(imagedir, UBOOT.format(uboot_board))
        uuu_target = os.path.join(outdir, UUU_SCRIPT.format(args.board))

        if not os.path.isfile(image_target) and not os.path.isfile(uboot_target):
            raise FileNotFoundError

        if not os.path.isabs(image_target):
            image_target = os.path.join('..', image_target)
        if not os.path.isabs(uboot_target):
            uboot_target = os.path.join('..', uboot_target)

        os.symlink(image_target, os.path.join(outdir, IMAGE.format(args.board)))
        os.symlink(uboot_target, os.path.join(outdir, UBOOT.format(uboot_board)))

        write_uuu_script(uuu_target, image_target, uboot_target, uuu_mods)
        flash_image(uuu_target, args.debug)
    except KeyboardInterrupt:
        logging.error("CTRL-C pressed.")
        return 1
    finally:
        if args.dir != outdir:
            logging.info("Cleaning up.")
            shutil.rmtree(outdir)

    return 0


if __name__ == '__main__':
    sys.exit(main())
