#!/usr/bin/env vpython3
# Copyright 2022 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Implements commands for flashing a Fuchsia device."""

import argparse
import logging
import os
import subprocess
import sys

from typing import Optional, Tuple

import common
from boot_device import BootMode, StateTransitionError, boot_device
from common import get_system_info, find_image_in_sdk, register_device_args
from compatible_utils import get_sdk_hash, running_unattended
from lockfile import lock

# Flash-file lock. Used to restrict number of flash operations per host.
# File lock should be marked as stale after 15 mins.
_FF_LOCK = os.path.join('/tmp', 'flash.lock')
_FF_LOCK_STALE_SECS = 60 * 15
_FF_LOCK_ACQ_TIMEOUT = _FF_LOCK_STALE_SECS


def _get_system_info(
    target: Optional[str], serial_num: Optional[str]
) -> Tuple[str, str]:
    """Retrieves installed OS version from device.

    Args:
        target: Target to get system info of.
        serial_num: Serial number of device to get system info of.
    Returns:
        Tuple of strings, containing (product, version number).
    """

    if running_unattended():
        try:
            boot_device(target, BootMode.REGULAR, serial_num)
        except (subprocess.CalledProcessError, StateTransitionError):
            logging.warning('Could not boot device. Assuming in fastboot')
            return ('', '')

    return get_system_info(target)


def _update_required(
    os_check,
    system_image_dir: Optional[str],
    target: Optional[str],
    serial_num: Optional[str] = None,
) -> Tuple[bool, Optional[str]]:
    """Returns True if a system update is required and path to image dir."""

    if os_check == 'ignore':
        return False, system_image_dir
    if not system_image_dir:
        raise ValueError('System image directory must be specified.')
    if not os.path.exists(system_image_dir):
        logging.warning(
            'System image directory does not exist. Assuming it\'s '
            'a product-bundle name and dynamically searching for '
            'image directory'
        )
        path = find_image_in_sdk(system_image_dir)
        if not path:
            raise FileNotFoundError(
                f'System image directory {system_image_dir} could notbe found'
            )
        system_image_dir = path
    if os_check == 'check' and get_sdk_hash(
        system_image_dir
    ) == _get_system_info(target, serial_num):
        return False, system_image_dir
    return True, system_image_dir


def _run_flash_command(system_image_dir: str, target_id: Optional[str]):
    """Helper function for running `ffx target flash`."""
    logging.info('Flashing %s to %s', system_image_dir, target_id)
    # Flash only with a file lock acquired.
    # This prevents multiple fastboot binaries from flashing concurrently,
    # which should increase the odds of flashing success.
    with lock(_FF_LOCK, timeout=_FF_LOCK_ACQ_TIMEOUT):
        # The ffx.fastboot.inline_target has negative impact when ffx
        # discovering devices in fastboot, so it's inlined here to limit its
        # scope. See the discussion in https://fxbug.dev/issues/317228141.
        logging.info(
            'Flash result %s',
            common.run_ffx_command(
                cmd=(
                    'target',
                    'flash',
                    '-b',
                    system_image_dir,
                    '--no-bootloader-reboot',
                ),
                target_id=target_id,
                configs=['ffx.fastboot.inline_target=true'],
                capture_output=True,
            ).stdout,
        )


def update(
    system_image_dir: str,
    os_check: str,
    target: Optional[str],
    serial_num: Optional[str] = None,
) -> None:
    """Conditionally updates target given.

    Args:
        system_image_dir: string, path to image directory.
        os_check: <check|ignore|update>, which decides how to update the device.
        target: Node-name string indicating device that should be updated.
        serial_num: String of serial number of device that should be updated.
    """
    needs_update, actual_image_dir = _update_required(
        os_check, system_image_dir, target, serial_num
    )
    logging.info(
        'update_required %s, actual_image_dir %s',
        needs_update,
        actual_image_dir,
    )
    if not needs_update:
        return
    if serial_num:
        boot_device(target, BootMode.BOOTLOADER, serial_num)
        _run_flash_command(system_image_dir, serial_num)
    else:
        _run_flash_command(system_image_dir, target)


def register_update_args(
    arg_parser: argparse.ArgumentParser,
    default_os_check: Optional[str] = 'check',
) -> None:
    """Register common arguments for device updating."""
    serve_args = arg_parser.add_argument_group(
        'update', 'device updating arguments'
    )
    serve_args.add_argument(
        '--system-image-dir',
        help='Specify the directory that contains the '
        'Fuchsia image used to flash the device. Only '
        'needs to be specified if "os_check" is not '
        '"ignore".',
    )
    serve_args.add_argument(
        '--serial-num',
        default=os.environ.get('FUCHSIA_FASTBOOT_SERNUM'),
        help='Serial number of the device. Should be '
        'specified for devices that do not have an image '
        'flashed.',
    )
    serve_args.add_argument(
        '--os-check',
        choices=['check', 'update', 'ignore'],
        default=default_os_check,
        help='Sets the OS version enforcement policy. If '
        '"check", then the deployment process will halt '
        'if the target\'s version does not match. If '
        '"update", then the target device will '
        'be reflashed. If "ignore", then the OS version '
        'will not be checked.',
    )


def main():
    """Stand-alone function for flashing a device."""
    parser = argparse.ArgumentParser()
    register_device_args(parser)
    register_update_args(parser, default_os_check='update')
    args = parser.parse_args()
    update(
        args.system_image_dir, args.os_check, args.target_id, args.serial_num
    )


if __name__ == '__main__':
    sys.exit(main())