#!/usr/bin/env python
#
# Copyright 2021 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""A wrapper script for flashing combined bootloader image"""

import argparse
import os
import subprocess
import sys
import time

MODE_UPDATE = "update"
MODE_FASTBOOT = "fastboot"


def parse_args():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("image", help="path to the combined bootloader image")
    parser.add_argument("mode",
                        choices=[MODE_UPDATE, MODE_FASTBOOT],
                        help="flash mode")
    parser.add_argument("--update_tool",
                        "-u",
                        help="path to the amlogic usb update executable")
    return parser.parse_args()


def update(update_tool, image):
    subprocess.check_call(
        [update_tool, "write", image, "0xfffa0000", "0x10000"])
    subprocess.check_call([update_tool, "run", "0xfffa0000"])
    subprocess.check_call([update_tool, "bl2_boot", image])
    print("Waiting TPL plug-in...")
    time.sleep(7)
    subprocess.check_call([update_tool, "partition", "bootloader", image])


def fastboot(image):
    subprocess.check_call(["fastboot", "flash", "bootloader", image])


def main():
    args = parse_args()
    update_tool = args.update_tool if args.update_tool else "update"
    if args.mode == MODE_UPDATE:
        update(update_tool, args.image)
    elif args.mode == MODE_FASTBOOT:
        fastboot(args.image)
    return 0


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