#!/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 multiprocessing
import os
import shutil
import subprocess
import sys
from tegraboardconfigs import *

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

out_tools_dir = os.path.abspath(os.path.join(scripts_parent_dir, '_out_tools'))
if os.path.exists(out_tools_dir):
    os.environ['PATH'] = out_tools_dir + ':' + os.environ['PATH']

if not os.environ.has_key('CROSS_COMPILE'):
    os.environ['CROSS_COMPILE'] = 'arm-linux-gnueabi-'

cbootimage_configs_dir = os.path.join(scripts_parent_dir, 'cbootimage-configs')
uboot_dir = os.path.join(scripts_parent_dir, 'u-boot')

def gen_cbimage_dir(boardname):
    return os.path.join(cbootimage_configs_dir, boards[boardname]['soc'], boards[boardname]['vendor'], boardname)

out_dir = os.path.join(scripts_parent_dir, '_out')
def gen_out_board_dir(boardname):
    return os.path.join(out_dir, boardname)

build_dir = os.path.join(scripts_parent_dir, '_build')
def gen_build_board_dir(boardname):
    return os.path.join(build_dir, boardname)
def gen_build_uboot_dir(build_board_dir):
    return os.path.join(build_board_dir, 'u-boot')

makejobs = '-j' + str(multiprocessing.cpu_count() + 1)

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 rm(fn):
    if (os.path.exists(fn)):
        os.unlink(fn)

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 all_enabled_confignames():
    for configname in configs.keys():
        if configs[configname].get('disabled', False):
            continue
        yield configname

def all_enabled_confignames_for_boardname(boardname):
    for configname in all_enabled_confignames():
        if configs[configname]['board'] != boardname:
            continue
        yield configname

def all_enabled_boardnames():
    seen = {}
    for configname in all_enabled_confignames():
        boardname = configs[configname]['board']
        if seen.has_key(boardname):
            continue
        seen[boardname] = True
        yield boardname

def user_restrict_socs(enabled_socs):
    enabled_socs = enabled_socs.split(',')
    for socname in socs.keys():
        socs[socname]['disabled'] = not socname in enabled_socs

def user_restrict_boards(enabled_boards):
    if enabled_boards[0] == '+':
        disable_others = False
        enabled_boards = enabled_boards[1:]
    else:
        disable_others = True
    enabled_boards = enabled_boards.split(',')
    for boardname in boards.keys():
        if boardname in enabled_boards:
            boards[boardname]['disabled'] = False
        elif disable_others:
            boards[boardname]['disabled'] = True

def restrict_boards():
    for board in boards.values():
        if socs[board['soc']].get('disabled', False):
            board['disabled'] = True

def restrict_configs():
    for config in configs.values():
        if boards[config['board']].get('disabled', False):
            config['disabled'] = True

def dtb_filename(config):
    extra = configs[config]['dtbfn-extra']
    boardname = configs[config]['board']
    return boards[boardname]['soc'] + '-' + boardname + extra + '.dtb'

def import_uboot_one_board(boardname, build_uboot_dir):
    out_board_dir = gen_out_board_dir(boardname)
    mkdir(out_board_dir)

    uboot = os.path.join(build_uboot_dir, 'u-boot')
    cmd = [os.environ['CROSS_COMPILE'] + 'size', '-A', uboot]
    lines = subprocess.check_output(cmd)
    bss_size = None
    for line in lines.split('\n'):
        fields = line.split()
        if fields[0] == '.bss':
            bss_size = fields[1]
            break
    if not bss_size:
        raise Exception('Could not find bss size in `size`')
    dst = os.path.join(out_board_dir, 'u-boot-bss-size')
    f = open(dst, 'wt')
    f.write(bss_size)
    f.close()

    get_entry = os.path.join(scripts_dir, 'elf-get-entry')
    uboot_spl = os.path.join(build_uboot_dir, 'spl', 'u-boot-spl')
    dst = os.path.join(out_board_dir, 'u-boot-spl-entry')
    with open(dst, "wt") as dst_f:
        cmd = [get_entry, uboot_spl]
        subprocess.check_call(cmd, stdout=dst_f)

    src = os.path.join(build_uboot_dir, 'u-boot-nodtb-tegra.bin')
    dst = os.path.join(out_board_dir, 'u-boot-nodtb-tegra.bin')
    cp(src, dst)

    src = os.path.join(build_uboot_dir, 'u-boot.dtb')
    if not os.path.exists(src):
        src = os.path.join(build_uboot_dir, 'dts/dt.dtb')
    dst = os.path.join(out_board_dir, 'u-boot.dtb')
    cp(src, dst)

    src = os.path.join(build_uboot_dir, 'u-boot-dtb-tegra.bin')
    dst = os.path.join(out_board_dir, 'u-boot-dtb-tegra.bin')
    cp(src, dst)

