#!/usr/bin/env python3

import sys
import os
import subprocess
from argparse import ArgumentParser
from common import *

def splitList(l, value):
    try:
        i = l.index(value)
        return l[0:i], l[i+1:]
    except ValueError:
        return l, []


def esp32GetTableFile(img):
    """
    Locate partition table for the image, if it doesn't exists, return the
    default one
    """
    candidate = os.path.splitext(img.path)[0] + ".table"
    if os.path.exists(candidate):
        return candidate
    return os.path.join(buildDir, img.suite, "partition_table", "partition-table.bin")

def esp32GetAppOffset(tablefile):
    import gen_esp32part
    """
    Read partition file, find the first app partition and return its offset
    """
    with open(tablefile, "rb") as f:
        content = f.read()
    table = gen_esp32part.PartitionTable.from_binary(content)
    for partition in table:
        if partition.type == gen_esp32part.APP_TYPE:
            return partition.offset
    raise RuntimeError(f"No APP partition found in {tablefile}")


def flashEsp32Image(img, extra):
    command = ["esptool.py"]
    # Adjust baudrate based on extra argument
    if "-b" not in extra and "--baud" not in extra:
        command += ["--baud", "921600"]
    else:
        try:
            i = extra.index("-b")
        except ValueError:
            i = extra.index("--baud")
        command += extra[i:i+2]
        del extra[i:i+2]
    if "-p" in extra or "--port" in extra:
        try:
            i = extra.index("-p")
        except ValueError:
            i = extra.index("--port")
        command += extra[i:i+2]
        del extra[i:i+2]
    command += [
        "--chip", "esp32",
        "--before=default_reset",
        "--after=hard_reset",
        "write_flash",
        "--flash_mode", "dio",
        "--flash_freq", "40m",
        "--flash_size", "detect"]
    command += extra

    pTable = esp32GetTableFile(img)
    appOffset = esp32GetAppOffset(pTable)
    command += [
        "0x8000", os.path.join(buildDir, img.suite, "partition_table", pTable),
        "0x1000", os.path.join(buildDir, img.suite, "bootloader", "bootloader.bin"),
        f"0x{appOffset:02x}", img.path]

    print(f"Flashing image {img.name} from {img.suite}: {img.path}")
    print(f"Using partition table: {pTable}, app offset: 0x{appOffset:02x}")
    print(command)
    retcode = subprocess.call(command)
    return retcode

def enumerateStLinks():
    proc = subprocess.run(["st-info", "--probe"], capture_output=True)
    if proc.returncode != 0:
        raise RuntimeError(f"Cannot probe ST-Links: {proc.stderr.decode('utf-8')}")
    stlinks = [l.split(":")[1].strip() for l in proc.stdout.decode('utf-8').splitlines() if "serial:" in l]
    return stlinks

def flashStm32ImageAsync(img, extra, id=None):
    if not img.name.endswith(".hex"):
        raise RuntimeError("Only HEX images are supported")
    command = ["st-flash", "--format", "ihex", "--freq=3500"]
    if id is not None:
        command += ["--serial", id]
    command += extra
    command += ["write", img.path]
    return subprocess.Popen(command)

def flashStm32Image(img, extra):
    proc = flashStm32ImageAsync(img, extra)
    proc.wait()
    return proc.returncode

def flashStm32Images(img, all, extra):
    if not all:
        return flashStm32Image(img, extra)
    procs = [flashStm32ImageAsync(img, extra, id) for id in enumerateStLinks()]
    for p in procs:
        p.wait()
    for p in procs:
        if p.returncode != 0:
            return 1
    return 0

def run():
    parser = ArgumentParser(
        prog="rflash",
        description="Flash embedded images from RoFI suite"
    )
    parser.add_argument("image", type=str)
    parser.add_argument("-a", "--all", action='store_true')

    rflashArgs, extraArgs = splitList(sys.argv, "--")
    args = parser.parse_args(rflashArgs[1:])


    imgs = collectImages()
    try:
        img = imgs[args.image]
    except KeyError:
        sys.exit(f"Unknown image {args.image}")
    if img.suite == "esp32":
        retcode = flashEsp32Image(img, extraArgs)
        sys.exit(retcode)
    if img.suite == "stm32":
        retcode = flashStm32Images(img, all, extraArgs)
        enumerateStLinks() # Forces restart of all MCUs
        sys.exit(retcode)
    else:
        sys.exit(f"Unsupported platform {args.image}")

if __name__ == "__main__":
    try:
        run()
    except KeyboardInterrupt:
        sys.exit("Interrupted.")
