#!/usr/bin/env python3.8
# Copyright 2022 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import argparse
import sys
import subprocess
import pathlib
import os
import json
import unittest
from typing import List

# These are variables will be substituted with the real values at run time.
# We do not know where the test artifacts are stored until after `ffx test`
# completes. So we offer this place holder to substitute the values once they
# are evaluated.
VARIABLE_SUBSTITUTION = {"{test_artifact_dir}": ""}


def run_target_test(
        ffx_bin: str, test_url: str, outdir: str,
        ffx_test_args: List[str]) -> subprocess.CompletedProcess:
    """Runs 'ffx test run <url> --output-directory outdir [args]'"""
    root_build_dir = os.getcwd()

    # use the same configuration in //src/developer/ffx/build/ffx_action.gni
    base_config = [
        'analytics.disabled=true',
        'assembly_enabled=true',
        'sdk.root=' + root_build_dir,
        'sdk.type=in-tree',
        'sdk.module=host_tools.modular',
    ]

    args = [ffx_bin]
    for c in base_config:
        args += ['--config', c]
    args += [
        '--env', './.ffx.env', 'test', 'run', test_url, '--output-directory',
        outdir
    ]
    args += ffx_test_args
    return subprocess.run(args, stderr=sys.stdout)


def get_test_suite_artifact_dir(
        output_root: pathlib.Path, test_suite_name: str):
    """Parses run_summary.json and reads back 'artifact_dir'.

    Raises:
       KeyError: if json keys for parsing test artifact directory does not exist.

    Returns:
       subdir path where test suite artifact path is stored.
    """
    summary_json = os.path.join(output_root, 'run_summary.json')
    with open(summary_json, 'r') as f:
        test_summary = json.load(f)
        test_summary = test_summary['data']
        for suite in test_summary['suites']:
            if test_suite_name in suite['name']:
                return suite['artifact_dir']
    return ""


def do_variable_substitution(input: str) -> str:
    """Substitute input with values defined in VARIABLE_SUBSTITUTION."""
    for var in VARIABLE_SUBSTITUTION:
        input = input.replace(var, VARIABLE_SUBSTITUTION[var])
    return input


def run_host_script(
        script_bin: str, script_args: List[str]) -> subprocess.CompletedProcess:
    """Runs the host script."""
    substituted_args = []
    for arg in script_args:
        substituted_args.append(do_variable_substitution(arg))
    return subprocess.run([script_bin] + substituted_args, stderr=sys.stdout)


class RunTargetTestWithHostScript(unittest.TestCase):
    """Executes the target test via `ffx test`, then executes the host script in tearDown."""

    def __init__(self, args=None, test_output_dir=None):
        super().__init__()
        self.args = args
        self.test_output_dir = test_output_dir
        self.test_status = None

    def runTest(self):
        """Runs test on target"""
        self.test_status = run_target_test(
            self.args.ffx_bin, self.args.test_url, self.test_output_dir,
            self.args.ffx_test_args)
        self.assertEqual(
            self.test_status.returncode, 0,
            "test returned non-zero status, error: %s" %
            self.test_status.stderr)

    def tearDown(self):
        if self.args.host_script_bin and (self.test_status.returncode == 0 or
                                          self.args.run_host_script_on_fail):
            artifact_subdir = get_test_suite_artifact_dir(
                self.test_output_dir, self.args.test_url)
            VARIABLE_SUBSTITUTION['{test_artifact_dir}'] = os.path.join(
                self.test_output_dir, artifact_subdir)
            script_status = run_host_script(
                self.args.host_script_bin, self.args.host_script_args)
            self.assertEqual(
                script_status.returncode, 0,
                "host script returned non-zero status, error: %s" %
                script_status.stderr)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--ffx-bin',
        type=pathlib.Path,
        required=True,
        help='Path to ffx binary.')
    parser.add_argument('--test-url', type=str, required=True, help='Test URL')
    parser.add_argument(
        '--test-outdir', type=pathlib.Path, help='Dir to store test outputs')
    parser.add_argument(
        '--ffx-test-args',
        action='append',
        default=[],
        help='List of args to be forwarded to ffx test')
    parser.add_argument(
        '--host-script-bin',
        type=pathlib.Path,
        help='Path to host script, optional')
    parser.add_argument(
        '--host-script-args',
        action='append',
        type=str,
        default=[],
        help='List of args to be forwarded to host script')
    parser.add_argument(
        '--run-host-script-on-fail',
        help='whether to run the host script if test fails')
    args = parser.parse_args()
    test_output_dir = os.getenv('FUCHSIA_TEST_OUTDIR')
    if not test_output_dir:
        test_output_dir = args.test_outdir
    print('storing test output to: ', test_output_dir)
    test = RunTargetTestWithHostScript(
        args=args, test_output_dir=test_output_dir)
    result = unittest.TextTestRunner().run(test)
    sys.exit(not result.wasSuccessful())


if __name__ == '__main__':
    main()