def cmd_import_uboot():
    boardname = None
    for i in all_enabled_boardnames():
        if boardname:
            raise Exception('import-uboot only allows one enabled board')
        boardname = i
    import_uboot_one_board(boardname, args.builddir)

def build_uboot_one_board(boardname):
    build_board_dir = gen_build_board_dir(boardname)
    build_uboot_dir = gen_build_uboot_dir(build_board_dir)
    mkdir(build_uboot_dir)

    out_board_dir = gen_out_board_dir(boardname)
    mkdir(out_board_dir)

    # U-Boot commit 51148790f26e "kconfig: switch to Kconfig" changed the
    # build commands from "make ${board}_config; make" to
    # "make ${board}_defconfig; make". That commit also introduced some new
    # text into the "make help" output. Check for that text here in order to
    # tell how to compile U-Boot:-/
    ret = subprocess.call('make help 2>/dev/null | grep -q "Configuration targets"', shell=True, cwd=uboot_dir)
    if ret:
        config_target_suffix = '_config'
    else:
        config_target_suffix = '_defconfig'
    run(uboot_dir, 'make O=' + build_uboot_dir + ' ' + boardname + config_target_suffix)
    run(uboot_dir, 'make O=' + build_uboot_dir + ' -s ' + makejobs)

    import_uboot_one_board(boardname, build_uboot_dir)

def cmd_build_uboots():
    for boardname in all_enabled_boardnames():
        build_uboot_one_board(boardname)

def build_bct_img_one_board(boardname):
    cbimage_dir = gen_cbimage_dir(boardname)

    out_board_dir = gen_out_board_dir(boardname)
    mkdir(out_board_dir)

    src = os.path.join(out_board_dir, 'u-boot-dtb-tegra.bin')
    dst = os.path.join(cbimage_dir, 'u-boot.bin')
    cp(src, dst)
    run(cbimage_dir, './build.sh')

    for configname in all_enabled_confignames_for_boardname(boardname):
        bct_fn = configs[configname]['bct']
        src = os.path.join(cbimage_dir, bct_fn)
        dst = os.path.join(out_board_dir, bct_fn)
        cp(src, dst)

        img_fn = configs[configname]['flash-image']
        src = os.path.join(cbimage_dir, img_fn)
        dst = os.path.join(out_board_dir, img_fn)
        cp(src, dst)

def cmd_build_bcts_imgs():
    for boardname in all_enabled_boardnames():
        build_bct_img_one_board(boardname)

def cmd_build_configs():
    mkdir(os.path.join(out_dir, 'configs'))
    boardnames = {}
    for configname in all_enabled_confignames():
        fn = configname + '.config'
        src = os.path.join(scripts_dir, 'configs', fn)
        dst = os.path.join(out_dir, 'configs', fn)
        cp(src, dst)
        boardnames[configs[configname]['board']] = True
    socnames = {}
    for boardname in boardnames:
        fn = boardname + '.board'
        src = os.path.join(scripts_dir, 'configs', fn)
        dst = os.path.join(out_dir, 'configs', fn)
        cp(src, dst)
        socnames[boards[boardname]['soc']] = True
    for socname in socnames:
        fn = socname + '.soc'
        src = os.path.join(scripts_dir, 'configs', fn)
        dst = os.path.join(out_dir, 'configs', fn)
        cp(src, dst)

def cmd_build():
    cmd_build_uboots()
    cmd_build_bcts_imgs()
    cmd_build_configs()

def cmd_help():
    print 'usage: build <command>'
    for cmd in sorted(cmdmap.keys()):
        print ' ', cmd

parser = argparse.ArgumentParser(description='Build U-Boot, BCT, and flash ' +
    'images for Tegra boards.')

parser.add_argument('--socs', type=str,
    help='Restrict the build to a (comma-separated) list of SoCs.')
parser.add_argument('--boards', type=str,
    help='Restrict the build to a (comma-separated) list of boards.')

subparsers = parser.add_subparsers()

subparser = subparsers.add_parser('import-uboot',
    help='Import U-Boot binaries from an external build tree')
subparser.add_argument('builddir', type=str,
    help='The external build directory to import from.')
subparser.set_defaults(func = cmd_import_uboot)

subparser = subparsers.add_parser('build-uboots',
    help='Build U-Boot binaries')
subparser.set_defaults(func = cmd_build_uboots)

subparser = subparsers.add_parser('build-bcts-imgs',
    help='Build BCT and flash images')
subparser.set_defaults(func = cmd_build_bcts_imgs)

subparser = subparsers.add_parser('build-configs',
    help='Build config files')
subparser.set_defaults(func = cmd_build_configs)

subparser = subparsers.add_parser('build',
    help='Build everything')
subparser.set_defaults(func = cmd_build)

if __name__ == '__main__':
    args = parser.parse_args()
    load_configs(os.path.join(scripts_dir, 'configs'))
    if args.socs:
        user_restrict_socs(args.socs)
    if args.boards:
        user_restrict_boards(args.boards)
    restrict_boards()
    restrict_configs()
    args.func()