# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
import json
import os
import pathlib
import platform
import shutil
import signal
import subprocess
import sys
import time

import mozprocess

from mozperftest.layers import Layer
from mozperftest.utils import ON_TRY, archive_folder, install_package, temp_dir

"""
Python dependencies needed for mozperftest have to be installed when running
via shellscript there is an issue with the way the shellscript runner does
not have all of the environment variables and system settings
"""
INTERNAL_PYPI = "https://pypi.pub.build.mozilla.org/pub/"
NUMPY_DEPENDENCY = "numpy<2"
OPENCV_DEPENDENCY = "opencv-python==4.10.0.84"


class UnknownScriptError(Exception):
    """Triggered when an unknown script type is encountered."""

    pass


class ShellScriptData:
    SUBMETRICS_SUFFIX = "_submetrics"

    def open_data(self, data):
        return {
            "name": "shellscript",
            "subtest": data["name"],
            "data": [
                {"file": "custom", "value": value, "xaxis": xaxis}
                for xaxis, value in enumerate(data["values"])
            ],
            "shouldAlert": data.get("shouldAlert", True),
            "unit": data.get("unit", "ms"),
            "lowerIsBetter": data.get("lowerIsBetter", True),
        }

    def transform(self, data):
        return data

    merge = transform

    def summary(self, suite):
        """Use the primary submetric's value as the grouped suite summary.

        Submetric suites are named `<primary>_submetrics` and gather a set of
        related measurements as subtests, one of which (named `<primary>`) is
        the representative value for the group. Reporting that subtest's value
        as the suite summary keeps the grouped suite anchored to a single,
        meaningful number. Returns None for regular suites so the default
        summary (mean of the subtests) applies.

        Only available in the Perfherder layer.
        """
        if not suite["name"].endswith(self.SUBMETRICS_SUFFIX):
            return None
        primary_name = suite["name"][: -len(self.SUBMETRICS_SUFFIX)]
        for subtest in suite["subtests"]:
            if subtest["name"] == primary_name:
                return subtest["value"]
        return None


