#!/usr/bin/env python3

import sys
import struct
import subprocess
import re
import os
import os.path
import argparse


UF2_MAGIC_START0 = 0x0A324655 # "UF2\n"
UF2_MAGIC_START1 = 0x9E5D5157 # Randomly selected
UF2_MAGIC_END    = 0x0AB16F30 # Ditto

INFO_FILE = "/INFO_UF2.TXT"

appstartaddr = 0x2000
familyid = 0x0


def is_uf2(buf):
    w = struct.unpack("<II", buf[0:8])
    return w[0] == UF2_MAGIC_START0 and w[1] == UF2_MAGIC_START1

def is_hex(filename):
    with open(filename, mode='r') as file:
        try:
            for line in file:
                line = line.strip()
                if not line:
                    continue
                if line[0] == ':':
                    continue
                return False
            return True
        except:
            return False

def convert_from_uf2(buf):
    global appstartaddr
    numblocks = len(buf) // 512
    curraddr = None
    outp = []
    for blockno in range(numblocks):
        ptr = blockno * 512
        block = buf[ptr:ptr + 512]
        hd = struct.unpack(b"<IIIIIIII", block[0:32])
        if hd[0] != UF2_MAGIC_START0 or hd[1] != UF2_MAGIC_START1:
            print("Skipping block at " + str(ptr))
            continue
        if hd[2] & 1:
            # NO-flash flag set; skip block
            continue
        datalen = hd[4]
        if datalen > 476:
            assert False, "Invalid UF2 data size at " + str(ptr)
        newaddr = hd[3]
        if curraddr == None:
            appstartaddr = newaddr
            curraddr = newaddr
        padding = newaddr - curraddr
        if padding < 0:
            assert False, "Block out of order at " + str(ptr)
        if padding > 10*1024*1024:
            assert False, "More than 10M of padding needed at " + str(ptr)
        if padding % 4 != 0:
            assert False, "Non-word padding size at " + str(ptr)
        while padding > 0:
            padding -= 4
            outp.append(b"\x00\x00\x00\x00")
        outp.append(block[32:32 + datalen])
        curraddr = newaddr + datalen
    return b"".join(outp)

def convert_to_carray(file_content):
    outp = "const unsigned char bindata_len = %d;\n" % len(file_content)
    outp += "const unsigned char bindata[] __attribute__((aligned(16))) = {"
    for i in range(len(file_content)):
        if i % 16 == 0:
            outp += "\n"
        outp += "0x%02x, " % file_content[i]
    outp += "\n};\n"
    return bytes(outp, "utf-8")

def convert_to_uf2(file_content):
    global familyid
    datapadding = b""
    while len(datapadding) < 512 - 256 - 32 - 4:
        datapadding += b"\x00\x00\x00\x00"
    numblocks = (len(file_content) + 255) // 256
    outp = []
    for blockno in range(numblocks):
        ptr = 256 * blockno
        chunk = file_content[ptr:ptr + 256]
        flags = 0x0
        if familyid:
            flags |= 0x2000
        hd = struct.pack(b"<IIIIIIII",
            UF2_MAGIC_START0, UF2_MAGIC_START1,
            flags, ptr + appstartaddr, 256, blockno, numblocks, familyid)
        while len(chunk) < 256:
            chunk += b"\x00"
        block = hd + chunk + datapadding + struct.pack(b"<I", UF2_MAGIC_END)
        assert len(block) == 512
        outp.append(block)
    return b"".join(outp)

class Block:
    def __init__(self, addr):
        self.addr = addr
        self.bytes = bytearray(256)

    def encode(self, blockno, numblocks):
        global familyid
        flags = 0x0
        if familyid:
            flags |= 0x2000
        hd = struct.pack("<IIIIIIII",
            UF2_MAGIC_START0, UF2_MAGIC_START1,
            flags, self.addr, 256, blockno, numblocks, familyid)
        datapadding = b"\x00" * (512 - 256 - 32 - 4)
        block = hd + self.bytes + datapadding + struct.pack("<I", UF2_MAGIC_END)
        return block

def convert_from_hex_to_uf2(records):
    global appstartaddr
    appstartaddr = None
    upper = 0
    blocks = {}
    for line in records:
        if line[0] != ':':
            continue
        (lenstr, addrstr, typestr, data, chkstr) = (line[1:3], line[3:7], line[7:9], line[9:-2], line[-2:])
        if int(chkstr, 16) != (-(sum(int(data[i:i+2], 16) for i in range(0, len(data), 2)) + int(typestr, 16) + int(addrstr, 16) + int(lenstr, 16)) & 0xff):
            assert False, "Invalid hex checksum for line: " + line
        tp = int(typestr, 16)
        if tp == 4:
            upper = int(data, 16) << 16
        elif tp == 2:
            upper = int(data, 16) << 4
        elif tp == 1:
            break
        elif tp == 0:
            addr = upper + int(addrstr, 16)
            if appstartaddr == None:
                appstartaddr = addr
            i = 0
            while i < len(data):
                if addr in blocks:
                    block = blocks[addr]
                else:
                    block = Block(addr & ~0xff)
                    blocks[addr & ~0xff] = block
                block.bytes[addr & 0xff] = int(data[i:i+2], 16)
                addr += 1
                i += 2

    blocks = sorted(blocks.values(), key=lambda x: x.addr)
    return b"".join(block.encode(i, len(blocks)) for i, block in enumerate(blocks))

