#!/usr/bin/env python3

# 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 argparse
import logging
import typing
import os
import datetime

from driver import DriverContext
import scenarios
import browsers
import plug
import utils

def IterScenarios(
    scenario_names: typing.List[str],
    browser_driver_factory: typing.Callable[[], browsers.BrowserDriver],
    **kwargs):
  for scenario_and_browser_name in scenario_names:
    scenario_name, _, browser_name = scenario_and_browser_name.partition(':')
    browser_name, _, variation = browser_name.partition(':')
    browser_driver = browser_driver_factory(browser_name, variation)
    scenario_driver = scenarios.MakeScenarioDriver(scenario_name,
                                                   browser_driver, **kwargs)
    if scenario_driver is None:
      logging.error(f"Skipping invalid scenario {scenario_and_browser_name}.")
    else:
      yield scenario_driver


def main():
  parser = argparse.ArgumentParser(description='Runs browser power benchmarks')
  parser.add_argument("--output_dir", help="Output dir")
  parser.add_argument('--no-checks',
                      dest='no_checks',
                      action='store_true',
                      help="Invalid environment doesn't throw")
  parser.add_argument(
      '--skip-wait-for-battery-not-full',
      dest='wait_for_battery_not_full',
      action='store_false',
      help=("Skip waiting until the battery isn't full before recording a "
            "scenario (for debugging only)"))
  mode_group = parser.add_mutually_exclusive_group()
  mode_group.add_argument(
      '--tracing_mode',
      dest='tracing_mode',
      action='store_true',
      help="Grab a trace instead of a profile or benchmark.")

  # Profile related arguments
  mode_group.add_argument(
      '--profile_mode',
      dest='profile_mode',
      action='store',
      choices=["wakeups", "cpu_time"],
      help="Profile the application in one of two modes: wakeups, cpu_time.")
  parser.add_argument("--power_sampler",
                      help="Path to power sampler binary",
                      required=True)
  parser.add_argument(
      '--scenarios',
      dest='scenarios',
      action='store',
      required=True,
      nargs='+',
      help="List of scenarios and browsers to run in the format"
      "<scenario_name>:<browser_name>, e.g. idle_on_wiki:safari")
  parser.add_argument('--meet-meeting-id',
                      dest='meet_meeting_id',
                      action='store',
                      help='The meeting ID to use for the Meet benchmarks')
  parser.add_argument(
      '--chrome-user-dir',
      dest='chrome_user_dir',
      action='store',
      help='The user data dir to pass to Chrome via --user-data-dir')

  parser.add_argument('--verbose',
                      action='store_true',
                      default=False,
                      help='Print verbose output.')

  parser.add_argument(
      "--brightness_level",
      type=int,
      required=False,
      # This is the average brightness from UMA data.
      default=65,
      help="Desired brightness level.")

  # If an ip is provided for the Kasa switch it needs to be fully set up
  # (see plug.py). It will be used to keep the machine charged between
  # scenarios.
  parser.add_argument(
      "--kasa_switch_ip",
      required=False,
      help="IP address of the kasa power switch controlling the current device."
  )

  parser.add_argument('--extra-command-line',
                      dest='extra_command_line',
                      action='append',
                      help="Multiple values are suported.")

  parser.add_argument('--tag',
                      dest='tag',
                      default="",
                      action='store',
                      help='Tag to be added to metada to identify run.')

  args = parser.parse_args()

  if args.verbose:
    log_level = logging.DEBUG
  else:
    log_level = logging.INFO
  logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s',
                      level=log_level)

  output_dir = args.output_dir
  if not output_dir:
    output_dir = os.path.join("output",
                              datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))

  kasa_plug_controller = None
  if args.kasa_switch_ip:
    kasa_plug_controller = plug.get_plug_controller(args.kasa_switch_ip)
    # Turn off power to pass environment checks below.
    kasa_plug_controller.turn_off()

  logging.info(f'Outputing results in {os.path.abspath(output_dir)}')
  with DriverContext(output_dir, args.power_sampler) as driver:
    driver.CheckEnv(not args.no_checks)
    driver.SetMainDisplayBrightness(args.brightness_level)

    # Measure or Profile all defined scenarios.
    def BrowserFactory(browser_name, variation):
      return browsers.MakeBrowserDriver(
          browser_name,
          variation,
          chrome_user_dir=args.chrome_user_dir,
          output_dir=output_dir,
          tracing_mode=args.tracing_mode,
          extra_command_line=args.extra_command_line)

    for scenario in IterScenarios(args.scenarios,
                                  BrowserFactory,
                                  meet_meeting_id=args.meet_meeting_id):

      scenario.tag = args.tag

      if kasa_plug_controller:
        kasa_plug_controller.charge_or_discharge_to(80)

      if args.tracing_mode:
        logging.info(f'Tracing scenario {scenario.name} ...')
        driver.Trace(scenario)
      elif args.profile_mode:
        logging.info(f'Profiling scenario {scenario.name} ...')
        driver.Profile(scenario, profile_mode=args.profile_mode)
      else:
        # This returns immediately after an IOPMPowerSource notification, which
        # is required for power measurements that cover precisely the benchmark
        # interval (if the benchmark starts n seconds after an IOPMPowerSource
        # notification, power measurements will implicitly include these n
        # seconds during which the benchmark wasn't running).
        if args.wait_for_battery_not_full:
          driver.WaitBatteryNotFull()

        logging.info(f'Recording scenario {scenario.name} ...')
        driver.Record(scenario)


if __name__ == "__main__":
  main()