#!/usr/bin/env python
# ==========================================================================
#
# Copyright (C) 2012-2014 RidgeRun, LLC (http://www.ridgerun.com)
#
# Authors: Jose Pablo Carballo <jose.carballo@ridgerun.com>
#          Diego Benavides <diego.benavides@ridgerun.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# The installer tool objective is to provide several ways to deploy the
# firmware to the target boards. 
#
# ==========================================================================

"""
Main entry point for the installer.

Please refer to the **User's Guide** for detailed information on how to call
this installer.
"""

# ==========================================================================
# Imports
# ==========================================================================

import sys
import argparse
import signal
import logging

from openfd.utils.args import ArgCheckerError
from openfd.boards import BoardFactory
from openfd.methods.board import *
from openfd.methods.sdcard import *
from openfd.methods.usb import *
from openfd.storage import DeviceException

# ==========================================================================
# Global variables
# ==========================================================================

_logger  = None

# ==========================================================================
# Constants
# ==========================================================================

# Modes
MODE_SD = 'sd'
MODE_SD_IMG = 'sd-img'
MODE_SD_SCRIPT = 'sd-script'
MODE_SD_SCRIPT_IMG = 'sd-script-img'
MODE_NAND = 'nand'
MODE_RAM = 'ram'
MODE_ENV = 'env'
MODE_USB_SCRIPT = 'usb-script'

# Components
COMP_IPL = "ipl"
COMP_BOOTLOADER = "bootloader"
COMP_KERNEL = "kernel"
COMP_FS = "fs"

# ==========================================================================
# Functions
# ==========================================================================

def _init_logging(args):
    global _logger
    _program_name = os.path.basename(sys.argv[0])
    _logger = utils.logger.init_global_logger(_program_name)
    _logger.setLevel(logging.DEBUG)
    streamhandler = logging.StreamHandler()
    streamhandler.setFormatter(logging.Formatter('%(msg)s'))
    if args.verbose:
        streamhandler.setLevel(logging.DEBUG)
    else:
        streamhandler.setLevel(logging.INFO)
    if args.quiet:
        streamhandler.setLevel(logging.CRITICAL)
    _logger.addHandler(streamhandler)
    if args.log_filename:
        filehandler = logging.FileHandler(args.log_filename, mode='a')
        filehandler.setLevel(logging.DEBUG)
        if args.verbose:
            filehandler.setFormatter(logging.Formatter('%(levelname)s:'
                                           '%(filename)s:%(lineno)s: %(msg)s'))
        else:
            filehandler.setFormatter(logging.Formatter('%(msg)s'))
        _logger.addHandler(filehandler)
        _logger.debug('-' * 80)
        _logger.debug('Log date: %s' % time.strftime("%Y%m%d %H:%M:%S"))
        command = ''
        for arg in sys.argv:
            if ' ' in arg:
                command += ' "%s"' % arg
            else:
                command += ' %s' % arg
        _logger.debug('Command: %s' % command)

def _init_executer(args):
    utils.executer.init_global_executer(dryrun=args.dryrun,
                                    enable_colors=True, verbose=args.verbose)

def _clean_exit(code=0):
    if code != 0: _logger.debug('Exiting with code %d' % code)
    exit(code)

def _abort_install():
    _logger.error('Installation aborted')
    _clean_exit(-1)

def _abort_install_user():
    _logger.error('Installation aborted by user interaction')
    _clean_exit(0)

def _sigint_handler(signal, frame):
    _logger.error('\nInstallation interrupted')
    _clean_exit(0)

def _get_args():
    parser = argparse.ArgumentParser()

    parser.add_argument('-y', '--assume-yes',
                       help='Automatic \'yes\' to prompts; '
                       'runs non-interactively',
                       dest='interactive',
                       action='store_false',
                       default=True)
    
    parser.add_argument('-v', '--verbose',
                       help='Enable debug',
                       dest='verbose',
                       action='store_true',
                       default=False)
    
    parser.add_argument('-q', '--quiet',
                       help='Be as quiet as possible',
                       dest='quiet',
                       action='store_true',
                       default=False)
    
    parser.add_argument('-l', '--log',
                       help="Log to file",
                       metavar='<file>',
                       dest='log_filename')
    
    parser.add_argument('--dryrun',
                       help='Sets the dryrun mode On (system and uboot '
                            'commands will be logged, but not executed)',
                       dest='dryrun',
                       action='store_true',
                       default=False)
    
    board_subparsers = parser.add_subparsers(help="board (--help available)",
                                             dest="board")
    for board_name in BoardFactory().supported_boards():
        board = BoardFactory().make(board_name)
        board_parser = board_subparsers.add_parser(board_name)
        board.add_args(board_parser)
    
    return parser.parse_args()
    
