#!/usr/bin/env python2

# Copyright (c) 2013-2015 NVIDIA CORPORATION. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.

import argparse
import binascii
import os
import os.path
import shutil
import stat
import subprocess
import sys
import tempfile
from tegraboardconfigs import *

def mkdir(path):
    if not os.path.isdir(path):
        os.makedirs(path)

def cp(src, dst):
    print '+ cp', src, dst
    shutil.copy(src, dst)

def rmtree(path):
    if os.path.exists(path):
        shutil.rmtree(path)

def run(cd, cmd):
    print '+ cd', cd
    print '+', ' '.join(cmd)
    subprocess.check_call(cmd, cwd=cd)

def gen_flashcmd_mmc(flash_image_addr, readback_addr, flash_img_size):
    flash_id = config['flash-id-uboot']
    flash_img_size_sectors = flash_img_size / 512
    flashcmd = 'echo >>> Selecting MMC device... ; '
    flashcmd += 'mmc dev %d 1 ; ' % flash_id
    flashcmd += 'echo >>> Writing image to MMC... ; '
    flashcmd += 'mmc write 0x%08x 0 0x%x ; ' % (flash_image_addr, flash_img_size_sectors)
    flashcmd += 'echo >>> Reading image back from MMC... ; '
    flashcmd += 'mmc read 0x%08x 0 0x%x ; ' % (readback_addr, flash_img_size_sectors)
    return flashcmd

def gen_flashcmd_nand(flash_image_addr, readback_addr, flash_img_size):
    flashcmd = 'echo >>> Erasing NAND... ; '
    flashcmd += 'nand erase.chip ; '
    flashcmd += 'echo >>> Writing image to NAND... ; '
    flashcmd += 'nand write 0x%08x 0 0x%08x ; ' % (flash_image_addr, flash_img_size)
    flashcmd += 'echo >>> Reading image back from NAND... ; '
    flashcmd += 'nand read 0x%08x 0 0x%08x ; ' % (readback_addr, flash_img_size)
    return flashcmd

def gen_flashcmd_spi(flash_image_addr, readback_addr, flash_img_size):
    flash_id = config.get('flash-id-uboot', '0')
    flashcmd = 'echo >>> Selecting SPI device... ; '
    flashcmd += 'sf probe %s ; ' % flash_id
    flashcmd += 'echo >>> Erasing SPI... ; '
    flashcmd += 'sf erase 0 0x%08x ; ' % config['flash-erase-size']
    flashcmd += 'echo >>> Writing image back to SPI... ; '
    flashcmd += 'sf write 0x%08x 0 0x%08x ; ' % (flash_image_addr, flash_img_size)
    flashcmd += 'echo >>> Reading image back from SPI... ; '
    flashcmd += 'sf read 0x%08x 0 0x%08x ; ' % (readback_addr, flash_img_size)
    return flashcmd

gen_flashcmds = {
    'emmc': gen_flashcmd_mmc,
    'nand': gen_flashcmd_nand,
    'spi': gen_flashcmd_spi,
}

def get_loadaddr():
    spl_entry_fn = os.path.join(out_board_dir, 'u-boot-spl-entry')
    with open(spl_entry_fn, 'rt') as spl_entry_f:
        spl_entry_s = spl_entry_f.read().strip()
    spl_entry = int(spl_entry_s, 0)
    if args.debug:
        print 'spl_entry %d 0x%x' % (spl_entry, spl_entry)
    return spl_entry

def gen_tegrarcm_cmd(bootloader, loadaddr):
    if type(loadaddr) != str:
        loadaddr = "0x%08x" % loadaddr
    cmd = ['tegrarcm', '--bct=' + bct, '--bootloader=' + bootloader , '--loadaddr=' + loadaddr]
    if args.usb_port_path:
        cmd.extend(['--usb-port-path', args.usb_port_path])
    return cmd

