#!/usr/bin/env python3
import sys
import os
import time
import re
from select import select
from glob import glob
from subprocess import check_output, STDOUT, PIPE, Popen, TimeoutExpired

IMAGES_HOST = 'l_samenv@linse-c'
IMAGES_DIR = 'boxes/rpi_images'
IMAGES_PAT = '*.lz4'

extdisks = {}


def prt(text):
    print(time.strftime('%H:%M:%S'), text)


def cmd(command, verbose=False):
    if verbose:
        print(f'> {command}')
    return check_output(command.split(), stderr=STDOUT).decode('latin-1')


def find_raspi(kind):
    # collect info about disks
    dlist = '\n' + cmd('/usr/sbin/diskutil list')
    extdisks.clear()
    proposed = []
    noname = re.compile(r' +0: +\*\d+\.?\d* GB +disk')
    for dinfo in dlist.split('\n/dev/'):
        if not dinfo:
            continue
        disk, info = dinfo.split(maxsplit=1)
        if 'external' in info.split('\n')[0]:
            extdisks[disk] = info
            if kind == 'emmc':
                if 'Windows_FAT_32 bootfs' in info and 'Linux' in info:
                    proposed.append(disk)
                elif noname.search(info):
                    proposed.append(disk)
            elif kind == 'sd':
                if 'Apple' not in info:
                    proposed.append(disk)
    return proposed


def rpiboot():
    if os.path.exists('/Volumes/bootfs'):
        prt('unmount bootfs')
        os.system('diskutil unmount bootfs')
    lines = []
    with Popen(['rpiboot'], stdout=PIPE) as p:
        tmo = 1
        t = 0
        while select([p.stdout], [], [], tmo)[0]:
            rawline = p.stdout.readline()
            if not rawline:
                return tmo != 1
            line = rawline.decode('latin-1').strip()
            if lines is None:
                print(line)
            elif 'Waiting' in line:
                tmo = 2
                lines.append(line)
                t = time.time()
            elif tmo == 1:
                lines.append(line)
            else:
                print('\n'.join(lines))
                print(line)
                lines = None
        else:
            p.terminate()
            return False


def list_external():
    print('--- external disks:')
    for disk, info in extdisks.items():
       print(disk, info)


rescue = """
Remark: not all cables might work.
Example on Markus Mac:
- does not work: black USB C to USB C cable
- works: white USB A to USB C cable plugged into docking station.

Please unpower the box with compute module,
start this script again and put power immediately."""


def write_image(kind):
    if kind == 'eemc':
        print("\nmake sure microUSB 'SLAVE' is connected to this computer")
        print("and microUSB 'POWER' is connected to a USB power supply\n")
        prt('try to find CM (compute module on PoE device) ...')
    else:
        pass
    proposed = find_raspi(kind)
    if not proposed:
        if kind == 'emmc':
            # we may install rpiboot here if not available
            # git clone --depth=1 https://github.com/raspberrypi/usbboot
            # cd usbboot
            # make
            # ./rpiboot

            print('not found - put rpi in bootloader mode')
            if not rpiboot() and not rpiboot():
                print(rescue)
                return
            for _ in range(10):
                proposed = find_raspi(kind)
                if proposed:
                    break
                time.sleep(0.1)
            else:
                list_external()
                print('can not find CM')
                return
    list_external()
    if len(proposed) == 1:
        dev = proposed[0]
    elif len(proposed) > 1:
        if kind == 'emmc':
            print('several potential devices for CM:', proposed)
            return
        dev = input('select device to write to {" ".join(proposed)}: ')
    else:
        print('none of above disks seems to be the sd card')
        return
    images = check_output(['sudo', '-u', 'zolliker', 'ssh', IMAGES_HOST, f'cd {IMAGES_DIR} ; ls {IMAGES_PAT}']).decode('latin-1').split('\n')
    for file in sorted(images):
        print(file)
    print('')
    image = input(f'select image to write to {dev} [{file}]: ') or file
    prt(f'unmount {dev} ...')
    print('you may be prompted for your account password for sudo')
    os.system(f'sudo diskutil unmountDisk /dev/{dev}')
    print('')
    prt('get, uncompress and write... (takes about 15 min)')
    os.system(f'ssh {IMAGES_HOST} "dd if={IMAGES_DIR}/{image} bs=512k" | lz4 -d | sudo dd of=/dev/{dev} bs=512k')
    prt('unmount bootfs')
    for _ in range(3):
        if os.path.exists('/Volumes/bootfs'):
           prt('unmount bootfs')
           os.system('diskutil unmount bootfs')
           break
        time.sleep(1.0)
    prt('done')
    os.system(f'diskutil unmountDisk /dev/{dev}')

if len(sys.argv) == 2:
    write_image(sys.argv[1])
else:
    print("""
Usage:

> write_to emmc   # write to compute module using PoE board"
> write_to sd     # write to sd card

""")