import argparse
import os.path
import threading
import time
import subprocess

import colorama
from colorama import Fore
import serial


def printc(line, mode):
    if isinstance(line, bytes):
        try:
            line = line.decode()
        except:
            pass

    print(Fore.BLUE + mode.ljust(10) + Fore.RESET + line.rstrip())


def printe(event, line):
    print(Fore.YELLOW + event.ljust(10) + Fore.RESET + line.rstrip())


class Flasher(threading.Thread):
    def __init__(self, source, target, control):
        threading.Thread.__init__(self)
        self.source = source
        self.target = target
        self.control = control

    def run(self):
        # Start dd to flash the image to the device
        printe('flasher', f'Write {self.source} to {self.target}')

        while not os.path.exists(self.target):
            time.sleep(0.2)

        printe('flasher', f'{self.target} appeared')

        subprocess.run([
            'sudo', 'dd', f'if={self.source}', f'of={self.target}',
            'bs=1M', 'oflag=direct,sync', 'status=progress'
        ])

        # Run a filesystem sync to make sure
        printe('flasher', f'Running sync')
        subprocess.run(['sudo', 'sync'])

        # Press the power button
        printe('flasher', f'Flashing complete, shutting down')
        self.control.press_power(True)
        time.sleep(8)
        self.control.press_power(False)
        time.sleep(1)
        printe('flasher', f'Powering on')
        self.control.press_power(True)


class ButtonController:
    def __init__(self, device):
        self.device = device
        self.port = serial.Serial(self.device, 115200)

        # Disable all gpios
        self.port.write(b'pbr')

        self.power_key = False
        self.bootloader_key = False
        self.power = False

    def press_power(self, state):
        self.power_key = state
        self.port.write(b'B' if state else b'b')

    def press_bootloader(self, state):
        self.bootloader_key = state
        self.port.write(b'R' if state else b'r')

    def set_power(self, state):
        self.power = state
        self.port.write(b'P' if state else b'p')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('uartdevice')
    parser.add_argument('controldevice')
    parser.add_argument('diskdevice')
    parser.add_argument('image')
    parser.add_argument('--baudrate', '-b', default=115200, type=int)
    args = parser.parse_args()

    colorama.init()

    boot_to_ums = True

    uart = serial.Serial(args.uartdevice, args.baudrate, timeout=5)
    control = ButtonController(args.controldevice)

    time.sleep(1)

    # Hard reset device and launch tow-boot UMS mode.
    printc('Hard reset (15 seconds)', 'sys')
    control.press_power(True)
    time.sleep(15)
    uart.flushInput()
    control.press_power(False)
    printc('Hard reset completed, starting device...', 'sys')
    time.sleep(1)
    control.press_power(True)
    mode = 'unknown'

    while True:
        oldmode = mode
        line = uart.readline()
        if line.startswith(b'U-Boot SPL '):
            if oldmode == 'spl':
                printe("error", "Boot loop detected")
            mode = 'spl'
            if control.power_key:
                control.press_power(False)
                if boot_to_ums:
                    control.press_bootloader(True)
        elif line.startswith(b'Tow-Boot '):
            mode = 'towboot'
        elif mode == 'towboot' and line.startswith(b'Starting kernel'):
            mode = 'kernel'
        elif mode == 'towboot' and line.startswith(b'Allwinner mUSB OTG'):
            mode = 'ums'
            # Don't hold the bootloader key on the next boot to get into the new kernel
            boot_to_ums = False
            control.press_bootloader(False)

            # Start writing the image in a second thread
            thread = Flasher(args.image, args.diskdevice, control)
            thread.daemon = True
            thread.start()
        elif line.startswith(b'[    0.000000] Booting Linux on physical CPU'):
            mode = 'kernel'

        if mode != oldmode:
            if mode == 'spl':
                printe("state", "Reset detected")
            else:
                printe("state", "Moved from [" + Fore.BLUE + oldmode + Fore.RESET + "] to [" +
                       Fore.BLUE + mode + Fore.RESET + "]")

        if mode == 'towboot' and b'to enter the boot menu' in line:
            printe("state", "Boot menu prompt")

        if mode == 'ums' and len(line) < 4:
            # Don't try to print the spinner animation in UMS mode
            continue

        printc(line, mode=mode)


if __name__ == '__main__':
    main()