def _check_args(args):
    try:
        board = BoardFactory().make(args.board)
        board.check_args(args)
    except ArgCheckerError as e:
        _logger.error(e)
        _abort_install()

def _check_sudo():
    _logger.warning("This installation mode requires to execute commands via sudo")
    ret = utils.get_global_executer().prompt_sudo()
    if ret is False:
        _logger.error("Failed obtaining superuser access via sudo")
        _clean_exit(-1)

# ==========================================================================
# Installation modes
# ==========================================================================

def _get_uboot_telnet(args):
    uboot = UbootExpect()
    uboot.console_logger = _logger
    uboot.dryrun = args.dryrun
    try:
        cmd = 'termnet %s %s' % (args.telnet_host, args.telnet_port)
        ret = uboot.open_comm(cmd)
        if ret is False: _abort_install()
        ret = uboot.sync()
        if ret is False: _abort_install()
    except (pexpect.ExceptionPexpect, UbootTimeoutException) as e:
        uboot.close_comm()
        _logger.error(e)
        _abort_install()
    return uboot

def _get_uboot_serial(args):
    uboot = Uboot()
    uboot.serial_logger = _logger
    uboot.dryrun = args.dryrun
    uboot.uboot_mode(mode=args.uboot_comm_mode)

    try:
        ret = uboot.open_comm(port=args.serial_port, 
                              baud=args.serial_baud, 
                              termnet_host=args.termnet_host, 
                              termnet_port=args.termnet_port)

        if ret is False: _abort_install()
        ret = uboot.sync()
        if ret is False: _abort_install()
    except (serial.SerialException, UbootTimeoutException) as e:
        uboot.close_comm()
        _logger.error(e)
        _abort_install()
    return uboot

def _get_uboot(args):
    if args.serial_port:
        return _get_uboot_serial(args)
    elif args.telnet_host:
        return _get_uboot_telnet(args)

def _get_tftp_loader(args, uboot):
    tftp_loader = TftpRamLoader(uboot, args.board_net_mode)
    tftp_loader.dir = args.tftp_dir
    tftp_loader.port = args.tftp_port
    tftp_loader.host_ipaddr = args.host_ip_addr
    tftp_loader.net_mode = args.board_net_mode
    if args.board_net_mode == TftpRamLoader.MODE_STATIC:
        tftp_loader.board_ipaddr = args.board_ip_addr
    tftp_loader.dryrun = args.dryrun
    return tftp_loader

def _mode_sd(args):
    _check_sudo()
    try:
        board = BoardFactory().make(args.board)
        board.sd_init_comp_installer(args)
        sd_installer = SDCardInstaller(board=board)
        sd_installer.interactive = args.interactive
        sd_installer.dryrun = args.dryrun
        sd_installer.device = args.device
        sd_installer.read_partitions(args.mmap_file)
        sd_installer.format()
        sd_installer.mount_partitions(args.workdir)
        sd_installer.install_components()
        sd_installer.release()
    except (SDCardInstallerError, SDCardInstallerCanceled, DeviceException, RamLoaderException) as e:
	if str(e) == 'User canceled':
            _abort_install_user()
        else:
            _logger.error(e)
            _abort_install()