def main():
    global appstartaddr, familyid
    def error(msg):
        print(msg)
        sys.exit(1)
    parser = argparse.ArgumentParser(description='Convert to UF2 or flash directly.')
    parser.add_argument('input', metavar='INPUT', type=str, nargs='?',
                        help='input file (HEX, BIN or UF2)')
    parser.add_argument('-b' , '--base', dest='base', type=str,
                        default="0x2000",
                        help='set base address of application for BIN format (default: 0x2000)')
    parser.add_argument('-o' , '--output', metavar="FILE", dest='output', type=str,
                        help='write output to named file; defaults to "flash.uf2" or "flash.bin" where sensible')
    parser.add_argument('-d' , '--device', dest="device_path",
                        help='select a device path to flash')
    parser.add_argument('-l' , '--list', action='store_true',
                        help='list connected devices')
    parser.add_argument('-c' , '--convert', action='store_true',
                        help='do not flash, just convert')
    parser.add_argument('-D' , '--deploy', action='store_true',
                        help='just flash, do not convert')
    parser.add_argument('-f' , '--family', dest='family', type=str,
                        default="0x0",
                        help='specify familyID - number or name (default: 0x0)')
    parser.add_argument('-C' , '--carray', action='store_true',
                        help='convert binary file to a C array, not UF2')
    args = parser.parse_args()
    appstartaddr = int(args.base, 0)

    if args.family.upper() in ["RP2040"]:
        familyid = 0xe48bff56
    else:
        try:
            familyid = int(args.family, 0)
        except ValueError:
            error("Family ID needs to be a number or one of: RP2040")

    if args.list:
        drives = get_drives()
        if len(drives) == 0:
            error("No drives found.")
        for d in drives:
            print(d, info_uf2(d))
        return
    
    if not args.input:
        error("Need input file")
    with open(args.input, mode='rb') as f:
        inpbuf = f.read()
    from_uf2 = is_uf2(inpbuf)
    ext = os.path.splitext(args.input)[1].lower()
    if from_uf2:
        outbuf = convert_from_uf2(inpbuf)
    elif is_hex(args.input):
        outbuf = convert_from_hex_to_uf2(inpbuf.decode("utf-8").split('\n'))
    elif ext == ".bin":
        if args.carray:
            outbuf = convert_to_carray(inpbuf)
        else:
            outbuf = convert_to_uf2(inpbuf)
    else:
        error("Extension %s not supported." % ext)

    if args.deploy:
        drives = get_drives()
        if len(drives) == 0:
            error("No drives to deploy.")
        for d in drives:
            print("Flashing %s (%s)" % (d, info_uf2(d)))
            with open(d + "NEW.UF2", "wb") as f:
                f.write(outbuf)
    elif args.output == None:
        if args.carray:
            print(outbuf.decode("utf-8"))
        else:
            drives = get_drives()
            if len(drives) == 1:
                args.output = drives[0] + "NEW.UF2"
            else:
                if from_uf2:
                    args.output = "flash.bin"
                else:
                    args.output = "flash.uf2"
    if args.output:
        with open(args.output, mode='wb') as f:
            f.write(outbuf)
        print("Wrote %d bytes to %s." % (len(outbuf), args.output))

def get_drives():
    def check_errors(r):
        if r.returncode != 0:
            return []
        return r.stdout.split('\n')

    if sys.platform == "win32":
        return [r + "\\" for r in check_errors(subprocess.run(
            ['wmic', 'logicaldisk', 'get', 'size,freespace,caption'],
            capture_output=True, text=True)) if r and not r.startswith("Caption")]
    elif sys.platform == "darwin":
        def parse_os_x_mount_output(mount_output):
            drives = []
            for line in mount_output:
                m = re.match(r'^/dev/disk.*? on (.*?) \([^/]*\)$', line)
                if m:
                    drives.append(m.group(1) + "/")
            return drives
        return parse_os_x_mount_output(check_errors(subprocess.run(['mount'], capture_output=True, text=True)))
    else:
        def parse_linux_mount_output(mount_output):
            drives = []
            for line in mount_output:
                words = line.split()
                if len(words) >= 3:
                    drives.append(words[2] + "/")
            return drives
        return parse_linux_mount_output(check_errors(subprocess.run(['mount'], capture_output=True, text=True)))

def info_uf2(d):
    try:
        with open(d + INFO_FILE, mode='r') as f:
            return f.read()
    except:
        return "UF2 Bootloader"


if __name__ == "__main__":
    main()