#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
flash_package.py - erase the chip and program the merged package via
NS-LINK (NSpyocd). Streams the vendor tool's output with the vendor
banner replaced by the CAIIC one (the banner is baked into the
PyInstaller-packed NSpyocd.exe and cannot be edited).

    usage:  flash_package.bat [image]
    [image] optional, default: tools/out/mothercup_ble_prod.hex

Connect NS-LINK to SWD (PA4/PA5) and the reset pin first, and CLOSE
Keil (it holds the probe).
"""

import os
import subprocess
import sys

# NSpyocd is a PyInstaller-packed pyocd: its banner is inside the
# compressed bundle, so we rebrand it on the fly. The console encoding
# varies (UTF-8 vs GBK), match the byte pattern in both.
BANNER_PATTERNS = [
    ("国民技术NSLINK上位机".encode("utf-8"), b"CAIIC NSLINK UMP"),
    ("国民技术NSLINK上位机".encode("gbk"), b"CAIIC NSLINK UMP"),
]

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PYOCD = os.path.join(ROOT, "nations-tec", "N32WB03x_SDK_V2.0.0",
                     "utilities", "dfu", "NSpyocd", "NSpyocd.exe")
# The builtin "n32wb031" target only maps 256KB of flash (fails at
# 0x01040000 on the dual-bank package); the DFP pack target
# "n32wb031keq6_2" maps the full 512KB of the KEQ6-2 we actually use.
PACK = os.path.join(ROOT, "nations-tec", "N32WB03x_DFP.1.4.0.pack")
TARGET = "n32wb031keq6_2"


def run_step(args):
    """Run one NSpyocd step, streaming output with the banner replaced.
    Reads in small chunks (not line-buffered) so the tool's \\r-based
    progress bar redraws live. Returns the process exit code."""
    proc = subprocess.Popen([PYOCD] + args,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.STDOUT)
    out = sys.stdout.buffer
    keep = max(len(p) for p, _ in BANNER_PATTERNS) - 1
    tail = b""
    while True:
        chunk = proc.stdout.read(256)
        if not chunk:
            break
        buf = tail + chunk
        # replace on the COMBINED buffer first (a pattern straddling the
        # chunk boundary only exists contiguously here), then hold back
        # the last maxpat-1 bytes so a pattern split by the next chunk
        # is still seen whole
        for old, new in BANNER_PATTERNS:
            buf = buf.replace(old, new)
        emit, tail = buf[:-keep], buf[-keep:]
        out.write(emit)
        out.flush()
    for old, new in BANNER_PATTERNS:
        tail = tail.replace(old, new)
    out.write(tail)
    out.flush()
    return proc.wait()


def main():
    image = sys.argv[1] if len(sys.argv) > 1 else \
        os.path.join(os.path.dirname(os.path.abspath(__file__)),
                     "out", "mothercup_ble_prod.hex")
    if not os.path.isfile(image):
        print("%s not found - run tools\\make_package.bat first" % image)
        return 1

    # -M under-reset: connect while holding the chip in reset
    # -f 1000000:     1 MHz SWD clock (slow but tolerant of wiring)
    print("Erasing chip...", flush=True)
    rc = run_step(["erase", "--chip", "--pack", PACK, "-M", "under-reset",
                   "-f", "1000000", "-t", TARGET])
    if rc != 0:
        return fail()

    print("Programming %s ..." % image, flush=True)
    # -O smart_flash=false: with the pack target the pre-program diff pass
    # faults reading 0x01000000; programming+verify itself is fine.
    rc = run_step(["load", "--pack", PACK, "-M", "under-reset",
                   "-f", "1000000", "-t", TARGET,
                   "-O", "smart_flash=false", image])
    if rc != 0:
        return fail()

    print("Program Finish!")
    return 0


def fail():
    print()
    print("*** FLASH FAILED ***")
    print("Troubleshooting:")
    print("  1. Close Keil / any tool holding the probe, replug NS-LINK USB")
    print("  2. List probes:  \"%s\" list" % PYOCD)
    print("  3. Check wiring: SWDIO=PA5 SWCLK=PA4 GND and the RESET pin")
    print("  4. If several probes are attached, add  -u <probe_uid>  "
          "to the commands above")
    return 1


if __name__ == "__main__":
    sys.exit(main())