#!/usr/bin/env python3
"""Build the quickboot _update.mcs for the SMOD ATSC exciter.

Runs exactly the three scripts the project has always used, in order:

  1. tcl_1/post_bitstream.tcl                    impl_1/*.bit,.bin,.ltx -> bitstream/
  2. tcl_1/runScript_GenRu.tcl                   write_cfgmem -size 64 -interface SPIx1
                                                 -> bitstream/atsc_exciter.mcs
  3. tcl_1/MakeSpiFlashProgrammerMcsFiles_backup.pl -imagesize 64 atsc_exciter.mcs
                                                 -> bitstream/atsc_exciter_update.mcs

Step 2 does ``glob *.bit`` inside bitstream/, so every stale .bit sitting there gets
concatenated into one giant image. bitstream/ accumulates iteration bits, so this
script moves the old ones aside for the duration and puts them back afterwards.

The resulting _update.mcs is ONLY for the network remote update (scripts/fpga_ru.py).
It cannot be JTAG-flashed -- see references/quickboot.md.
"""

import argparse
import glob
import os
import re
import shutil
import subprocess
import sys
import tempfile

REPO = r"E:\Xilinx_design\smod\SMOD_ATSC\my_exciter"
TCL_DIR = os.path.join(REPO, "my_exciter.srcs", "tcl_1")
BITSTREAM = os.path.join(REPO, "bitstream")
IMPL = os.path.join(REPO, "my_exciter.runs", "impl_1")

POST_BITSTREAM = os.path.join(TCL_DIR, "post_bitstream.tcl")
GEN_RU = os.path.join(TCL_DIR, "runScript_GenRu.tcl")
PERL_SCRIPT = os.path.join(TCL_DIR, "MakeSpiFlashProgrammerMcsFiles_backup.pl")
DEFAULT_VIVADO = r"D:\Xilinx\Vivado\2018.1\bin\vivado.bat"

# runScript_GenRu.tcl hardcodes this name.
BASE_MCS = "atsc_exciter.mcs"
UPDATE_MCS = "atsc_exciter_update.mcs"

# The perl script emits these alongside the one we want.
BYPRODUCTS = ("atsc_exciter_initial.mcs", "atsc_exciter_golden.mcs", "atsc_exciter_total.mcs",
              "atsc_exciter_update2.mcs", "atsc_exciter_simple.mcs")

# ":02000004HHHH" is the Intel-HEX extended-linear-address record.
ELA_RE = re.compile(rb"^:02000004([0-9A-Fa-f]{4})")
EXPECT_ELA = 0x800000


def sh(cmd, cwd=None):
    print(f"  $ {' '.join(cmd)}", flush=True)
    if subprocess.run(cmd, cwd=cwd).returncode != 0:
        sys.exit("[FAIL] command failed; see output above")


def first_ela(path):
    with open(path, "rb") as fh:
        for line in fh:
            m = ELA_RE.match(line)
            if m:
                return int(m.group(1), 16) << 16
    return None


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--vivado", default=DEFAULT_VIVADO)
    ap.add_argument("--tag", help="rename the result to <tag>_update.mcs so it survives the next run")
    ap.add_argument("--keep-byproducts", action="store_true")
    args = ap.parse_args()

    for path, what in ((args.vivado, "vivado"), (POST_BITSTREAM, "post_bitstream.tcl"),
                       (GEN_RU, "runScript_GenRu.tcl"), (PERL_SCRIPT, "perl script")):
        if not os.path.exists(path):
            sys.exit(f"[FAIL] {what} not found: {path}")

    impl_bits = glob.glob(os.path.join(IMPL, "*.bit"))
    if len(impl_bits) != 1:
        sys.exit(f"[FAIL] expected exactly one .bit in {IMPL}, found {len(impl_bits)}: "
                 f"{[os.path.basename(b) for b in impl_bits]}")
    print(f"[1/4] implementation bit: {os.path.basename(impl_bits[0])}")

    # Park every stale .bit so runScript_GenRu.tcl's `glob *.bit` can only see the fresh one.
    stash = tempfile.mkdtemp(prefix="smod_bit_stash_")
    parked = glob.glob(os.path.join(BITSTREAM, "*.bit"))
    for b in parked:
        shutil.move(b, os.path.join(stash, os.path.basename(b)))
    if parked:
        print(f"[2/4] parked {len(parked)} stale .bit files from bitstream/ (restored at the end)")

    try:
        # One vivado launch, both tcl scripts: each re-derives its own path from [info script].
        wrapper = os.path.join(stash, "run_both.tcl")
        with open(wrapper, "w") as fh:
            fh.write(f'source {{{POST_BITSTREAM}}}\n')
            fh.write(f'source {{{GEN_RU}}}\n')
        print(f"[3/4] vivado: post_bitstream.tcl + runScript_GenRu.tcl -> {BASE_MCS}")
        sh([args.vivado, "-mode", "batch", "-nojournal", "-notrace", "-source", wrapper], cwd=stash)

        base = os.path.join(BITSTREAM, BASE_MCS)
        if not os.path.exists(base):
            sys.exit(f"[FAIL] {base} was not produced")

        print(f"[4/4] perl MakeSpiFlashProgrammerMcsFiles_backup.pl -imagesize 64 {BASE_MCS}")
        sh(["perl", PERL_SCRIPT, "-imagesize", "64", BASE_MCS], cwd=BITSTREAM)
    finally:
        for name in os.listdir(stash):
            if not name.endswith(".bit"):
                continue
            dest = os.path.join(BITSTREAM, name)
            # post_bitstream.tcl just copied the fresh impl_1 bit here under this same
            # name; putting the parked (stale) one back would silently undo that.
            if not os.path.exists(dest):
                shutil.move(os.path.join(stash, name), dest)
        shutil.rmtree(stash, ignore_errors=True)

    update = os.path.join(BITSTREAM, UPDATE_MCS)
    if not os.path.exists(update):
        sys.exit(f"[FAIL] {UPDATE_MCS} was not produced")

    # -imagesize 64 must place the image at 0x800000. If it doesn't, the payload is not
    # 8 MB and the FPGA's CRC check over 0x400000..0xC00000 will fail after a long upload.
    ela = first_ela(update)
    if ela != EXPECT_ELA:
        got = "none" if ela is None else f"0x{ela:X}"
        sys.exit(f"[FAIL] {UPDATE_MCS} starts at {got}, expected 0x{EXPECT_ELA:X}. Do not upload it.")

    if not args.keep_byproducts:
        for name in BYPRODUCTS:
            p = os.path.join(BITSTREAM, name)
            if os.path.exists(p):
                os.remove(p)

    if args.tag:
        tagged = os.path.join(BITSTREAM, f"{args.tag}_update.mcs")
        shutil.move(update, tagged)
        update = tagged

    print(f"\n[OK] {update}  ({os.path.getsize(update)/1e6:.1f} MB)")
    print(f"     base mcs kept at {os.path.join(BITSTREAM, BASE_MCS)}")
    print("\nNext: python scripts/fpga_ru.py --mcs \"%s\"" % update)
    print("      Turn DPD off first -- see references/protocol.md.")


if __name__ == "__main__":
    main()