#!/usr/bin/env python3
"""Install (or update) the Wes World app on a Pimoroni Tufty 2350 over USB.

    python3 install.py                  # fetch the latest app from the site and install it
    python3 install.py --local DIR      # install the app folder DIR instead (development)
    python3 install.py --port /dev/cu.usbmodem1101

What it does: makes sure `mpremote` is available (pip-installs it if not), downloads
__init__.py / room.py / games.py / icon.png from https://wes.imadeyousomethingwithai.com, remounts the badge's
read-only /system partition read-write (the same trick the Walker app's OTA uses), copies the
files to /system/apps/wesworld/, remounts read-only and resets the badge. The app then shows
up in the Badgeware menu as "Wesworld". Needs python 3.8+ and the badge on USB, no Thonny.
"""
import glob
import os
import subprocess
import sys
import tempfile
import urllib.request

SITE = "https://wes.imadeyousomethingwithai.com/"
FILES = ["__init__.py", "room.py", "games.py", "icon.png"]
APP_DIR = "/system/apps/wesworld"


def say(msg):
    print("\033[36m»\033[0m", msg)


def ensure_mpremote():
    if subprocess.run(["mpremote", "version"], capture_output=True).returncode == 0:
        return ["mpremote"]
    r = subprocess.run([sys.executable, "-m", "mpremote", "version"], capture_output=True)
    if r.returncode == 0:
        return [sys.executable, "-m", "mpremote"]
    say("installing mpremote (the MicroPython USB tool)...")
    for args in (["install", "--user", "mpremote"], ["install", "--user", "--break-system-packages", "mpremote"]):
        if subprocess.run([sys.executable, "-m", "pip", *args], capture_output=True).returncode == 0:
            break
    r = subprocess.run([sys.executable, "-m", "mpremote", "version"], capture_output=True)
    if r.returncode != 0:
        sys.exit("could not install mpremote. try: python3 -m pip install mpremote   (or brew install mpremote)")
    return [sys.executable, "-m", "mpremote"]


def find_port(explicit):
    if explicit:
        return explicit
    cands = sorted(glob.glob("/dev/cu.usbmodem*") + glob.glob("/dev/ttyACM*") + glob.glob("/dev/ttyUSB*"))
    if sys.platform.startswith("win"):
        return None          # mpremote auto-detects on Windows
    if not cands:
        sys.exit("no badge found on USB. plug the Tufty in (data cable, not charge-only) and try again.")
    if len(cands) > 1:
        say("several serial ports found, using %s (pass --port to pick another)" % cands[0])
    return cands[0]


def fetch(url, dest):
    say("downloading " + url)
    with urllib.request.urlopen(url, timeout=30) as r, open(dest, "wb") as f:
        f.write(r.read())


# One exec does the whole write: each mpremote connection may soft-reset the board, which
# would undo a read-write remount made in a previous call. Files are staged to /state first.
INSTALL = '''
import os, vfs, rp2
def remount(ro):
    os.chdir("/")
    try:
        vfs.umount("/system")
    except OSError:
        pass
    user = rp2.Flash().ioctl(4, 0) * rp2.Flash().ioctl(5, 0)
    fat = vfs.VfsFat(rp2.Flash(start=0, len=user - 1024 * 1024))
    vfs.mount(fat, "/system", readonly=ro)
remount(False)
try:
    try:
        os.mkdir("%(app)s")
    except OSError:
        pass
    for name in %(files)r:
        with open("%(stage)s/" + name, "rb") as src, open("%(app)s/" + name + ".new", "wb") as dst:
            while True:
                chunk = src.read(4096)
                if not chunk:
                    break
                dst.write(chunk)
        try:
            os.remove("%(app)s/" + name)
        except OSError:
            pass
        os.rename("%(app)s/" + name + ".new", "%(app)s/" + name)
        print("installed", name, os.stat("%(app)s/" + name)[6], "bytes")
finally:
    remount(True)
for name in %(files)r:
    try:
        os.remove("%(stage)s/" + name)
    except OSError:
        pass
st = os.statvfs("/system")
print("system free", st[0] * st[3])
'''
STAGE = "/state/wesworld_stage"


def main():
    args = sys.argv[1:]
    port = args[args.index("--port") + 1] if "--port" in args else None
    local = args[args.index("--local") + 1] if "--local" in args else None
    mp = ensure_mpremote()
    port = find_port(port)
    conn = mp + (["connect", port] if port else [])

    def run(*a, check=True, **kw):
        return subprocess.run(conn + list(a), text=True, check=check, **kw)

    say("checking the badge...")
    r = run("exec", "import sys; print(sys.implementation._machine)", capture_output=True)
    machine = (r.stdout or "").strip()
    if "Tufty 2350" not in machine:
        sys.exit("this does not look like a Tufty 2350 with Badgeware (saw: %r). nothing changed." % machine)
    say("found " + machine)

    tmp = tempfile.mkdtemp(prefix="wesworld-")
    if local:
        src = {f: os.path.join(local, f) for f in FILES}
    else:
        src = {}
        for f in FILES:
            dest = os.path.join(tmp, f)
            fetch(SITE + "?a=app&f=" + f, dest)
            src[f] = dest
    with open(src["__init__.py"], "rb") as fh:
        head = fh.read(400)
    if b"Wes World" not in head:
        sys.exit("the downloaded app does not look right; not installing.")

    say("staging files on the badge...")
    run("fs", "mkdir", ":" + STAGE, check=False, capture_output=True)
    for f in FILES:
        run("fs", "cp", src[f], ":" + STAGE + "/" + f, capture_output=True)
    say("writing %s (remount read-write, copy, remount read-only)..." % APP_DIR)
    r = run("exec", INSTALL % {"app": APP_DIR, "files": FILES, "stage": STAGE}, capture_output=True, check=False)
    print(r.stdout.strip())
    if r.returncode != 0 or "installed __init__.py" not in r.stdout:
        sys.exit("install failed:\n" + (r.stderr or r.stdout))
    say("resetting the badge. look for 'Wesworld' in the menu.")
    run("reset", check=False, capture_output=True)
    say("done. build your lil guy and press B.")


if __name__ == "__main__":
    main()