def find_config_dir():
    if not configs.has_key(args.configname):
        print 'Unknown config "%s"' % args.configname
        sys.exit(1)

    global config, boardname, board, socname, soc, out_board_dir, bct

    config = configs[args.configname]
    boardname = config['board']
    board = boards[boardname]
    socname = board['soc']
    soc = socs[socname]
    out_board_dir = os.path.join(args.data_dir, boardname)
    bct = os.path.join(out_board_dir, config['bct'])

def func_list_configs():
    for configname in sorted(configs.keys()):
        print configname

def func_flash():
    find_config_dir()

    u_boot_no_dtb = os.path.join(out_board_dir, 'u-boot-nodtb-tegra.bin')
    u_boot_no_dtb_size = os.path.getsize(u_boot_no_dtb)
    if args.debug:
        print 'u_boot_no_dtb_size %d 0x%x' % (u_boot_no_dtb_size, u_boot_no_dtb_size)

    u_boot_dtb = os.path.join(out_board_dir, 'u-boot.dtb')
    u_boot_dtb_size = os.path.getsize(u_boot_dtb)
    if args.debug:
        print 'u_boot_dtb_size %d 0x%x' % (u_boot_dtb_size, u_boot_dtb_size)

    if args.flash_image:
        flash_img = args.flash_image
    else:
        flash_img = os.path.join(out_board_dir, config['flash-image'])
    flash_img_size = os.path.getsize(flash_img)
    if args.debug:
        print 'flash_img_size %d 0x%x' % (flash_img_size, flash_img_size)

    imgf = file(flash_img, 'rb')
    imgd = imgf.read()
    imgf.close()
    flash_img_crc32 = binascii.crc32(imgd)
    if args.debug:
        print 'flash_img_crc32 %x' % flash_img_crc32
    flash_img_crc32_bs = (
        ((flash_img_crc32 & 0xff) << 24) |
        ((flash_img_crc32 & 0xff00) << 8) |
        ((flash_img_crc32 & 0xff0000) >> 8) |
        ((flash_img_crc32 & 0xff000000) >> 24)
    )

    u_boot_plus_dtb_size = u_boot_no_dtb_size + u_boot_dtb_size
    if args.debug:
        print 'u_boot_plus_dtb_size %d 0x%x' % (u_boot_plus_dtb_size, u_boot_plus_dtb_size)

    bss_size_fn = os.path.join(out_board_dir, 'u-boot-bss-size')
    bss_size_f = open(bss_size_fn, 'rt')
    bss_size_s = bss_size_f.read()
    bss_size_f.close()
    bss_size = int(bss_size_s)
    if args.debug:
        print 'bss_size %d 0x%x' % (bss_size, bss_size)

    # Avoid U-Boot's BSS, so the BSS-zeroing doesn't trash the appended flash image
    # Add 4KB in case the DTB size changes due to fdtput
    # Align to 4k, so flash writes don't need a bounce buffer for DMA
    padded_size = (u_boot_plus_dtb_size + bss_size + (2 * 4 * 1024) - 1) & ~((4 * 1024) - 1)
    if args.debug:
        print 'padded_size %d 0x%x' % (padded_size, padded_size)

    pad_size = padded_size - u_boot_plus_dtb_size
    if args.debug:
        print 'pad_size %d 0x%x' % (pad_size, pad_size)

    loadaddr = get_loadaddr()
    flash_image_addr = loadaddr + padded_size
    if args.debug:
        print 'flash_image_addr %d 0x%x' % (flash_image_addr, flash_image_addr)
    readback_addr = flash_image_addr + flash_img_size
    if args.debug:
        print 'readback_addr %d 0x%x' % (readback_addr, readback_addr)

    flash_type = config['flash-type']
    if not gen_flashcmds.has_key(flash_type):
        print 'flash-type "%s" not yet supported' % flash_type
        sys.exit(1)
    gen_flashcmd = gen_flashcmds[flash_type]

    if args.work_dir:
        mkdir(workdir)
    else:
        workdir = tempfile.mkdtemp()
    try:
        u_boot_dtb_runflash = os.path.join(workdir, 'u-boot-runflash.dtb')
        cp(u_boot_dtb, u_boot_dtb_runflash)

        # 0xfffffffe==-2; never delay or interrupt
        cmd = ['fdtput', '-p', '-t', 'x', u_boot_dtb_runflash, '/config', 'bootdelay', '0xfffffffe']
        run(workdir, cmd)

        crc_addr = soc['ram-base'] + 4
        bootcmd = ''
        bootcmd += 'echo >>> Verifying image in RAM... ; '
        bootcmd += 'crc32 0x%08x 0x%08x 0x%08x ; ' % (flash_image_addr, flash_img_size, crc_addr)
        bootcmd += 'if itest.l *0x%08x != 0x%x; then echo CRC MISMATCH of initial image; exit; fi ; ' % (crc_addr, flash_img_crc32_bs)
        bootcmd += gen_flashcmd(flash_image_addr, readback_addr, flash_img_size)
        bootcmd += 'echo >>> Verifying image from flash... ; '
        bootcmd += 'crc32 0x%08x 0x%08x 0x%08x ; ' % (readback_addr, flash_img_size, crc_addr)
        bootcmd += 'if itest.l *0x%08x != 0x%x; then echo CRC MISMATCH of readback image; exit; fi ; ' % (crc_addr, flash_img_crc32_bs)
        bootcmd += 'echo >>> Setting up environment... ; '
        bootcmd += 'env default -f -a ; '
        # Perhaps U-Boot should set $boardname based on the ID EEPROM; then we wouldn't need this
        if config['dtbfn-extra'] != '':
            bootcmd += 'setenv board ' + boardname + config['dtbfn-extra'] + ' ; '
        if args.env:
            for (var, value) in args.env:
                bootcmd += 'setenv %s \'%s\' ; ' % (var, value)
        bootcmd += 'saveenv ; '
        bootcmd += 'echo >>> Flashing OK, rebooting... ; '
        # To update the bootloader, reset.
        # If wanting to run installer, set installer_args.configname in environment, 'run bootcmd'
        if args.post_flash_cmd:
            bootcmd += args.post_flash_cmd
        else:
            bootcmd += 'reset'
        print 'bootcmd:', bootcmd
        cmd = ['fdtput', '-p', '-t', 's', u_boot_dtb_runflash, '/config', 'bootcmd', bootcmd]
        run(workdir, cmd)

        u_boot_dtb_runflash_size = os.path.getsize(u_boot_dtb_runflash)
        if args.debug:
            print 'u_boot_dtb_runflash_size %d 0x%x' % (u_boot_dtb_runflash_size, u_boot_dtb_runflash_size)
        pad_size -= (u_boot_dtb_runflash_size - u_boot_dtb_size)
        if args.debug:
            print 'pad_size %d 0x%x' % (pad_size, pad_size)

        uboot_flasher = os.path.join(workdir, 'u-boot-flasher.bin')
        f = open(uboot_flasher, 'wb')
        shutil.copyfileobj(open(u_boot_no_dtb, 'rb'), f)
        shutil.copyfileobj(open(u_boot_dtb_runflash, 'rb'), f)
        f.write(chr(0) * pad_size)
        shutil.copyfileobj(open(flash_img, 'rb'), f)
        f.close()

        cmd = gen_tegrarcm_cmd(uboot_flasher, loadaddr)

        flasher_sh = os.path.join(workdir, 'flasher.sh')
        f = open(flasher_sh, 'wt')
        os.fchmod(f.fileno(), stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO)
        f.write("#!/bin/sh\n")
        f.write("\n")
        f.write(' '.join(cmd))
        f.write("\n")
        f.close()

        if not args.gen_only:
            run(workdir, [flasher_sh])
    except:
        raise
    finally:
        if args.save_work_dir:
            print 'Not removing work directory:', workdir
        else:
            rmtree(workdir)

