#!/usr/bin/python

import argparse
import os
import subprocess
import tempfile

parser = argparse.ArgumentParser(description='Writes a file to NOR flash')
parser.add_argument('filename', help='file to write to NOR flash')
parser.add_argument('-o', '--offset', default=0, help='Offset in flash to start the write')
args = parser.parse_args()

uboot = "u-boot-librem5-uart4.imx"

memLoc0 = 0x43000000
memLoc1 = 0x44000000

chunkSize = 64 * 1024
fileStats = os.stat(args.filename)
if fileStats.st_size % chunkSize == 0:
    chunks = int(fileStats.st_size / chunkSize)
else:
    chunks = int(fileStats.st_size / chunkSize + 1)
size = chunks * chunkSize
hexWords = int(size / 4)
chunkWords = int(chunkSize / 4)
cwd = os.getcwd()
print("CWD ", cwd)
uuuFile = tempfile.NamedTemporaryFile(dir=cwd,delete=False)

verify = ""
for i in range(chunks):
        offset = i * chunkSize
        idx0 = memLoc0 + offset
        idx1 = memLoc1 + offset
        verify = verify + f'FB: ucmd sf erase {args.offset+offset:x} {chunkSize:x}\n'
        verify = verify + f'FB: ucmd sf write {idx0:x} {args.offset+offset:x} {chunkSize:x}\n' 
        verify = verify + f'FB: ucmd sf read {idx1:x} {args.offset+offset:x} {chunkSize:x}\n'
        verify = verify + f'FB: ucmd cmp {idx0:x} {idx1:x} {chunkWords:x}\n'

if not os.path.isfile(f'files/{uboot}'):
    print("u-boot not found, downloading...")
    try:
        os.mkdir('files/')
    except:
        pass
    
    os.system('wget -Ofiles/u-boot-librem5.imx https://arm01.puri.sm/job/u-boot_builds/job/uboot_librem5_build/lastSuccessfulBuild/artifact/output/uboot-librem5/u-boot-librem5.imx')


uuuScript = \
'uuu_version 1.0.1\n' \
'# Use uboot to write a file to the rootfs\n' \
'CFG: FB:  -vid 0x316d -pid 0x4c05 -t 9000 -T 9000\n' \
'CFG: SDP: -chip MX8MQ -compatible MX8MQ -vid 0x316d -pid 0x4c05\n' \
'\n' \
f'SDP: boot -f files/{uboot}\n' \
'# This command will be run when use SPL\n' \
'SDPU: delay 1000\n' \
f'SDPU: write -f files/{uboot} -offset 0x57c00\n' \
'SDPU: jump\n' \
'SDPV: delay 1000\n' \
f'SDPV: write -f files/{uboot} -skipspl\n' \
'SDPV: jump\n' \
f'FB: ucmd setenv fastboot_buffer {memLoc0:x}\n' \
f'FB: download -f {args.filename}\n' \
'FB: ucmd sf probe\n' \
'FB: delay 100\n' \
f'{verify}\n' \
'FB: delay 100\n' \
'FB: Done\n'

uuuFile.write(uuuScript.encode())
uuuFile.close()

print(f'About to run {uuuFile.name}\n{uuuScript}')

subprocess.run(['/usr/local/bin/uuu', '-v', f'{uuuFile.name}'], cwd=cwd)
os.remove(f'{uuuFile.name}')
