#!/usr/bin/env python3

import subprocess
import time
import fourletterphat as flp

#--
drive_blacklist = {"/boot","/"}
#--

def disp_string(string):
    for letter in range(0,4):
        try:
            flp.set_digit(letter, string[letter].upper())
        except IndexError:
            pass
    flp.show()

def get_mountedlist():
    return [item[item.find("/"):] for item in subprocess.check_output(
            ["/bin/bash", "-c", "lsblk"]).decode("utf-8").split("\n") if "/" in item]

def get_mounted_devs():
    cmd = ["/bin/bash", "-c", "lsblk"]
    output = subprocess.check_output(cmd).decode("utf-8").split("\n")
    res = set()
    for item in output:
        if item.startswith("sd"):
            # print(f"Found /dev/{item.split(' ')[0]}")
            res.add(f"/dev/{item.split(' ')[0]}")
    return res

def flash(dev):
    flash_cmd = ["/bin/dd", "if=/home/pi/Downloads/2019-07-10-raspbian-buster-lite.img", f"of={dev}"]
    output = subprocess.check_output(flash_cmd).decode("utf-8")
    print(output)
    return True

disp_string("GO")

done = set()
while True:
    try:
        devs = get_mounted_devs()
        done = done.intersection(devs)
        for dev in devs:
            if dev not in done:
                print(f"Flashing {dev}")
                disp_string(dev.split("/")[2])
                if flash(dev):
                    print(f"Flashed {dev}")
                    disp_string("DONE")
                    time.sleep(.5)
                    done.add(dev)
        time.sleep(1)
    except BaseException as err:
        disp_string("err")