#!/usr/bin/env python3

import threading
import loader
from queue import Queue
import subprocess
import time
import fourletterphat as flp
import os
import tempfile
from shutil import copyfile

#--
q = Queue()
t = loader.bgload(3, .1)
IMAGE = "image.img"
PWD = os.getcwd()
#--

def disp_string(string):
    for letter in range(0,3) if len(string) > 3 else range(0, len(string)):
        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):
    loader.q.put(dev)
    try:
        flash_cmd = ["/bin/dd", f"if={IMAGE}", f"of={dev}"]
        
        print("flashing")
        subprocess.run(flash_cmd)
        print("flashed")
        
        with tempfile.TemporaryDirectory() as tempdir:
            mount_cmd = ["mount",/dev/f"{dev}2", tempdir]
            umount_cmd = ["umount", tempdir]

            print("mounting")
            subprocess.run(mount_cmd)
            print(f"changing dir to {tempdir}")
            os.chdir(tempdir)
            print("making ssh dir")
            os.mkdir("ssh")
            print(f"changing dir to {PWD}")
            os.chdir(PWD)
            print(f"copying wpa_supplicant to {tempdir}/")
            copyfile(f"wpa_supplicant.conf", f"{tempdir}/wpa_supplicant.conf")
            print("unmounting")
            subprocess.run(umount_cmd)

        return True

    except BaseException as err:
        print(err)
        disp_string('ERR')

    finally:
        loader.q.get()

if __name__=="__main__":

    while True:
        if os.path.isfile(IMAGE):
            break
        else:
            disp_string("IMG")
            time.sleep(1)

    disp_string("GO")

    done = set()
    while True:
        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("GO")
                    time.sleep(.5)
                    done.add(dev)
        time.sleep(1)