#!/usr/bin/env python3

import argparse
import logging
import os
import shutil
import sys
import time
import traceback

from firmware_update import download_file, get_installation_candidate
from logging.handlers import SysLogHandler
from utils import get_command_output, run_command, extract_to

DEVID_BIN = '/usr/sbin/devid'
HARDWARE_NAME = 'mrc-signalq2'
ROOTFS_SLOT = 'rootfs'
LOG_DIR = '/media/sd/logs'
SYSLOG_TAG = 'signalq2_updater'

logging.basicConfig(format='%(levelname)s: %(message)s', level=logging.DEBUG)
logger = logging.getLogger('mrc-drive-updater')
syslog_handler = SysLogHandler('/dev/log', facility=SysLogHandler.LOG_LOCAL0)
syslog_handler.setLevel(logging.DEBUG)
formatter = logging.Formatter(SYSLOG_TAG + ":%(levelname)s: %(message)s")
syslog_handler.setFormatter(formatter)
logger.addHandler(syslog_handler)


def get_device_id():
    return get_command_output([DEVID_BIN, '--serial'])


def get_current_rtos_version():
    return get_command_output([DEVID_BIN, '--rtos_ver'])


def get_current_linux_version():
    return get_command_output([DEVID_BIN, '--linux_ver'])


def get_current_mrc_version():
    return get_command_output([DEVID_BIN, '--mrc_ver'])


def get_firmware_version():
    '''Firmware version is composed from linux version and mrc version'''

    # linux version example: 0.0.3-2.linux-yandex_2.1-dev_69883f2d_202011271412_master
    linux_version_full = get_current_linux_version()
    linux_version_short = linux_version_full.split('_')[0]
    mrc_version = get_current_mrc_version()
    return linux_version_short + '_' + mrc_version


def get_user_agent():
    rtos_ver = get_current_rtos_version()
    linux_ver = get_current_linux_version()
    mrc_ver = get_current_mrc_version()
    return 'rtos/{} linux/{} mrc/{}'.format(rtos_ver, linux_ver, mrc_ver)


def download_and_install(firmware):
    try:
        download_dir = '/data/download/updates'
        firmware_dir = os.path.join(download_dir, 'firmware-{}'.format(firmware.version))
        install_dir = '/media/emmc'

        shutil.rmtree(download_dir, ignore_errors=True)
        os.makedirs(download_dir, mode=0o755, exist_ok=True)

        logger.info('Download firmware update to dir: {}'.format(download_dir))
        archive_path = os.path.join(download_dir, 'firmware-{}.tar.gz'.format(firmware.version))
        download_file(firmware.url, archive_path, firmware.md5sum)

        os.makedirs(firmware_dir, mode=0o755, exist_ok=True)

        logger.info('Extract firmware archive {} to {}'.format(archive_path, firmware_dir))
        extract_to(archive_path, firmware_dir)

        # Extracted tar.gz should contain firmware and signature files
        logger.info('Copy firmware files to {}'.format(install_dir))
        for file in os.listdir(firmware_dir):
            if file.endswith('.bin'):
                shutil.move(os.path.join(firmware_dir, file),
                            os.path.join(install_dir, 'yandex.bin'))
            if file.endswith('.sig'):
                shutil.move(os.path.join(firmware_dir, file),
                            os.path.join(install_dir, 'yandex.bin.sha256.signed'))

        logger.info('Raise fwupdate_flag')
        return run_command(['yandex_api', 'fwupdate_flag', 'set'])
    finally:
        shutil.rmtree(firmware_dir, ignore_errors=True)


def do_firmware_update():
    logger.info('Start firmware update')

    firmware = get_installation_candidate(get_user_agent(), HARDWARE_NAME, get_device_id(), ROOTFS_SLOT)
    if not firmware:
        logger.info('No firmware update found, quit')
        return
    logger.info('Installation candidate: {}'.format(firmware))

    current_version = get_firmware_version()
    logger.info('Current version: {}'.format(current_version))

    if firmware.version == current_version:
        logger.info('Target firmware version is already installed, quit')
        return

    success = download_and_install(firmware)

    if success:
        logger.info('Firmware update done, reboot')
        run_command(['/sbin/reboot', 'now'])


def daemonize():
    if os.fork() != 0:
        exit()

    os.setsid()

    # Ensure that fork process is not a session leader
    if os.fork() != 0:
        exit()

    os.umask(0)
    os.chdir("/")

    sys.stdout.flush()
    sys.stderr.flush()

    devnull = open("/dev/null", "r+")
    os.dup2(devnull.fileno(), sys.stdin.fileno())
    os.dup2(devnull.fileno(), sys.stdout.fileno())
    os.dup2(devnull.fileno(), sys.stderr.fileno())


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--update_interval', type=int, required=False, default=300,
                        help='Update interval in seconds')

    parser.add_argument('-d', '--daemon', help='Run program as daemon', action='store_true')

    args = parser.parse_args()

    if args.daemon:
        daemonize()
        os.makedirs(LOG_DIR, mode=0o755, exist_ok=True)

    while True:
        try:
            logger.info('Start update')
            do_firmware_update()
        except:
            logger.error(traceback.format_exc())

        time.sleep(args.update_interval)


if __name__ == '__main__':
    main()
