# Copyright 2021 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import ctypes
import io
import json
import logging
import os
import pandas as pd
import signal
import subprocess
import sys
import datetime
import time
import typing

import browsers
import scenarios
import utils


class DriverContext:
  """Class in charge of running the measurements and keeping track
  of the global state needed to do so.
  """

  def __init__(self, output_dir: str, power_sample_path: str):
    """
    Args:
      output_dir: A string path of Where the results should be stored.
    """

    self._output_dir = output_dir
    self._power_sample_path = power_sample_path

    # Make sure there is somewhere to put  results.
    os.makedirs(f"{self._output_dir}", exist_ok=True)

  def __enter__(self):

    self._caffeinate_process = subprocess.Popen([
        "caffeinate",
        "-d",  # Prevent the display from sleeping.
    ])

    # Launch a power_sampler whose job is only to simulate an active user. This
    # is applicable to both profiling and measuring so should always be done.
    power_sampler_args = [
        self._power_sample_path, "--no-samplers", "--simulate-user-active"
    ]
    self._power_sampler_process = subprocess.Popen(power_sampler_args,
                                                   stdout=subprocess.PIPE,
                                                   stdin=subprocess.PIPE)

    return self

  def __exit__(self, exc_type, exc_val, exc_tb):
    utils.TerminateProcess(self._caffeinate_process)
    utils.TerminateProcess(self._power_sampler_process)

  def SetMainDisplayBrightness(self, brightness_level: int):
    # This function imitates the open-source "brightness" tool at
    # https://github.com/nriley/brightness.
    # Since the benchmark doesn't care about older MacOSen, multiple displays
    # or other complications that tool has to consider, setting the brightness
    # level boils down to calling this function for the main display.
    CoreGraphics = ctypes.CDLL(
        "/System/Library/Frameworks/CoreGraphics.framework/CoreGraphics")
    main_display = CoreGraphics.CGMainDisplayID()
    DisplayServices = ctypes.CDLL(
        "/System/Library/PrivateFrameworks/DisplayServices.framework"
        "/DisplayServices")
    DisplayServices.DisplayServicesSetBrightness.argtypes = [
        ctypes.c_int, ctypes.c_float
    ]
    DisplayServices.DisplayServicesSetBrightness(main_display,
                                                 brightness_level / 100)

  def WaitBatteryNotFull(self):
    # Empirical evidence has shown that right after a full battery charge, the
    # current capacity stays equal to the maximum capacity for several minutes,
    # despite the fact that power is definitely consumed. To ensure that power
    # consumption estimates from battery level are meaningful, wait until the
    # battery is no longer reporting being fully charged before benchmarking.

    power_sampler_args = [
        self._power_sample_path, "--sample-on-notification",
        "--samplers=battery", "--sample-count=1"
    ]

    logging.info('Waiting for battery to no longer be full')

    while True:
      power_sampler_output = subprocess.check_output(power_sampler_args)
      power_sampler_data = pd.read_csv(io.BytesIO(power_sampler_output))
      max_capacity = power_sampler_data.iloc[0]['battery_max_capacity(Ah)']
      current_capacity = power_sampler_data.iloc[0][
          'battery_current_capacity(Ah)']

      logging.info(
          f'Battery level is {100 * current_capacity / max_capacity:.2f}%')

      if max_capacity != current_capacity:
        return

    logging.info('Battery is no longer full')

  def CheckEnv(self, throw_on_bad_env: bool):
    """Verifies that the environment is conducive to proper profiling or
    measurements.

    Args:
      throw_on_bad_env: False if executions continues no matter what and
      only warnings are printed.

    Raises:
      SystemExit: When the environment is invalid and throw_on_bad_env is
      True.
    """

    if throw_on_bad_env:
      logging_function = logging.error
    else:
      logging_function = logging.warning

    logging.warning("Trying sudo access. Possibly enter password:")
    sudo_check = subprocess.Popen(["sudo", "ls"],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.STDOUT)
    sudo_check.wait()

    try:
      check_env = subprocess.run([
          'zsh', '-c',
          os.path.join(os.path.dirname(__file__), 'check_env.sh')
      ],
                                 check=throw_on_bad_env,
                                 capture_output=True)
      logging_function(check_env.stdout.decode('ascii'))
    except subprocess.CalledProcessError as e:
      logging_function(e.stdout.decode('ascii'))
      sys.exit(-1)

    # Make sure that no browsers are running which would affect the
    # tests.
    for browser_name in browsers.PROCESS_NAMES:
      if utils.FindProcess(browser_name):
        logging_function(f"{browser_name} already running. "
                         "Make sure to close it before running again.")

        if throw_on_bad_env:
          sys.exit(-1)

  def Record(self, scenario_driver: scenarios.ScenarioOSADriver):
    """Cover the running of the scenario with powermetrics and save
    the results

    Args:
      scenario_config: A dictionary describing the scenario.
    """

    self.WriteScenarioSummary(scenario_driver)

    powermetrics_output = os.path.join(self._output_dir, scenario_driver.name,
                                       "powermetrics.plist")
    power_sampler_output = os.path.join(self._output_dir, scenario_driver.name,
                                        "power_sampler.json")

    powermetrics_process = None
    power_sampler_process = None
    browser_process = None

    cycle_length_in_secs = scenario_driver.CycleDuration().total_seconds()
    if cycle_length_in_secs < 60:
      raise ValueError("Cycle length must be more than 60 seconds!")
    elif cycle_length_in_secs % 60 != 0:
      logging.warning("Cycle has a duration not divisible by 60 secs."
                      "This is suboptimal for measurment variance.")
    try:
      scenario_driver.Launch()
      if hasattr(scenario_driver, 'browser'):
        browser_process = scenario_driver.browser.browser_process

      # "-i 60000" to emit a sample every minute. This is the same frequency as
      # power_sampler, which emits a sample on IOPMPowerSource notification,
      # which happens every minute.
      powermetrics_args = [
          "sudo", "powermetrics", "-f", "plist", "--samplers",
          "tasks,cpu_power,gpu_power,thermal,disk,network",
          "--show-process-coalition", "--show-process-gpu",
          "--show-process-energy", "-i", "60000", "--output-file",
          powermetrics_output
      ]

      powermetrics_process = subprocess.Popen(powermetrics_args)

      power_sampler_args = [
          self._power_sample_path, f"--sample-on-notification",
          f"--initial-sample",
          "--samplers=battery,smc,user_idle_level,main_display",
          f"--timeout={int(scenario_driver.duration.total_seconds())}",
          f"--json-output-file={power_sampler_output}"
      ]

      if browser_process is not None:
        power_sampler_args += [
            f"--resource-coalition-pid={browser_process.pid}"
        ]
      power_sampler_process = subprocess.Popen(power_sampler_args)

      scenario_driver.Wait()

      logging.debug("Waiting for power_sampler to exit")
      power_sampler_process.wait()
      logging.debug(
          f"power_sampler returned {power_sampler_process.returncode}")

    finally:
      scenario_driver.TearDown()
      if power_sampler_process:
        utils.TerminateProcess(power_sampler_process)
      if powermetrics_process:
        # Force powermetrics to flush data.
        utils.SendSignalToRootProcess(powermetrics_process, signal.SIGIO)
        utils.TerminateRootProcess(powermetrics_process)

  def Profile(self, scenario_driver: scenarios.ScenarioWithBrowserOSADriver,
              profile_mode: str):
    """Cover the running of the scenario with DTrace and save the
    results.

    Args:
      profile_mode: A string describing the Profile mode between "wakeups"
      and "cpu_time".

    Raises:
      TimeoutExpired: When a DTrace process takes more than 30 seconds to
      terminate after the end of the scenario.
    """

    if scenario_driver.browser is None:
      raise ValueError("Scenario must have an associated browser.")
    if scenario_driver.browser.name not in ["chromium", "canary", "chrome"]:
      raise ValueError("Only Chromium can be profiled! Skipping.")

    self.WriteScenarioSummary(scenario_driver)

    dtraces_output_dir = os.path.join(self._output_dir, scenario_driver.name,
                                      f"dtraces_{profile_mode}")
    os.makedirs(dtraces_output_dir, exist_ok=True)
    scenario_driver.Launch()
    browser_process = scenario_driver.browser.browser_process

    # Set up the environment for correct dtrace execution.
    dtrace_env = os.environ.copy()
    dtrace_env["DYLD_SHARED_REGION"] = "avoid"

    try:
      with open(
          os.path.join(self._output_dir, scenario_driver.name,
                       f'dtrace_{profile_mode}_log.txt'), "w") as dtrace_log:

        scripts_dir = os.path.join(os.path.dirname(__file__), "dtrace_scripts")
        if profile_mode == "wakeups":
          script = os.path.join(scripts_dir, "iwakeups.d")
        else:
          script = os.path.join(scripts_dir, "profile.d")

        pid = browser_process.pid
        iteration = 1

        # Capture until the scenario is done running.
        while True:
          output_filename = os.path.join(dtraces_output_dir,
                                         f"{pid}_{iteration}.txt")
          dtrace_args = [
              'sudo', 'dtrace', '-p', f"{pid}", "-o", output_filename, '-s',
              script, f"{pid}"
          ]

          dtrace_process = subprocess.Popen(dtrace_args,
                                            env=dtrace_env,
                                            stdout=dtrace_log,
                                            stderr=dtrace_log)

          # This timeout was chosen experimentally to on an M1 MBA to avoid
          # running for too long and losing samples (because of suspected
          #  bug in Dtrace) and running too little and thus having too much
          # overhead.
          time.sleep(7.199)
          utils.TerminateRootProcess(dtrace_process)
          iteration = iteration + 1

          if scenario_driver.script_process.poll() is not None:
            break

    finally:
      scenario_driver.TearDown()

    logging.debug(f"Waiting for dtrace to exit")

  def Trace(self, scenario_driver: scenarios.ScenarioOSADriver):
    self.WriteScenarioSummary(scenario_driver)

    try:
      scenario_driver.Launch()
      scenario_driver.Wait()
    finally:
      scenario_driver.TearDown()

  def WriteScenarioSummary(
      self, scenario_driver: scenarios.ScenarioWithBrowserOSADriver):
    """Outputs a json file describing `scenario_driver` arguments into the
        output directory
    """
    os.makedirs(os.path.join(self._output_dir, scenario_driver.name),
                exist_ok=True)
    with open(
        os.path.join(self._output_dir, scenario_driver.name, 'metadata.json'),
        'w') as summary_file:
      json.dump(scenario_driver.Summary(), summary_file, indent=2)