def func_exec():
    find_config_dir()

    if args.bootloader:
        bootloader = args.bootloader
    else:
        bootloader = os.path.join(out_board_dir, 'u-boot-dtb-tegra.bin')

    if args.loadaddr:
        loadaddr = args.loadaddr
    else:
        loadaddr = get_loadaddr()

    cmd = gen_tegrarcm_cmd(bootloader, loadaddr)
    run(scripts_dir, cmd)

parser = argparse.ArgumentParser(description='Execute a bootloader on a ' +
    'Tegra board, possibly modifying it prior to download so as to execute ' +
    'commands, such as writing an image to flash.')

parser.add_argument('--debug', action='store_true',
    help='Turn on debugging prints')
parser.add_argument('--data-dir', type=str,
    help='The directory containing board data')
parser.add_argument('--force-no-out-dir', action='store_true',
    help='Don\'t check for ../_out* directories used in source tree')

subparsers = parser.add_subparsers()

parser_list_configs = subparsers.add_parser('list-configs',
    help='List known board configurations')
parser_list_configs.set_defaults(func = func_list_configs)

parser_flash = subparsers.add_parser('flash',
    help='Write an image, usually U-Boot itself, to flash on the device')
parser_flash.set_defaults(func = func_flash)
parser_flash.add_argument('--work-dir', type=str,
    help='The temporary directory used during operation')
