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

VERSION = "0.0.5"

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
    have_jenkins = True
except ImportError:
    have_jenkins = False

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


from urllib.parse import urljoin

IMAGES = {
    'stable': {
        'url': 'https://storage.puri.sm/librem5/images/',
        'u-boot': 'https://storage.puri.sm/librem5/u-boot/',
    }
}

JENKINS = 'https://build.puri.sm'
BOARD_TYPE = 'librem5r4'
VALID_PHONE_BOARD_TYPES = ['librem5r2', 'librem5r3', 'librem5r4']
VALID_DEVKIT_BOARD_TYPES = ['devkit']
DIST = 'current'
IMAGE = '{}.img'
IMAGE_VARIANT = 'luks'
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: download -f {uboot}
FB: ucmd mmc dev 0 1
FB: ucmd mmc write 0x43000000 0x42 0x1000
FB: ucmd mmc partconf 0 0 {bootpart} 0
FB: ucmd mmc dev 0 0
FB: ucmd setenv fastboot_dev mmc
FB: ucmd setenv mmcdev 0
FB: flash -raw2sparse all {image}
FB: reboot
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"
'''

# python3-urllib 2.x breaks python3-jenkins <1.8.1; it no longer accepts
# socket._GLOBAL_DEFAULT_TIMEOUT as the default timeout.  This combination of
# packages appears in Ubuntu 24.04 LTS and Debian Trixie (as of 2024-12-11).
# Work around it by specifying our own timeout value.
#
# python-jenkins commit: https://opendev.org/jjb/python-jenkins/commit/ba9f06e0c2381f6506a92e9c0553b2f475975a18
# python-jenkins bug: https://bugs.launchpad.net/python-jenkins/+bug/2018567
JENKINS_TIMEOUT = 60


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)},
                                timeout=10
                                )
            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,
                requests.exceptions.Timeout,
                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_jenkins(jobname, type, variant, dist):
    server = jenkins.Jenkins(JENKINS, timeout=JENKINS_TIMEOUT)
    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'] is None:
            continue
        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_image_stable(board, variant, dist):
    remote = IMAGES['stable']
    logging.info("Looking for {} {} {} image".format(board, variant, dist))
    found = None

    path = f"{dist}/latest/{board}/{variant}/"
    url = urljoin(remote['url'], f"{path}/artifact/{IMAGE.format(board)}.xz")
    try:
        resp = requests.head(url, timeout=10)
        if resp.ok:
            d = datetime.datetime.strptime(resp.headers['Last-Modified'],
                                           '%a, %d %b %Y %H:%M:%S %Z')
            return {
                'url': urljoin(remote['url'], path),
                'timestamp': d.timestamp() * 1000,
                'id': '"stable"',
                'description': f"Last stable {board} build",
            }
        resp.raise_for_status()

    except (requests.exceptions.ConnectionError,
            requests.exceptions.Timeout):
        logging.error("Failed to find image at {} - connection failed".format(url))
        found = None
    return found


def find_image(jobname, board, variant, dist, stable):
    if stable:
        return find_image_stable(board, variant, dist)
    else:
        return find_image_jenkins(jobname, board, variant, dist)


def find_uboot_stable(board):
    remote = IMAGES['stable']
    logging.info("Looking for {} u-boot".format(board))
    found = None

    path = "latest/"
    url = urljoin(remote['u-boot'], f"{path}/artifact/output/uboot-{board}/{UBOOT.format(board)}")
    try:
        resp = requests.head(url, timeout=10)
        if resp.ok:
            d = datetime.datetime.strptime(resp.headers['Last-Modified'],
                                           '%a, %d %b %Y %H:%M:%S %Z')
            return {
                'url': urljoin(remote['u-boot'], path),
                'timestamp': d.timestamp() * 1000,
                'id': '"stable"',
                'description': f"Last stable {board} build",
            }
        resp.raise_for_status()

    except (requests.exceptions.ConnectionError,
            requests.exceptions.Timeout):
        logging.error("Failed to find u-boot at {} - connection failed".format(url))
        found = None
    return found


def find_uboot_jenkins(jobname):
    server = jenkins.Jenkins(JENKINS, timeout=JENKINS_TIMEOUT)

    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 find_uboot(jobname, board, stable):
    if stable:
        return find_uboot_stable(board)
    else:
        return find_uboot_jenkins(jobname)


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, bootpart, mods):
    with open(target, 'w+') as f:
        f.write(UUU_SCRIPT_TMPL.format(image=os.path.basename(image),
                                       uboot=os.path.basename(uboot),
                                       bootpart=bootpart,
                                       mods=mods))


def wait_for_librem(librem_msg=None):
    found = False
    printed_msg = False

    if librem_msg is None:
        librem_msg = """
            Enter the flashing mode by holding volume-up button while turning the phone on.

            If it's not detected, follow these steps:
            - Ensure that the phone is powered off
            - Turn all Hardware-Kill-Switches off
            - Unplug the USB cable if connected
            - Remove battery
            - Hold volume-up button
            - Insert the USB-C cable (red light blinks, no green light)
            - Reinsert the battery (red and green lights constantly on, the script will continue)
            - Release volume-up button\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('--version', action='version', version=VERSION)
    parser.add_argument('--dir', type=str, default=None,
                        help='download files to DIR (instead of a temporary directory)')
    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('--skip-download', action='store_true', default=False,
                        help='use already downloaded files (specified with --dir) instead of downloading')
    parser.add_argument('--download-attempts', type=int, default=RETRIES,
                        help="maximum number of attempts to resume "
                        "image download. 0: unlimited, default is {}".format(RETRIES))

    group = parser.add_argument_group(title='image options')
    group.add_argument('--dist', type=str, default=DIST,
                       help="download an image for this distribution ( byzantium, crimson, dawn, ... ) "
                       "instead of the current release")
    group.add_argument('--variant', choices=['plain', 'luks'], default=IMAGE_VARIANT,
                       help="image variant to download ( plain, luks ), "
                       "default is '{}'".format(IMAGE_VARIANT))
    group.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.add_argument('--stable', action='store_true', default=True,
                       help="download the latest stable image (default, use --snapshot to override)")
    group.add_argument('--snapshot', action='store_true', default=False,
                       help="download the latest CI snapshot instead of \"blessed\" images "
                       "(takes precedence over --stable)")
    group.add_argument('--image', type=str,
                       help="flash the given uncompressed image file instead of downloading one")
    group.add_argument('--uboot', type=str,
                       help="use and flash the given bootloader file instead of downloading one")

    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.add_argument('--embedded-boot', action='store_true', default=False,
                       help="set the bootloader embedded in the flashed image as bootable "
                       "instead of using separate boot partition")

    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()

    if args.snapshot:
        args.stable = False
        if not have_jenkins:
            logging.error("Cannot load Jenkins\nTry running: sudo apt install python3-jenkins")
            return 1

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

    # Check available downloads upfront so it's less likely we fail
    # later:
    if not args.image:
        image_ref = find_image(args.image_job, board, args.variant, args.dist, args.stable)
        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

    uboot_board = board

    if not args.uboot:
        if args.snapshot and args.uboot_job:
            uboot_ref = find_uboot_jenkins(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("librem5_mainline" if uboot_board == "librem5" else uboot_board), uboot_board, args.stable)

        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"

    dirprefix = 'tmp_' + os.path.basename(sys.argv[0]) + '_'
    outdir = args.dir if args.dir is not None else tempfile.mkdtemp(prefix=dirprefix, 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))

        if args.skip_download:
            if args.dir is None:
                logging.error("Skipping download requires a directory to be specified with --dir")
                return 1
            if (not os.path.exists(image_target)) or (not os.path.exists(uboot_target)):
                logging.error(f"Requested to skip download but at least one of image ({image_target}) or uboot ({uboot_target}) is missing")
                return 1
            logging.info(f"Skipping download and using local image '{image_target}'")
        else:
            if args.image:
                if not os.path.exists(args.image):
                    logging.error(f"Given image {args.image} does not exist")
                    return 1
                os.symlink(os.path.abspath(args.image), image_target)
            else:
                download_image(urljoin(image_ref['url'], 'artifact/{}.xz').format(IMAGE.format(board)),
                               image_target, args.download_attempts)

            if args.uboot:
                if not os.path.exists(args.uboot):
                    logging.error(f"Given bootloader {args.uboot} does not exist")
                    return 1
                os.symlink(os.path.abspath(args.uboot), uboot_target)
            else:
                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, 0 if args.embedded_boot else 1, 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)
            print("Flashing complete.")
    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())