#!/usr/bin/env python3
#
# SPDX-License-Identifier: GPL-3.0-or-later
# Copyright (C) 2018-2021 Purism SPC

import argparse
import datetime
import hashlib
import itertools
import logging
import lzma
import os
import re
import requests
import shutil
import subprocess
import sys
import tempfile
import time
import tqdm
import yaml

try:
    import jenkins
except ImportError:
    logging.error("Cannot load Jenkins\nTry running: sudo apt install python3-jenkins")
    exit(1)

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


from urllib.parse import urljoin

JENKINS = 'https://arm01.puri.sm'
BOARD_TYPE = 'librem5r4'
VALID_PHONE_BOARD_TYPES = ['librem5r2', 'librem5r3', 'librem5r4']
VALID_DEVKIT_BOARD_TYPES = ['devkit']
DIST = 'byzantium'
IMAGE = '{}.img'
IMAGE_VARIANT = 'plain'
META_YAML = 'files/meta.yml'
IMAGE_JOB_NAME = 'Images/Image Build'
UBOOT = 'u-boot-{}.imx'
UBOOT_JOB_NAME = 'u-boot_builds/uboot_{}_build'
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_buffer 0x43000000
FB: ucmd setenv fastboot_dev mmc
FB: ucmd setenv mmcdev 0
FB: flash -raw2sparse all {image}
FB: Done
'''
BLOCK_SIZE = 8192
RETRIES = 10

UDEV_RULES = '''
SUBSYSTEM!="usb", GOTO="librem5_devkit_rules_end"
# Devkit USB flash
ATTR{idVendor}=="1fc9", ATTR{idProduct}=="012b", GROUP+="plugdev", TAG+="uaccess"
ATTR{idVendor}=="0525", ATTR{idProduct}=="a4a5", GROUP+="plugdev", TAG+="uaccess"
ATTR{idVendor}=="0525", ATTR{idProduct}=="b4a4", GROUP+="plugdev", TAG+="uaccess"
ATTR{idVendor}=="316d", ATTR{idProduct}=="4c05", GROUP+="plugdev", TAG+="uaccess"
LABEL="librem5_devkit_rules_end"
'''


class VerifyImageException(Exception):
    pass


class PrematureEndException(Exception):
    pass


def verify_image(image, meta):
    m = hashlib.sha256()
    size = int(meta['image']['size'])
    hexdigest = meta['image']['sha256sum']

    filesize = os.path.getsize(image)
    if filesize != size:
        raise VerifyImageException(
            "Image file \"{}\" size {} does not match {}".format(
                os.path.basename(image), filesize, size))

    logging.info("Calculating sha256sum of {}".format(image))
    bar = tqdm.tqdm(total=size,
                    desc='Checking',
                    leave=False)
    with open(image, 'rb') as f:
        while True:
            data = f.read(BLOCK_SIZE)
            if data:
                m.update(data)
                bar.update(n=len(data))
            else:
                break
    bar.close()
    if m.hexdigest() != hexdigest:
        raise VerifyImageException("Checksum of image {} "
                                   "does not match {}".format(m.hexdigest(), hexdigest))


def resuming_stream(url, expected_size, max_attempts):
    position = 0

    if max_attempts < 1:
        retries = itertools.count()
    else:
        retries = range(max_attempts)

    for i in retries:
        try:
            resp = requests.get(url,
                                stream=True,
                                headers={'Range': 'bytes={}-'.format(position)}
                                )
            resp.raise_for_status()

            if resp.status_code != requests.codes.partial_content:
                position = 0
            logging.debug('Proceeding from {} bytes'.format(position))

            for data in resp.iter_content(BLOCK_SIZE):
                position += len(data)
                yield data

            if position < expected_size:
                raise PrematureEndException()
            return
        except (requests.exceptions.ConnectionError, PrematureEndException):
            if i == max_attempts - 1:
                logging.error("Max connection errors reached, aborting")
                raise
            logging.info("Connection error, retrying")
            time.sleep(5)


def stream_file(url, attempts):
    resp = requests.head(url, stream=True)
    resp.raise_for_status()
    ts = int(resp.headers.get('content-length', 0))
    return resuming_stream(url, ts, attempts), ts


def needs_download(target, meta):
    if not os.path.exists(target):
        return True

    try:
        verify_image(target, meta)
    except VerifyImageException:
        return True

    return False


def download_image(url, target, attempts):
    decomp = lzma.LZMADecompressor()

    # We expect metadata to be right next to the image
    meta_yml_url = "{}/{}".format(url.rsplit('/', 1)[0], META_YAML)
    resp = requests.get(meta_yml_url)
    resp.raise_for_status()
    meta = yaml.safe_load(resp.text)
    uncompressed_size = int(meta['image']['size'])
    logging.debug("Image size is %d", uncompressed_size)

    if not needs_download(target, meta):
        logging.info("Image already up to date - no download needed.")
        return

    logging.info("Downloading image from {}".format(url))
    stream, ts = stream_file(url, attempts)
    download_bar = tqdm.tqdm(total=ts,
                             desc='Download',
                             leave=False)
    decompress_bar = tqdm.tqdm(total=uncompressed_size,
                               desc='Decompr.',
                               leave=False)
    with open(target, 'wb+') as f:
        for data in stream:
            if data:
                out = decomp.decompress(data)
                decompress_bar.update(len(out))
                f.write(out)
            download_bar.update(n=len(data))
    download_bar.close()
    decompress_bar.close()
    verify_image(target, meta)


def find_image(jobname, type, variant, dist):
    server = jenkins.Jenkins(JENKINS)
    logging.info("Looking for {} {} {} image".format(type, variant, dist))
    try:
        info = server.get_job_info(jobname)
    except jenkins.NotFoundException:
        logging.error("Job %s not found", jobname)
        return None
    for build in info['builds']:
        resp = requests.get(build['url'] + '/api/json')
        resp.raise_for_status()
        json = resp.json()
        if (json['description'].startswith(variant + ' ' + type) and
                dist in json['description'] and
                json['result'] == 'SUCCESS'):
            found = json
            break
    else:
        found = None
    return found


def find_uboot(jobname):
    server = jenkins.Jenkins(JENKINS)

    try:
        info = server.get_job_info(jobname)
    except jenkins.NotFoundException:
        logging.error("Job %s not found", jobname)
        return None
    for build in info['builds']:
        resp = requests.get(build['url'] + '/api/json')
        resp.raise_for_status()
        json = resp.json()
        if (json['result'] == 'SUCCESS'):
            found = json
            break
    else:
        found = None
    return found


def download_uboot(url, target):
    logging.info("Downloading uboot from {}".format(url))
    resp = requests.get(url, stream=True)
    resp.raise_for_status()
    ts = int(resp.headers.get('content-length', 0))
    bar = tqdm.tqdm(total=ts, leave=False)
    with open(target, 'wb+') as f:
        for data in resp.iter_content(BLOCK_SIZE):
            if data:
                f.write(data)
            bar.update(n=len(data))


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 wait_for_librem(librem_msg="default"):
    found = False
    printed_msg = False

    if librem_msg == "default":
        librem_msg = """
            Turn all Hardware-Kill-Switches off
            Remove battery
            Hold volume-up
            Insert the USB-c cable: (red light blinks, no green light)
            Reinsert the battery: (red light is constantly on, the script will continue)\n\nSearching..."""

    while not found:
        returned_text = subprocess.check_output("lsusb", shell=True, universal_newlines=True)
        for line in returned_text.split("\n"):
            if "NXP Semiconductors" in line:
                found = True
        if not found and not printed_msg:
            print(librem_msg)
            printed_msg = True
        time.sleep(1)


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


def get_board(path):
    """
    Get board type from script name.

    In case of doubt return the default board.

    >>> get_board("/a/b/librem5-flash-image")
    'librem5r4'
    >>> get_board("/a/b/librem5-devkit-flash-image")
    'devkit'
    >>> get_board("/a/b/librem5r2-flash-image")
    'librem5r2'
    >>> get_board("/a/b/librem5r3-flash-image")
    'librem5r3'
    >>> get_board("/a/b/librem5r4-flash-image")
    'librem5r4'
    >>> get_board("/a/b/schwobel")
    'librem5r4'
    """
    name = os.path.split(path)[-1]
    parts = name.split('-')
    if len(parts) < 3 or len(parts) > 4:
        return BOARD_TYPE

    if parts[-2:] != ['flash', 'image']:
        return BOARD_TYPE

    parts = parts[:-2]
    if len(parts) == 1:
        if parts[0] in VALID_PHONE_BOARD_TYPES:
            return parts[0]
        else:
            return BOARD_TYPE

    if parts[0] == 'librem5' and parts[1] in VALID_DEVKIT_BOARD_TYPES:
        return parts[1]

    # People do odd things so use default in this case
    return BOARD_TYPE


def setup_udev():
    rules = "/etc/udev/rules.d/70-librem5-flash-image.rules"
    if os.path.exists(rules):
        logging.info(f"Udev rules '{rules}' already present")
        return 0

    try:
        with open(rules, 'w+') as f:
            f.write(UDEV_RULES)
    except PermissionError as e:
        logging.error(f"{e} - are you root?")
        return 1

    subprocess.check_call(['udevadm', 'control', '--reload'])
    logging.info(f"Udev rules '{rules}' updated. You can now flash without root permissions.")
    return 0


def main():
    uuu_mods = ''
    parser = argparse.ArgumentParser(description='Flash a Librem 5 or Librem 5 Devkit.')
    parser.add_argument('--dir', type=str, default=None,
                        help='Download files to dir (instead of a temporary directory)')
    parser.add_argument('--dist', type=str, default=DIST,
                        help="Download an image for this distribution, default is '{}'".format(DIST))
    parser.add_argument('--skip-cleanup', action='store_true', default=False,
                        help='Skip temporary directory cleanup')
    parser.add_argument('--skip-flash', action='store_true', default=False,
                        help='Do all the preparations but don\'t flash')
    parser.add_argument('--download-attempts', type=int, default=RETRIES,
                        help="Maximum number of attempts to resume "
                        "image download. 0: unlimited, default is {}".format(RETRIES))
    parser.add_argument('--variant', choices=['plain', 'luks'], default=IMAGE_VARIANT,
                        help="Image variant to download ( plain, luks ), "
                        "default is '{}'".format(IMAGE_VARIANT))
    parser.add_argument('--board', choices=VALID_DEVKIT_BOARD_TYPES + VALID_PHONE_BOARD_TYPES, default=None,
                        help="Type of the board to download ( devkit, librem5r2, librem5r3, librem5r4 ) "
                        "default is 'auto'")

    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('--image-job', type=str, default=IMAGE_JOB_NAME,
                       help='Jenkins job to download the image from')
    group.add_argument('--uboot-job', type=str,
                       help='Jenkins job to download the uboot from')

    group = parser.add_argument_group(title='Initial setup options')
    group.add_argument('--udev', action="store_true", default=False,
                       help='Setup udev rules to flash without root permissions')
    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')

    if args.udev:
        return setup_udev()

    board = args.board if args.board else get_board(sys.argv[0])

    # Check available downloads upfront so it's less likely we fail
    # later:
    image_ref = find_image(args.image_job, board, args.variant, args.dist)
    if image_ref:
        image_ref['ts'] = datetime.datetime.fromtimestamp(image_ref['timestamp'] / 1000).strftime('%c')
        logging.info("Found disk image Build {id} '{description}' from {ts}".format(**image_ref))
    else:
        logging.error("No matching image found")
        return 1

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

    if uboot_ref:
        uboot_ref['ts'] = datetime.datetime.fromtimestamp(uboot_ref['timestamp'] / 1000).strftime('%c')
        logging.info("Found uboot Build {id} from {ts}".format(**uboot_ref))
    else:
        logging.error("No matching uboot found")
        return 1

    if board in ["librem5r3", 'librem5r4']:
        uuu_mods = "CFG: SDP: -chip MX8MQ -compatible MX8MQ -vid 0x316d -pid 0x4c05"

    outdir = args.dir if args.dir is not None else tempfile.mkdtemp(prefix='devkit_image_', dir='.')
    try:
        logging.info("Downloading to {}".format(outdir))
        if args.dir == outdir:
            os.makedirs(args.dir, exist_ok=True)

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

        download_image(urljoin(image_ref['url'], 'artifact/{}.xz').format(IMAGE.format(board)),
                       image_target, args.download_attempts)
        download_uboot(urljoin(uboot_ref['url'],
                       'artifact/output/uboot-{}/{}'.format(uboot_board, UBOOT.format(uboot_board))),
                       uboot_target)
        write_uuu_script(uuu_target, image_target, uboot_target, uuu_mods)
        if not args.skip_flash:
            if board == "devkit":
                wait_for_librem(librem_msg="Please connect your devkit now.\nSearching...")
            else:
                wait_for_librem()
            flash_image(uuu_target, args.debug)
    except VerifyImageException as e:
        logging.error(e)
        return 1
    except KeyboardInterrupt:
        logging.error("CTRL-C pressed.")
        return 1
    finally:
        if args.dir != outdir and not args.skip_cleanup:
            logging.info("Cleaning up.")
            shutil.rmtree(outdir)

    return 0


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