#!/usr/bin/env python

# Copyright (c) 2013, 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 os
import os.path
import shutil
import stat
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(dir, cmd):
    oldcwd = os.getcwd()
    print '+ cd', dir
    os.chdir(dir)
    print '+', cmd
    ret = os.system(cmd)
    if ret:
        raise Exception('Command failed: %d' % ret)
    os.chdir(oldcwd)

def gen_flashcmd_mmc(flash_image_addr, flash_img_size):
    flash_id = config['flash-id-uboot']
    flash_img_size_sectors = flash_img_size / 512
    flashcmd = 'mmc dev %d 1 ; ' % flash_id
    flashcmd += 'mmc write 0x%08x 0 0x%x ; ' % (flash_image_addr, flash_img_size_sectors)
    return flashcmd

def gen_flashcmd_nand(flash_image_addr, flash_img_size):
    flashcmd = 'nand erase.chip ; '
    flashcmd += 'nand write 0x%08x 0 0x%08x ; ' % (flash_image_addr, flash_img_size)
    return flashcmd

def gen_flashcmd_spi(flash_image_addr, flash_img_size):
    flashcmd = 'sf probe 0 ; '
    flashcmd += 'sf erase 0 0x%08x ; ' % config['flash-erase-size']
    flashcmd += 'sf write 0x%08x 0 0x%08x ; ' % (flash_image_addr, flash_img_size)
    return flashcmd

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

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_img
    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)

    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)

    # Add 32k in case 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 + (32 * 1024) + (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)

    # 0x00108000 is CONFIG_SYS_TEXT_BASE in U-Boot, minus RAM base
    loadaddr = soc['ram-base'] + 0x00108000
    flash_image_addr = loadaddr + padded_size
    if args.debug:
        print 'flash_image_addr %d 0x%x' % (flash_image_addr, flash_image_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:
        workdir = os.path.abspath(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)

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

        bootcmd = ''
        if args.debug:
            bootcmd = 'crc32 0x%08x 0x%08x ; ' % (flash_image_addr, flash_img_size)
        bootcmd += gen_flashcmd(flash_image_addr, flash_img_size)
        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'] + ' ; '
        bootcmd += 'saveenv ; '
        # To update the bootloader, reset.
        # If wanting to run installer, set installer_args.configname in environment, 'run bootcmd'
        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 = 'tegrarcm --bct=' + bct + ' --bootloader=' + uboot_flasher + ' --loadaddr=0x%08x' % 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(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)

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('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)

if not args.force_no_out_dir:
    out_tools_dir = os.path.abspath(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.abspath(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:
    data_dir = '/usr/share/tegra_uboot_flasher'

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

args.func()