# 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.

from __future__ import annotations

import dataclasses
import datetime as dt
import enum
import logging
import os
import pathlib
import urllib.request
from typing import (TYPE_CHECKING, Any, Callable, Dict, Iterable, List,
                    Optional, Union)
from urllib.parse import urlparse

import colorama

from crossbench import compat, helper, plt

if TYPE_CHECKING:
  from crossbench.browsers.browser import Browser
  from crossbench.probes.probe import Probe
  from crossbench.runner.runner import Runner


def merge_bool(name: str, left: Optional[bool],
               right: Optional[bool]) -> Optional[bool]:
  if left is None:
    return right
  if right is None:
    return left
  if left != right:
    raise ValueError(f"Conflicting merge values for {name}: "
                     f"{left} vs. {right}")
  return left


NumberT = Union[float, int]


def merge_number_max(name: str, left: Optional[NumberT],
                     right: Optional[NumberT]) -> Optional[NumberT]:
  del name
  if left is None:
    return right
  if right is None:
    return left
  return max(left, right)


def merge_number_min(name: str, left: Optional[NumberT],
                     right: Optional[NumberT]) -> Optional[NumberT]:
  del name
  if left is None:
    return right
  if right is None:
    return left
  return min(left, right)


def merge_str_list(name: str, left: Optional[List[str]],
                   right: Optional[List[str]]) -> Optional[List[str]]:
  del name
  if left is None:
    return right
  if right is None:
    return left
  return left + right


@dataclasses.dataclass(frozen=True)
class HostEnvironmentConfig:
  IGNORE = None

  disk_min_free_space_gib: Optional[float] = IGNORE
  power_use_battery: Optional[bool] = IGNORE
  screen_brightness_percent: Optional[int] = IGNORE
  cpu_max_usage_percent: Optional[float] = IGNORE
  cpu_min_relative_speed: Optional[float] = IGNORE
  system_allow_monitoring: Optional[bool] = IGNORE
  browser_allow_existing_process: Optional[bool] = IGNORE
  browser_allow_background: Optional[bool] = IGNORE
  browser_is_headless: Optional[bool] = IGNORE
  require_probes: Optional[bool] = IGNORE
  system_forbidden_process_names: Optional[List[str]] = IGNORE
  screen_allow_autobrightness: Optional[bool] = IGNORE

  def merge(self, other: HostEnvironmentConfig) -> HostEnvironmentConfig:
    mergers: Dict[str, Callable[[str, Any, Any], Any]] = {
        "disk_min_free_space_gib": merge_number_max,
        "power_use_battery": merge_bool,
        "screen_brightness_percent": merge_number_max,
        "cpu_max_usage_percent": merge_number_min,
        "cpu_min_relative_speed": merge_number_max,
        "system_allow_monitoring": merge_bool,
        "browser_allow_existing_process": merge_bool,
        "browser_allow_background": merge_bool,
        "browser_is_headless": merge_bool,
        "require_probes": merge_bool,
        "system_forbidden_process_names": merge_str_list,
        "screen_allow_autobrightness": merge_bool,
    }
    kwargs = {}
    for name, merger in mergers.items():
      self_value = getattr(self, name)
      other_value = getattr(other, name)
      kwargs[name] = merger(name, self_value, other_value)
    return HostEnvironmentConfig(**kwargs)


@enum.unique
class ValidationMode(compat.StrEnumWithHelp):
  THROW = ("throw", "Strict mode, throw and abort on env issues")
  PROMPT = ("prompt", "Prompt to accept potential env issues")
  WARN = ("warn", "Only display a warning for env issue")
  SKIP = ("skip", "Don't perform any env validation")


class ValidationError(Exception):
  pass


_config_default = HostEnvironmentConfig()
_config_strict = HostEnvironmentConfig(
    cpu_max_usage_percent=98,
    cpu_min_relative_speed=1,
    system_allow_monitoring=False,
    browser_allow_existing_process=False,
    require_probes=True,
)
_config_battery = _config_strict.merge(
    HostEnvironmentConfig(power_use_battery=True))
_config_power = _config_strict.merge(
    HostEnvironmentConfig(power_use_battery=False))
_config_catan = _config_strict.merge(
    HostEnvironmentConfig(
        screen_brightness_percent=65,
        system_forbidden_process_names=["terminal", "iterm2"],
        screen_allow_autobrightness=False))

STALE_RESULT_ICONS = {
    75: "👻",
    100: "👾",
    125: "🎃",
    150: "👹",
    200: "💀",
    250: "😱",
    500: "🤯"
}