parser_flash.add_argument('--save-work-dir', action='store_true',
    help='Don\'t delete the work-dir after execution')
parser_flash.add_argument('--flash-image', type=str,
    help='The flash image to write, instead of U-Boot itself')
parser_flash.add_argument('--gen-only', action='store_true',
    help='Just create the work-dir; don\'t actually flash the image')
parser_flash.add_argument('--env', type=str, nargs=2, action='append',
    metavar=("VAR", "VALUE"),
    help='Set a U-Boot environment variable after flashing. This option ' +
        'may be given multiple times')
parser_flash.add_argument('--post-flash-cmd', type=str,
    help='The U-Boot command to run after flashing, defaults to \'reset\'')

parser_exec = subparsers.add_parser('exec',
    help='Download and execute an unmodified bootloader binary, named ' +
        'u-boot-dtb-tegra.bin by default. This can be useful to simply test ' +
        'a U-Boot build without flashing it, or to download something other ' +
        'than U-Boot')
parser_exec.set_defaults(func = func_exec)
parser_exec.add_argument('--bootloader', type=str,
    help='The bootloader to download. Defaults to u-boot-dtb-tegra.bin in ' +
        'the data directory')
parser_exec.add_argument('--loadaddr', type=str,
    help='Load address for the bootloader binary')

for p in (parser_flash, parser_exec):
    p.add_argument('--usb-port-path', metavar='USBPORTPATH', type=str,
        help='The USB port path of the Tegra device')
    p.add_argument('configname', metavar='CONFIG', type=str,
        help='The configuration name of the board')

args = parser.parse_args()
if args.debug: print args

scripts_dir = os.path.dirname(os.path.abspath(__file__))
scripts_parent_dir = os.path.dirname(scripts_dir)

path_args = (
    'bootloader',
    'data_dir',
    'flash_image',
    'work_dir',
)
for path_arg in path_args:
    if args.__dict__.has_key(path_arg) and args.__dict__[path_arg]:
        args.__dict__[path_arg] = os.path.abspath(args.__dict__[path_arg])

if not args.force_no_out_dir:
    out_tools_dir = os.path.join(scripts_parent_dir, '_out_tools')
    if os.path.exists(out_tools_dir):
        if args.debug: print 'Detected build tree; adding ' + out_tools_dir + ' to $PATH'
        os.environ['PATH'] = out_tools_dir + ':' + os.environ['PATH']

if not args.data_dir and not args.force_no_out_dir:
    out_data_dir = os.path.join(scripts_parent_dir, '_out')
    if os.path.exists(out_data_dir):
        if args.debug: print 'Detected build tree; using ' + out_data_dir + ' as data dir'
        args.data_dir = out_data_dir
if not args.data_dir:
    args.data_dir = '/usr/share/tegra_uboot_flasher'

load_configs(os.path.join(args.data_dir, 'configs'))

args.func()