Import("env")
import os
import subprocess

# Real, already-built bootloader image used as the reference for the presence
# check below -- a raw synthetic probe blob turned out unreliable with
# verify_image (silent no-op on tiny files), whereas verify_image against a
# real firmware .bin has been solid throughout this project's bring-up.
BOOTLOADER_BIN = os.path.join(
    "bootloader", ".pio", "build", "versapad_bootloader", "firmware.bin"
)


def _openocd_paths(env):
    pkg_dir = env.PioPlatform().get_package_dir("tool-openocd")
    return (
        os.path.join(pkg_dir, "bin", "openocd.exe"),
        os.path.join(pkg_dir, "scripts"),
    )


def _run_openocd(env, extra_cmd, capture=False):
    openocd, scripts = _openocd_paths(env)
    cmd = [
        openocd,
        "-s", scripts,
        "-f", "interface/cmsis-dap.cfg",
        "-f", "target/at91samdXX.cfg",
        "-c", extra_cmd,
    ]
    print(" ".join(cmd))
    if capture:
        return subprocess.run(cmd, capture_output=True, text=True)
    return subprocess.run(cmd)


def _bootloader_present(env):
    bootloader_bin = env.subst(
        os.path.join("$PROJECT_DIR", BOOTLOADER_BIN)
    )
    if not os.path.isfile(bootloader_bin):
        print("WARNING: {} not found (build it with".format(BOOTLOADER_BIN))
        print("'cd bootloader && pio run -e versapad_bootloader') -- can't check")
        print("for an installed bootloader, refusing to flash as a precaution.")
        print("Use 'pio run -e versapad -t erase-bootloader-and-flash' to override.")
        return True

    result = _run_openocd(
        env,
        'init; reset halt; verify_image "{}" 0x0; shutdown'.format(
            bootloader_bin.replace("\\", "/")
        ),
        capture=True,
    )
    output = result.stdout + result.stderr
    if "checksum mismatch" in output or "diff " in output:
        return False  # something else is at 0x0000 -- not this bootloader
    if "halted due to debug-request" in output and "Error" not in output:
        return True  # verify_image is silent on a match; absence of a
        # mismatch/error after a successful connect means it matched
    print("WARNING: could not read flash to check for an installed bootloader")
    print("(inconclusive) -- refusing to flash as a precaution.")
    print("Use 'pio run -e versapad -t erase-bootloader-and-flash' to override.")
    return True


def _flash(env, firmware):
    result = _run_openocd(
        env, 'program "{}" verify reset; shutdown'.format(firmware.replace("\\", "/"))
    )
    if result.returncode != 0:
        env.Exit(1)


def upload_via_openocd(source, target, env):
    firmware = str(source[0])  # .elf path

    # This script is shared with bootloader/platformio.ini's own upload
    # (env:versapad_bootloader), which legitimately writes 0x0000 every time
    # -- the guard below only makes sense for the app-without-bootloader
    # target (env:versapad).
    if env["PIOENV"] == "versapad" and _bootloader_present(env):
        print("=" * 78)
        print("REFUSING TO FLASH: a UF2 bootloader looks like it's installed at 0x0000.")
        print("")
        print("This target (env:versapad) writes the app starting at 0x0000 and")
        print("would silently overwrite it -- the board would lose its USB flashing")
        print("path (see doc/10_usb_bootloader.md).")
        print("")
        print("Use instead:")
        print("  pio run -e versapad_usb --target upload   # flash over USB, keeps the bootloader")
        print("or, if you deliberately want to erase the bootloader and go back to")
        print("standalone SWD-only operation:")
        print("  pio run -e versapad -t erase-bootloader-and-flash")
        print("=" * 78)
        env.Exit(1)

    _flash(env, firmware)


env.Replace(UPLOADCMD=upload_via_openocd)


def erase_bootloader_and_flash(*_args, **_kwargs):
    firmware = env.subst(os.path.join("$BUILD_DIR", "${PROGNAME}.elf"))
    print("Overwriting 0x0000..0x1FAFF -- any installed UF2 bootloader will be erased.")
    _flash(env, firmware)


env.AddCustomTarget(
    name="erase-bootloader-and-flash",
    dependencies=["buildprog"],
    actions=[erase_bootloader_and_flash],
    title="Erase bootloader + flash (standalone SWD)",
    description=(
        "Unconditionally overwrites 0x0000..0x1FAFF, wiping any installed UF2 "
        "bootloader. Use only to return to standalone SWD-only operation."
    ),
)