From da62fa98d438bd2dc545018b66d2b542aacc960d Mon Sep 17 00:00:00 2001
From: Alexander Bessman <alexander.bessman@gmail.com>
Date: Tue, 14 Nov 2023 15:58:59 +0100
Subject: [PATCH] Add preliminary support for dropping to bootloader

---
 pslab/cli.py            | 36 ++++++++++++++++++++----------------
 pslab/serial_handler.py |  3 +++
 2 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/pslab/cli.py b/pslab/cli.py
index 3d639548..208b266c 100644
--- a/pslab/cli.py
+++ b/pslab/cli.py
@@ -233,12 +233,12 @@ def main(args: argparse.Namespace):
         install(args)
         return
 
+    handler = SerialHandler(port=args.port)
+
     if args.function == "flash":
-        flash(args)
+        flash(handler, args.hexfile)
         return
 
-    handler = SerialHandler(port=args.port)
-
     if args.function == "collect":
         collect(handler, args)
     elif args.function == "wave":
@@ -524,15 +524,18 @@ def add_install_args(subparser: argparse._SubParsersAction):
     )
 
 
-def flash(args: argparse.Namespace):
-    """Flash firmware over USB.
+def flash(handler: SerialHandler, hexfile: str):
+    """Flash firmware over USB."""
+    handler.drop_to_boot()
+    bootattrs = mcbootflash.get_boot_attrs(handler)
+    mcbootflash.erase_flash(handler, bootattrs.memory_range, bootattrs.erase_size)
 
-    Parameters
-    ----------
-    args : :class:`argparse.Namespace`
-        Parsed arguments.
-    """
-    mcbootflash.flash(args)
+    for chunk in mcbootflash.chunked(hexfile, bootattrs):
+        mcbootflash.write_flash(handler, chunk)
+        mcbootflash.checksum(handler, chunk)
+
+    mcbootflash.self_verify(handler)
+    mcbootflash.reset(handler)
 
 
 def add_flash_args(subparser: argparse._SubParsersAction):
@@ -543,8 +546,9 @@ def add_flash_args(subparser: argparse._SubParsersAction):
     subparser : :class:`argparse._SubParsersAction`
         SubParser to add other arguments related to flash function.
     """
-    parser = mcbootflash.get_parser()
-    parser.prog = "pslab"
-    parser.usage = "Flash firmware to PSLab v6."
-    parser.add_argument("-b", "--baudrate", default=460800, help=argparse.SUPPRESS)
-    subparser.add_parser("flash", parents=[parser], add_help=False)
+    flash = subparser.add_parser("flash")
+    flash.add_argument(
+        "hexfile",
+        type=str,
+        help="an Intel HEX file containing application firmware",
+    )
diff --git a/pslab/serial_handler.py b/pslab/serial_handler.py
index acea7c9d..4018faa2 100755
--- a/pslab/serial_handler.py
+++ b/pslab/serial_handler.py
@@ -399,6 +399,9 @@ def wait_for_data(self, timeout: float = 0.2) -> bool:
 
         return False
 
+    def drop_to_boot(self):
+        self.send_int(0xB007)
+
 
 RECORDED_TRAFFIC = iter([])
 """An iterator returning (request, response) pairs.