Import("env", "projenv")
import os
import glob

print("##post script running")

def getString(buffer, offset, length):
    return buffer[offset:offset+length].rstrip(b'\0').decode('utf-8')

def getFirmwareInfo(fwfile):
    """
    ESP-IDF firmware descriptor

    typedef struct {
        uint32_t magic_word;
        uint32_t secure_version;
        uint32_t reserv1;
        uint32_t reserv2;
        char version[32];
        char project_name[32];
        char time[16];
        char date[16];
        char idf_ver[32];
        uint8_t app_elf_sha256[32];
    } esp_app_desc_t

    """
    print(f"checking firmware {fwfile}")
    with open(fwfile, "rb") as fh:
        MAGIC = b"\x32\x54\xcd\xab" # little endian
        buffer = fh.read(1024)
        idx = buffer.find(MAGIC, 288) # find MAGIC
        if idx < 0:
            raise Exception("no app descriptor found")
        else:
            print("found app descriptor at {}".format(idx))
        version = getString(buffer, idx + 16, 32)
        name = getString(buffer, idx + 48, 32)
        fwtime = getString(buffer, idx + 80, 16)
        fwdate = getString(buffer, idx + 96, 16)
        idf_ver = getString(buffer, idx + 112, 16)
    print(f"  name: '{name}'")
    print(f"  version: '{version}'")
    print(f"  time: '{fwtime}'")
    print(f"  date: '{fwdate}'")
    print(f"  idf: '{idf_ver}'")
    return (name, version)

def postbuild(source, target, env):
    print("##postbuild")
    print("source={}".format(source[0]))
    print("target={}".format(target[0]))
    firmware = env.subst("$BUILD_DIR/${PROGNAME}.bin")
    (fwname, fwversion) = getFirmwareInfo(firmware)

    esptool = env.get('UPLOADER')
    uploaderflags = env.subst("${UPLOADERFLAGS}")
    base = env.subst("$PIOENV")
    appoffset = env.subst("$ESP32_APP_OFFSET")
    python = env.subst("$PYTHONEXE")

    print("base=%s,esptool=%s,appoffset=%s,uploaderflags=%s"%(base,esptool,appoffset,uploaderflags))

    uploadparts = uploaderflags.split(" ")

    chip = "esp32"
    if len(uploadparts) < 6:
        print("uploaderflags does not have enough parameter")
        return
    for i in range(0, len(uploadparts)):
        if uploadparts[i] == "--chip":
            if i < (len(uploadparts) - 1):
                chip = uploadparts[i+1]
    uploadfiles = uploadparts[-6:]
    for i in range(1, len(uploadfiles), 2):
        if not os.path.isfile(uploadfiles[i]):
            print("file %s for combine not found"%uploadfiles[i])
            return
    offset = uploadfiles[0]

    outdir = env.subst("$BUILD_DIR")
    for f in glob.glob(os.path.join(outdir, base + "*.bin")):
        print("removing old file {}".format(f))
        os.unlink(f)

    outfile = os.path.join(outdir, "{}-all.bin".format((base)))
    cmd = [python, esptool, "--chip", chip, "merge_bin", "--target-offset", offset, "-o", outfile]
    cmd += uploadfiles
    cmd += [appoffset, firmware]
    print("running {}".format(' '.join(cmd)))
    env.Execute(' '.join(cmd), "#merging bin files")

    # Create symlinks to better identify firmware
    fw_update = os.path.join(outdir, f"{base}-{fwversion}-update.bin")
    os.symlink(firmware, fw_update)
    fw_full = os.path.join(outdir, f"{base}-{fwversion}-all.bin")
    os.symlink(outfile, fw_full)

env.AddPostAction(
    "$BUILD_DIR/${PROGNAME}.bin",
    postbuild
)