def _mode_sd_img(args):
    _check_sudo()
    try:
        board = BoardFactory().make(args.board)
        board.sd_init_comp_installer(args)
        ld_installer = LoopDeviceInstaller(board=board)
        ld_installer.dryrun = args.dryrun
        ld_installer.read_partitions(args.mmap_file)
        ld_installer.format(args.image, args.imagesize_mb)
        ld_installer.mount_partitions(args.workdir)
        ld_installer.install_components()
        ld_installer.release()
        _logger.info("<hint>")
        _logger.info("  You can use the image file to flash an SD card:")
        _logger.info("    1. Plug your SD card (say it is device /dev/sdX)")
        _logger.info("    2. Unmount it ('sudo umount' should help)")
        _logger.info("    3. Flash: sudo dd bs=1M if=%s of=/dev/sdX" % args.image)
        _logger.info("</hint>")
    except (LoopDeviceInstallerError, SDCardInstallerCanceled, DeviceException, RamLoaderException) as e:
	if str(e) == 'User canceled':
            _abort_install_user()
        else:
            _logger.error(e)
            _abort_install()
    
def _mode_nand(args):
    uboot = _get_uboot(args)
    tftp_loader = _get_tftp_loader(args, uboot)
    board = BoardFactory().make(args.board)
    try:
        tftp_loader.setup_uboot_network()
        nand_installer = NandInstaller(uboot=uboot, board=board,
                                       loader=tftp_loader)
        if args.nand_blk_size:
            nand_installer.nand_block_size = args.nand_blk_size
        if args.nand_page_size:
            nand_installer.nand_page_size = args.nand_page_size
        nand_installer.ram_load_addr = args.ram_load_addr
        nand_installer.dryrun = args.dryrun
        nand_installer.read_partitions(args.mmap_file)
        if args.nand_uboot_file:
            nand_installer.load_uboot_to_ram(args.nand_uboot_file,
                                             args.ram_load_addr)
            tftp_loader.setup_uboot_network()
        if args.component == COMP_IPL:
            nand_installer.install_ipl(force=args.ipl_force)
        if args.component == COMP_BOOTLOADER:
            nand_installer.install_bootloader()
        if args.component == COMP_KERNEL:
            nand_installer.install_kernel(force=args.kernel_force)
        if args.component == COMP_FS:
            nand_installer.install_fs(force=args.fs_force)
        if uboot.get_env('autostart') == 'no':
            uboot.set_env('autostart', 'yes')
            uboot.save_env()
        uboot.cmd('echo Installation complete', prompt_timeout=None)
    except (UbootTimeoutException, RamLoaderException, NandInstallerError) as e:
        uboot.close_comm()
        _logger.error(e)
        _abort_install()
    uboot.close_comm()

def _mode_ram(args):
    uboot = _get_uboot(args)
    tftp_loader = _get_tftp_loader(args, uboot)
    try:
        tftp_loader.setup_uboot_network()
        _logger.info("Loading %s to RAM address %s" %
                     (args.ram_file, args.ram_load_addr))
        tftp_loader.load_file_to_ram_and_boot(args.ram_file, args.ram_load_addr,
                      args.ram_boot_line, boot_timeout=args.ram_boot_timeout)
    except (UbootTimeoutException, RamLoaderException) as e:
        uboot.close_comm()
        _logger.error(e)
        _abort_install()
    uboot.close_comm()

def _mode_env(args):
    uboot = _get_uboot(args)
    try:
        env_installer = EnvInstaller(uboot=uboot)
        env_installer.dryrun = args.dryrun
        env_installer.install_variable(args.env_variable, args.env_value,
                                       args.env_force)
        uboot.cmd('echo Installation complete', prompt_timeout=None)
    except UbootTimeoutException as e:
        uboot.close_comm()
        _logger.error(e)
        _abort_install()
    uboot.close_comm()

def _mode_sd_script(args):
    _check_sudo()
    board = BoardFactory().make(args.board)
    ext_nand_installer = NandExternalInstaller(board=board)
    ext_nand_installer.read_partitions(args.flash_mmap_file)
    ext_nand_installer.install_boardinfo()
    ext_nand_installer.install_ipl()
    ext_nand_installer.install_bootloader()
    ext_nand_installer.install_kernel()
    ext_nand_installer.install_fs()
    ext_nand_installer.write(args.template_file, args.output_file)
    try:
        board.sd_init_comp_installer(args)
        ext_sd_installer = SDCardExternalInstaller(board=board)
        ext_sd_installer.interactive = args.interactive
        ext_sd_installer.dryrun = args.dryrun
        ext_sd_installer.device = args.device
        ext_sd_installer.read_partitions(args.sd_mmap_file)
        ext_sd_installer.format()
        ext_sd_installer.mount_partitions(args.workdir)
        imgs = ext_nand_installer.get_imgs()
        ext_sd_installer.install_components(args.workdir, imgs,
                                            args.mkimage_bin, args.output_file)
        ext_sd_installer.release()        
    except (SDCardInstallerError, DeviceException) as e:
        _logger.error(e)
        _abort_install()

