from __future__ import annotations

import argparse
from pathlib import Path


RELEASE_TAG = "matrix-axp313a1"
MARKER = "matrix-kernel-good"


class BootScriptError(RuntimeError):
    pass


def render(source: str) -> str:
    if "matrix_kernel_candidate=1" in source:
        raise BootScriptError("boot script is already candidate-managed")
    anchor = (
        'if test "${docker_optimizations}" = "on"; then setenv bootargs '
        '"${bootargs} cgroup_enable=memory swapaccount=1"; fi\n'
    )
    if source.count(anchor) != 1:
        raise BootScriptError("bootargs anchor is missing or ambiguous")
    candidate = f'''

# Matrix AXP313A candidate: consume the marker before boot so a failed reboot falls back.
setenv matrix_kernel_file "Image"
setenv matrix_dtb_suffix ""
if test -e ${{devtype}} ${{devnum}} ${{prefix}}{MARKER}; then
    if test -e ${{devtype}} ${{devnum}} ${{prefix}}Image-{RELEASE_TAG}; then
        fatrm ${{devtype}} ${{devnum}} ${{prefix}}{MARKER}
        setenv matrix_kernel_file "Image-{RELEASE_TAG}"
        setenv matrix_dtb_suffix "-{RELEASE_TAG}"
        setenv bootargs "${{bootargs}} matrix_kernel_candidate=1 panic=10"
        echo "Booting one-shot Matrix AXP313A candidate"
    fi
fi
'''
    replacements = {
        "${prefix}${fdtfile_emmc}.dtb": "${prefix}${fdtfile_emmc}${matrix_dtb_suffix}.dtb",
        "${prefix}${fdtfile}.dtb": "${prefix}${fdtfile}${matrix_dtb_suffix}.dtb",
        "${prefix}Image": "${prefix}${matrix_kernel_file}",
    }
    expected = {
        "${prefix}${fdtfile_emmc}.dtb": 1,
        "${prefix}${fdtfile}.dtb": 2,
        "${prefix}Image": 1,
    }
    for old, new in replacements.items():
        if source.count(old) != expected[old]:
            raise BootScriptError(f"boot path {old!r} is missing or ambiguous")
        source = source.replace(old, new)
    rendered = source.replace(anchor, anchor + candidate, 1)
    return rendered


def main() -> int:
    parser = argparse.ArgumentParser(description="Render a one-shot dual-kernel WalnutPi boot.cmd")
    parser.add_argument("source", type=Path)
    parser.add_argument("destination", type=Path)
    args = parser.parse_args()
    source = args.source.read_text(encoding="utf-8")
    rendered = render(source)
    args.destination.write_text(rendered, encoding="utf-8", newline="\n")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())