class ShellScriptRunner(Layer):
    """
    This is a layer that can be used to run custom shell scripts. They
    are expected to produce a log message prefixed with `perfMetrics` and
    contain the data from the test.
    """

    name = "shell-script"
    activated = True
    arguments = {
        "output-timeout": {
            "default": 120,
            "help": "Output timeout for the script, or how long to wait for a log message.",
        },
        "process-timeout": {
            "default": 600,
            "help": "Process timeout for the script, or the max run time.",
        },
    }

    def __init__(self, env, mach_cmd):
        super().__init__(env, mach_cmd)
        self.metrics = []
        self.timed_out = False
        self.output_timed_out = False

    def setup(self):
        # Install numpy first so opencv-python's numpy>=1.17.0 dep is already
        # satisfied, then install opencv-python, see similar fix in (bug 2033807)
        install_package(self.mach_cmd.virtualenv_manager, NUMPY_DEPENDENCY)
        subprocess.check_call([
            self.mach_cmd.virtualenv_manager.python_path,
            "-m",
            "pip",
            "install",
            OPENCV_DEPENDENCY,
            "--no-deps",
            "--no-index",
            "--find-links",
            INTERNAL_PYPI,
        ])

    def kill(self, proc):
        if "win" in platform.system().lower():
            proc.send_signal(signal.CTRL_BREAK_EVENT)
        else:
            os.killpg(proc.pid, signal.SIGKILL)
        proc.wait()

    def parse_metrics(self):
        parsed_metrics = []
        for metrics in self.metrics:
            prepared_metrics = (
                metrics
                .replace("perfMetrics:", "")
                .replace("{{", "{")
                .replace("}}", "}")
                .strip()
            )

            metrics_json = json.loads(prepared_metrics)
            if isinstance(metrics_json, list):
                parsed_metrics.extend(metrics_json)
            else:
                parsed_metrics.append(metrics_json)

        return parsed_metrics

    def line_handler_wrapper(self):
        """This function is used to gather the perfMetrics logs."""

        def _line_handler(proc, line):
            # NOTE: this hack is to workaround encoding issues on windows
            # a newer version of browsertime adds a `σ` character to output
            line = line.replace(b"\xcf\x83", b"")

            line = line.decode("utf-8")
            if "perfMetrics" in line:
                self.metrics.append(line)

            # Bug 1900056 - Use a different logger in mozperftest because the current
            # one can't handle messages with curly braces or JSONs in them
            print(line.strip())

        return _line_handler

    def run(self, metadata):
        test = metadata.script

        cmd = []
        if test["filename"].endswith(".sh"):
            cmd = ["bash", test["filename"]]
        else:
            raise UnknownScriptError(
                "Only `.sh` (bash) scripts are currently implemented."
            )

        def timeout_handler(proc):
            self.timed_out = True
            self.error("Process timed out")
            self.kill(proc)

        def output_timeout_handler(proc):
            self.output_timed_out = True
            self.error("Process output timed out")
            self.kill(proc)

        with temp_dir() as testing_dir:
            os.environ["APP"] = self.get_arg("app")
            os.environ["TESTING_DIR"] = testing_dir
            os.environ["BROWSER_BINARY"] = metadata.binary
            os.environ["PYTHON_PATH_SHELL_SCRIPT"] = sys.executable

            # Setup a python packages path for python scripts. This is required since
            # the sys.path modifications made by mozperftest don't get propagated
            # to the scripts. Furthermore, PYTHONPATH can't be used here since it
            # gets overrided by something like a ._pth file somewhere along the way.
            venv_site_lib = str(
                pathlib.Path(
                    self.mach_cmd.virtualenv_manager.bin_path, "..", "lib"
                ).resolve()
            )
            venv_site_packages = pathlib.Path(
                venv_site_lib,
                f"python{sys.version_info.major}.{sys.version_info.minor}",
                "site-packages",
            )
            if not venv_site_packages.exists():
                venv_site_packages = pathlib.Path(
                    venv_site_lib,
                    "site-packages",
                )
            os.environ["PYTHON_PACKAGES"] = str(venv_site_packages)

            mozprocess.run_and_wait(
                cmd,
                output_line_handler=self.line_handler_wrapper(),
                env=os.environ,
                timeout=self.get_arg("process-timeout"),
                timeout_handler=timeout_handler,
                output_timeout=self.get_arg("output-timeout"),
                output_timeout_handler=output_timeout_handler,
                text=False,
            )

            if self.get_arg("output"):
                if ON_TRY:
                    # In CI, make an archive of the files and upload those
                    archive_folder(
                        pathlib.Path(testing_dir),
                        pathlib.Path(self.get_arg("output")),
                        archive_name=test["name"],
                    )
                else:
                    output_dir = pathlib.Path(
                        self.get_arg("output"),
                        test["name"],
                        f"run-{str(time.time()).split('.')[0]}",
                    )
                    self.info(f"Copying testing directory to {output_dir}")
                    shutil.copytree(testing_dir, output_dir)
                    self.env.set_arg("output", output_dir)

        # Route each metric to a suite based on its optional `suite` tag,
        # defaulting to the test name. Grouping the submetrics and computing a
        # suite summary value is left to the metrics layer (see ShellScriptData),
        # keeping this layer focused on running the test and reporting raw data.
        default_suite = test["name"]
        suites = {}
        for m in self.parse_metrics():
            suite_name = m.pop("suite", default_suite)
            suites.setdefault(suite_name, []).append(m)

        for suite_name, suite_metrics in suites.items():
            metadata.add_result({
                "name": suite_name,
                "framework": {"name": "mozperftest"},
                "transformer": "mozperftest.test.shellscript:ShellScriptData",
                "shouldAlert": True,
                "results": suite_metrics,
            })

        return metadata