class HostEnvironment:
  """
  HostEnvironment can check and enforce certain settings on a host
  where we run benchmarks.

  Modes:
    skip:     Do not perform any checks
    warn:     Only warn about mismatching host conditions
    enforce:  Tries to auto-enforce conditions and warns about others.
    prompt:   Interactive mode to skip over certain conditions
    fail:     Fast-fail on mismatch
  """

  CONFIGS = {
      "default": _config_default,
      "strict": _config_strict,
      "battery": _config_battery,
      "power": _config_power,
      "catan": _config_catan,
  }

  def __init__(self,
               runner: Runner,
               config: Optional[HostEnvironmentConfig] = None,
               validation_mode: ValidationMode = ValidationMode.THROW):
    self._wait_until = dt.datetime.now()
    self._config = config or HostEnvironmentConfig()
    self._runner = runner
    self._platform = runner.platform
    self._validation_mode = validation_mode

  @property
  def runner(self) -> Runner:
    return self._runner

  @property
  def config(self) -> HostEnvironmentConfig:
    return self._config

  @property
  def validation_mode(self) -> ValidationMode:
    return self._validation_mode

  def _add_min_delay(self, seconds: float) -> None:
    end_time = dt.datetime.now() + dt.timedelta(seconds=seconds)
    if end_time > self._wait_until:
      self._wait_until = end_time

  def _wait_min_time(self) -> None:
    delta = self._wait_until - dt.datetime.now()
    if delta > dt.timedelta(0):
      self._platform.sleep(delta)

  def handle_validation_warning(self, message: str) -> None:
    message = f"Runner/Host environment requests cannot be fulfilled: {message}"
    self.handle_warning(message)

  def handle_warning(self,
                     message: str,
                     allow_interactive: bool = True) -> None:
    """Process a warning, depending on the requested mode, this will
    - throw an error,
    - log a warning,
    - prompts for continue [Yn], or
    - skips (and just debug logs) a warning.
    If returned True (in the prompt mode) the env validation may continue.
    """
    if self._validation_mode == ValidationMode.SKIP:
      logging.debug("Ignoring %s", message)
      return
    if self._validation_mode == ValidationMode.WARN:
      logging.warning(message)
      return
    if self._validation_mode == ValidationMode.PROMPT:
      if allow_interactive:
        result = input(f"{colorama.Fore.RED}{message} Continue?"
                       f"{colorama.Fore.RESET} [Yn]")
        # Accept <enter> as default input to continue.
        if result.lower() != "n":
          return
    elif self._validation_mode != ValidationMode.THROW:
      raise ValueError(
          f"Unknown environment validation mode={self._validation_mode}")
    raise ValidationError(message)

  def validate_url(self,
                   url: str,
                   platform: plt.Platform = plt.PLATFORM) -> bool:
    if self._validation_mode == ValidationMode.SKIP:
      return True
    result = urlparse(url)
    if result.scheme == "file":
      return platform.exists(pathlib.Path(result.path))
    if platform.is_remote and result.hostname in ("localhost", "127.0.0.1"):
      # TODO: support remote URL verification, for now we just assume that
      # checking a live site is ok.
      return True
    try:
      if not all([result.scheme in ["http", "https"], result.netloc]):
        return False
      if self._validation_mode != ValidationMode.PROMPT:
        return True
      with urllib.request.urlopen(url, timeout=5) as request:
        if request.getcode() == 200:
          return True
        logging.debug("Could not load URL '%s', got %s", url, request)
    except urllib.error.URLError as e:
      logging.debug("Could not parse URL '%s' got error: %s", url, e)
    return False

  def _check_system_monitoring(self) -> None:
    # TODO(cbruni): refactor to use list_... and disable_system_monitoring api
    if self._platform.is_macos:
      self._check_crowdstrike()

  def _check_crowdstrike(self) -> None:
    """Crowdstrike security monitoring (for googlers go/crowdstrike-falcon) can
    have quite terrible overhead for each file-access. Disable it to reduce
    flakiness. """
    is_disabled = False
    force_disable = self._config.system_allow_monitoring is False
    try:
      # TODO(cbruni): refactor to use list_... and disable_system_monitoring api
      is_disabled = self._platform.check_system_monitoring(force_disable)
      if force_disable:
        # Add cool-down period, crowdstrike caused CPU usage spikes
        self._add_min_delay(5)
    except plt.SubprocessError as e:
      self.handle_validation_warning(
          "Could not disable go/crowdstrike-falcon monitor which can cause"
          f" high background CPU usage: {e}")
      return
    if not is_disabled:
      self.handle_validation_warning(
          "Crowdstrike monitoring is running, "
          "which can impact startup performance drastically.\n"
          "Use the following command to disable it manually:\n"
          "sudo /Applications/Falcon.app/Contents/Resources/falconctl unload\n")

  def _check_disk_space(self) -> None:
    limit = self._config.disk_min_free_space_gib
    if limit is HostEnvironmentConfig.IGNORE:
      return
    # Check the remaining disk space on the FS where we write the results.
    usage = self._platform.disk_usage(self._runner.out_dir)
    free_gib = round(usage.free / 1024 / 1024 / 1024, 2)
    if free_gib < limit:
      self.handle_validation_warning(
          f"Only {free_gib}GiB disk space left, expected at least {limit}GiB.")

  def _check_power(self) -> None:
    use_battery = self._config.power_use_battery
    if use_battery is HostEnvironmentConfig.IGNORE:
      return
    battery_probes: List[Probe] = []
    # Certain probes may require battery power:
    for probe in self._runner.probes:
      if probe.BATTERY_ONLY:
        battery_probes.append(probe)
    if not use_battery and battery_probes:
      probes_str = ",".join(probe.name for probe in battery_probes)
      self.handle_validation_warning(
          "Requested battery_power=False, "
          f"but probes={probes_str} require battery power.")
    sys_use_battery = self._platform.is_battery_powered
    if sys_use_battery != use_battery:
      self.handle_validation_warning(
          f"Expected battery_power={use_battery}, "
          f"but the system reported battery_power={sys_use_battery}")

  def _check_cpu_usage(self) -> None:
    max_cpu_usage = self._config.cpu_max_usage_percent
    if max_cpu_usage is HostEnvironmentConfig.IGNORE:
      return
    cpu_usage_percent = round(100 * self._platform.cpu_usage(), 1)
    if cpu_usage_percent > max_cpu_usage:
      self.handle_validation_warning(
          f"CPU usage={cpu_usage_percent}% is higher than "
          f"requested max={max_cpu_usage}%.")

  def _check_cpu_temperature(self) -> None:
    min_relative_speed = self._config.cpu_min_relative_speed
    if min_relative_speed is HostEnvironmentConfig.IGNORE:
      return
    cpu_speed = self._platform.get_relative_cpu_speed()
    if cpu_speed < min_relative_speed:
      self.handle_validation_warning(
          "CPU thermal throttling is active. "
          f"Relative speed is {cpu_speed}, "
          f"but expected at least {min_relative_speed}.")

  def _check_forbidden_system_process(self) -> None:
    # Verify that no terminals are running.
    # They introduce too much overhead. (As measured with powermetrics)
    system_forbidden_process_names = self._config.system_forbidden_process_names
    if system_forbidden_process_names is HostEnvironmentConfig.IGNORE:
      return
    process_found = self._platform.process_running(
        system_forbidden_process_names)
    if process_found:
      self.handle_validation_warning(
          f"Process:{process_found} found."
          "Make sure not to have a terminal opened. Use SSH.")

  def _check_screen_autobrightness(self) -> None:
    auto_brightness = self._config.screen_allow_autobrightness
    if auto_brightness is not False:
      return
    if self._platform.check_autobrightness():
      self.handle_validation_warning(
          "Auto-brightness was found to be ON. "
          "Deactivate it in 'System Preferences/Displays'")

  def _check_cpu_power_mode(self) -> bool:
    # TODO Implement checks for performance mode
    return True

  def _check_running_binaries(self) -> None:
    if self._config.browser_allow_existing_process:
      return
    browser_binaries: Dict[str, List[Browser]] = helper.group_by(
        self._runner.browsers, key=lambda browser: str(browser.path.resolve()))
    own_pid = os.getpid()
    for proc_info in self._platform.processes(["cmdline", "exe", "pid",
                                               "name"]):
      if not browser_binaries:
        return
      # Skip over this python script which might have the binary path as
      # part of the command line invocation.
      if proc_info["pid"] == own_pid:
        continue
      cmdline = " ".join(proc_info["cmdline"] or "")
      exe = proc_info["exe"]
      for binary, browsers in list(browser_binaries.items()):
        # Add a white-space to get less false-positives
        if f"{binary} " not in cmdline and binary != exe:
          continue
        # Use the first in the group
        browser: Browser = browsers[0]
        logging.debug("Binary=%s", binary)
        logging.debug("PS status output:")
        logging.debug("proc(pid=%s, name=%s, cmd=%s)", proc_info["pid"],
                      proc_info["name"], cmdline)
        self.handle_validation_warning(
            f"{browser.app_name} {browser.version} seems to be already running."
        )
        # Avoid re-checking the same binary once we've allowed it to be running.
        del browser_binaries[binary]

  def _check_screen_brightness(self) -> None:
    brightness = self._config.screen_brightness_percent
    if brightness is HostEnvironmentConfig.IGNORE:
      return
    assert 0 <= brightness <= 100, f"Invalid brightness={brightness}"
    self._platform.set_main_display_brightness(brightness)
    current = self._platform.get_main_display_brightness()
    if current != brightness:
      self.handle_validation_warning(
          f"Requested main display brightness={brightness}%, "
          "but got {brightness}%")

  def _check_headless(self) -> None:
    # TODO: migrate to full viewport support
    requested_headless = self._config.browser_is_headless
    if requested_headless is HostEnvironmentConfig.IGNORE:
      return
    if self._platform.is_linux and not requested_headless:
      # Check that the system can run browsers with a UI.
      if not self._platform.has_display:
        self.handle_validation_warning(
            "Requested browser_is_headless=False, "
            "but no DISPLAY is available to run with a UI.")
    # Check that browsers are running in the requested display mode:
    for browser in self._runner.browsers:
      if browser.viewport.is_headless != requested_headless:
        self.handle_validation_warning(
            f"Requested browser_is_headless={requested_headless},"
            f"but browser {browser.unique_name} has conflicting "
            f"headless={browser.viewport.is_headless}.")

  def _check_probes(self) -> None:
    for probe in self._runner.probes:
      try:
        probe.validate_env(self)
      except Exception as e:
        raise ValidationError(
            f"Probe='{probe.NAME}' validation failed: {e}") from e
    require_probes = self._config.require_probes
    if require_probes is HostEnvironmentConfig.IGNORE:
      return
    if self._config.require_probes and not self._runner.probes:
      self.handle_validation_warning("No probes specified.")

  def _check_results_dir(self) -> None:
    results_dir = self._runner.out_dir.parent
    if not results_dir.exists():
      return
    results = [path for path in results_dir.iterdir() if path.is_dir()]
    num_results = len(results)
    if num_results < 20:
      return
    message = (f"Found {num_results} existing crossbench results. "
               f"Consider cleaning stale results in '{results_dir}'")
    for count, icon in reversed(STALE_RESULT_ICONS.items()):
      if num_results > count:
        message = f"{icon} {message}"
        break
    if num_results > 50:
      logging.error(message)
    else:
      logging.warning(message)

  def _check_macos_terminal(self) -> None:
    if not self._platform.is_macos or (
        self._platform.environ.get("TERM_PROGRAM") != "Apple_Terminal"):
      return
    any_not_headless = any(
        not browser.viewport.is_headless for browser in self._runner.browsers)
    if any_not_headless:
      self.handle_validation_warning(
          "Terminal.app does not launch apps in the foreground.\n"
          "Please use iTerm.app for a better experience.")

  def check_browser_focused(self, browser: Browser) -> None:
    if (self._config.browser_allow_background or not browser.pid or
        browser.viewport.is_headless):
      return
    info = browser.platform.foreground_process()
    if not info:
      return
    if info["pid"] != browser.pid:
      self.handle_warning(
          f"Browser(name={browser.unique_name} pid={browser.pid})) "
          "was not in the foreground at the end of the benchmark. "
          "Background apps and tabs can be heavily throttled.",
          allow_interactive=False)

  def setup(self) -> None:
    self.validate()

  def validate(self) -> None:
    logging.info("-" * 80)
    if self._validation_mode == ValidationMode.SKIP:
      logging.info("VALIDATE ENVIRONMENT: SKIP")
      return
    message = "VALIDATE ENVIRONMENT"
    if self._validation_mode != ValidationMode.WARN:
      message += " (--env-validation=warn for soft warnings)"
    message += ": %s"
    logging.info(message, self._validation_mode.name)
    self._check_system_monitoring()
    self._check_power()
    self._check_disk_space()
    self._check_cpu_usage()
    self._check_cpu_temperature()
    self._check_cpu_power_mode()
    self._check_running_binaries()
    self._check_screen_brightness()
    self._check_headless()
    self._check_results_dir()
    self._check_probes()
    self._wait_min_time()
    self._check_forbidden_system_process()
    self._check_screen_autobrightness()
    self._check_macos_terminal()

  def check_installed(self,
                      binaries: Iterable[str],
                      message: str = "Missing binaries: {}") -> None:
    assert not isinstance(binaries, str), "Expected iterable of strings."
    missing_binaries = list(
        binary for binary in binaries if not self._platform.which(binary))
    if missing_binaries:
      self.handle_validation_warning(message.format(missing_binaries))

  def check_sh_success(self,
                       *args: Union[str, pathlib.Path],
                       message: str = "Could not execute: {}") -> None:
    assert args, "Missing sh arguments"
    try:
      assert self._platform.sh_stdout(*args, quiet=True)
    except plt.SubprocessError as e:
      self.handle_validation_warning(message.format(e))