def _mode_sd_script_img(args):
    _check_sudo()
    board = BoardFactory().make(args.board)
    ext_nand_installer = NandExternalInstaller(board=board)
    ext_nand_installer.read_partitions(args.flash_mmap_file)
    ext_nand_installer.install_boardinfo()
    ext_nand_installer.install_ipl()
    ext_nand_installer.install_bootloader()
    ext_nand_installer.install_kernel()
    ext_nand_installer.install_fs()
    ext_nand_installer.write(args.template_file, args.output_file)
    try:
        board.sd_init_comp_installer(args)
        ld_installer = LoopDeviceExternalInstaller(board=board)
        ld_installer.dryrun = args.dryrun
        ld_installer.read_partitions(args.sd_mmap_file)
        ld_installer.format(args.image, args.imagesize_mb)
        ld_installer.mount_partitions(args.workdir)
        imgs = ext_nand_installer.get_imgs()
        ld_installer.install_components(args.workdir, imgs,
                                        args.mkimage_bin, args.output_file)
        ld_installer.release()
        _logger.info("<hint>")
        _logger.info("  You can use the image file to flash an SD card:")
        _logger.info("    1. Plug your SD card (say it is device /dev/sdX)")
        _logger.info("    2. Unmount it ('sudo umount' should help)")
        _logger.info("    3. Flash: sudo dd bs=1M if=%s of=/dev/sdX" % args.image)
        _logger.info("</hint>")
    except (LoopDeviceInstallerError, DeviceException) as e:
        _logger.error(e)
        _abort_install()

def _mode_usb_script(args):
    _check_sudo()
    board = BoardFactory().make(args.board)
    ext_nand_installer = NandExternalInstaller(board=board)
    ext_nand_installer.read_partitions(args.flash_mmap_file)
    ext_nand_installer.install_boardinfo()
    ext_nand_installer.install_ipl()
    ext_nand_installer.install_bootloader()
    ext_nand_installer.install_kernel()
    ext_nand_installer.install_fs()
    ext_nand_installer.write(args.template_file, args.output_file)
    try:
        board.usb_init_comp_installer(args)
        usb_installer = USBInstaller(board=board)
        usb_installer.interactive = args.interactive
        usb_installer.dryrun = args.dryrun
        usb_installer.device = args.device
        usb_installer.read_partitions(args.usb_mmap_file)
        usb_installer.format()
        usb_installer.mount_partitions(args.workdir)
        imgs = ext_nand_installer.get_imgs()
        usb_installer.install_components(args.workdir, imgs,
                                            args.mkimage_bin, args.output_file)
        usb_installer.release()        
    except (USBInstallerError, DeviceException) as e:
        _logger.error(e)
        _abort_install()


# ==========================================================================
# Main logic
# ==========================================================================

def main():
    signal.signal(signal.SIGINT, _sigint_handler)
    signal.signal(signal.SIGTERM, _sigint_handler)
    args = _get_args()
    _init_logging(args)
    _init_executer(args)
    _check_args(args)
    try:
        board = BoardFactory().make(args.board)
        board.check_args(args)
    except ArgCheckerError as e:
        _logger.error(e)
        _abort_install()
    if args.mode == MODE_SD:
        _mode_sd(args)
    if args.mode == MODE_SD_IMG:
        _mode_sd_img(args)
    if args.mode == MODE_NAND:
        _mode_nand(args)
    if args.mode == MODE_RAM:
        _mode_ram(args)
    if args.mode == MODE_ENV:
        _mode_env(args)
    if args.mode == MODE_SD_SCRIPT:
        _mode_sd_script(args)
    if args.mode == MODE_SD_SCRIPT_IMG:
        _mode_sd_script_img(args)
    if args.mode == MODE_USB_SCRIPT:
        _mode_usb_script(args)
    _logger.info('Installation complete')
    _clean_exit(0)
    
if __name__ == '__main__':
    main()
