# Copyright 2023 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 argparse
import datetime as dt
import enum
import logging
from typing import (TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple,
                    Type, Union, cast)

from crossbench import compat, helper
from crossbench.benchmarks.speedometer.speedometer import (
    ProbeClsTupleT, SpeedometerBenchmark, SpeedometerBenchmarkStoryFilter,
    SpeedometerProbe, SpeedometerStory)
from crossbench.browsers import viewport as vp
from crossbench.parse import DurationParser, NumberParser
from crossbench.stories.story import Story

if TYPE_CHECKING:
  from crossbench.cli.parser import CrossBenchArgumentParser
  from crossbench.runner.run import Run
  ShuffleSeedT = Optional[Union[str, int]]
  from crossbench.runner.actions import Actions
  from crossbench.types import Json


class Speedometer30Probe(SpeedometerProbe):
  """
  Speedometer3-specific probe (compatible with v3.0).
  Extracts all speedometer times and scores.
  """
  NAME: str = "speedometer_3.0"
  JS: str = "return window.benchmarkClient.metrics"

  @property
  def speedometer(self) -> Speedometer30Benchmark:
    return cast(Speedometer30Benchmark, self.benchmark)

  def to_json(self, actions: Actions) -> Json:
    return actions.js(self.JS)

  def process_json_data(self, json_data) -> Any:
    # Move aggregate scores to the end
    aggregate_keys = []
    for metric_key in json_data.keys():
      if metric_key.startswith("Iteration-"):
        aggregate_keys.append(metric_key)
    aggregate_keys.extend(["Geomean", "Score"])
    for metric_key in aggregate_keys:
      json_data[metric_key] = json_data.pop(metric_key)
    return json_data

  def flatten_json_data(self, json_data: Any) -> Json:
    result: Dict[str, float] = {}
    assert isinstance(json_data, dict), f"Expected dict, got {type(json_data)}"
    for name, metric in json_data.items():
      result[name] = metric["mean"]
    return result

  def _is_valid_metric_key(self, metric_key: str) -> bool:
    parts = metric_key.split("/")
    if len(parts) != 1:
      return False
    if self.speedometer.detailed_metrics:
      return True
    if metric_key.startswith("Iteration-"):
      return False
    if metric_key == "Geomean":
      return False
    return True


@enum.unique
class MeasurementMethod(compat.StrEnumWithHelp):
  RAF = ("raf", "requestAnimationFrame-based measurement")
  TIMER = ("timer", "setTimeout-based measurement")


def to_ms(duration: dt.timedelta) -> int:
  return int(round(duration.total_seconds() * 1000))


def parse_shuffle_seed(value: Optional[Any]) -> ShuffleSeedT:
  if value in (None, "off", "generate"):
    return value
  if isinstance(value, int):
    return value
  return NumberParser.any_int(value, "shuffle-seed")


# Generated by running this JS snippet and updating the bools:
#  JSON.stringify(
#    Suites.reduce((data, e) => {
#        data[e.name]={ tags:e.tags, enabled:!e.disabled};
#        return data}, {}))
#  .replaceAll(":true", ":True")
#  .replaceAll(":false", ":False");
SPEEDOMETER_3_STORY_DATA = {
    "TodoMVC-JavaScript-ES5": {
        "tags": ["all", "default", "todomvc"],
        "enabled": True
    },
    "TodoMVC-JavaScript-ES5-Complex-DOM": {
        "tags": ["all", "todomvc", "complex"],
        "enabled": False
    },
    "TodoMVC-JavaScript-ES6-Webpack": {
        "tags": ["all", "todomvc"],
        "enabled": False
    },
    "TodoMVC-JavaScript-ES6-Webpack-Complex-DOM": {
        "tags": ["all", "default", "todomvc", "complex", "complex-default"],
        "enabled": True
    },
    "TodoMVC-WebComponents": {
        "tags": ["all", "default", "todomvc", "webcomponents"],
        "enabled": True
    },
    "TodoMVC-WebComponents-Complex-DOM": {
        "tags": ["all", "todomvc", "webcomponents", "complex"],
        "enabled": False
    },
    "TodoMVC-React": {
        "tags": ["all", "todomvc"],
        "enabled": False
    },
    "TodoMVC-React-Complex-DOM": {
        "tags": ["all", "default", "todomvc", "complex", "complex-default"],
        "enabled": True
    },
    "TodoMVC-React-Redux": {
        "tags": ["all", "default", "todomvc"],
        "enabled": True
    },
    "TodoMVC-React-Redux-Complex-DOM": {
        "tags": ["all", "todomvc", "complex"],
        "enabled": False
    },
    "TodoMVC-Backbone": {
        "tags": ["all", "default", "todomvc"],
        "enabled": True
    },
    "TodoMVC-Backbone-Complex-DOM": {
        "tags": ["all", "todomvc", "complex"],
        "enabled": False
    },
    "TodoMVC-Angular": {
        "tags": ["all", "todomvc"],
        "enabled": False
    },
    "TodoMVC-Angular-Complex-DOM": {
        "tags": ["all", "default", "todomvc", "complex", "complex-default"],
        "enabled": True
    },
    "TodoMVC-Vue": {
        "tags": ["all", "default", "todomvc"],
        "enabled": True
    },
    "TodoMVC-Vue-Complex-DOM": {
        "tags": ["all", "todomvc", "complex", "complex-default"],
        "enabled": False
    },
    "TodoMVC-jQuery": {
        "tags": ["all", "default", "todomvc"],
        "enabled": True
    },
    "TodoMVC-jQuery-Complex-DOM": {
        "tags": ["all", "todomvc", "complex"],
        "enabled": False
    },
    "TodoMVC-Preact": {
        "tags": ["all", "todomvc"],
        "enabled": False
    },
    "TodoMVC-Preact-Complex-DOM": {
        "tags": ["all", "default", "todomvc", "complex", "complex-default"],
        "enabled": True
    },
    "TodoMVC-Svelte": {
        "tags": ["all", "todomvc"],
        "enabled": False
    },
    "TodoMVC-Svelte-Complex-DOM": {
        "tags": ["all", "default", "todomvc", "complex", "complex-default"],
        "enabled": True
    },
    "TodoMVC-Lit": {
        "tags": ["all", "todomvc", "webcomponents"],
        "enabled": False
    },
    "TodoMVC-Lit-Complex-DOM": {
        "tags": [
            "all", "default", "todomvc", "webcomponents", "complex",
            "complex-default"
        ],
        "enabled": True
    },
    "NewsSite-Next": {
        "tags": ["all", "default", "newssite", "language"],
        "enabled": True
    },
    "NewsSite-Nuxt": {
        "tags": ["all", "default", "newssite"],
        "enabled": True
    },
    "Editor-CodeMirror": {
        "tags": ["all", "default", "editor"],
        "enabled": True
    },
    "Editor-TipTap": {
        "tags": ["all", "default", "editor"],
        "enabled": True
    },
    "Charts-observable-plot": {
        "tags": ["all", "default", "chart"],
        "enabled": True
    },
    "Charts-chartjs": {
        "tags": ["all", "default", "chart"],
        "enabled": True
    },
    "React-Stockcharts-SVG": {
        "tags": ["all", "default", "chart", "svg"],
        "enabled": True
    },
    "Perf-Dashboard": {
        "tags": ["all", "default", "chart", "webcomponents"],
        "enabled": True
    }
}


class Speedometer30Story(SpeedometerStory):
  __doc__ = SpeedometerStory.__doc__
  NAME: str = "speedometer_3.0"
  URL: str = "https://chromium-workloads.web.app/speedometer/v3.0/"
  URL_OFFICIAL: str = "https://browserbench.org/Speedometer3.0/"
  URL_LOCAL: str = "http://127.0.0.1:7000"
  SUBSTORIES: Tuple[str, ...] = tuple(SPEEDOMETER_3_STORY_DATA.keys())

  @classmethod
  def default_story_names(cls) -> Tuple[str, ...]:
    return tuple(
        tuple(name for name, data in SPEEDOMETER_3_STORY_DATA.items()
              if data["enabled"]))

  def __init__(self,
               substories: Sequence[str] = (),
               iterations: Optional[int] = None,
               sync_wait: Optional[dt.timedelta] = None,
               sync_warmup: Optional[dt.timedelta] = None,
               measurement_method: Optional[MeasurementMethod] = None,
               viewport: Optional[vp.Viewport] = None,
               shuffle_seed: ShuffleSeedT = None,
               url: Optional[str] = None):
    self._sync_wait = DurationParser.positive_or_zero_duration(
        sync_wait or dt.timedelta(0), "sync_wait")
    self._sync_warmup = DurationParser.positive_or_zero_duration(
        sync_warmup or dt.timedelta(0), "sync_warmup")
    self._measurement_method: MeasurementMethod = (
        measurement_method or MeasurementMethod.RAF)
    self._viewport = None
    if viewport:
      self._viewport = vp.Viewport.parse_sized(viewport)
    self._shuffle_seed: ShuffleSeedT = parse_shuffle_seed(shuffle_seed)
    super().__init__(url=url, substories=substories, iterations=iterations)

  @property
  def single_substory_duration(self) -> dt.timedelta:
    return dt.timedelta(seconds=0.25)

  @property
  def sync_wait(self) -> dt.timedelta:
    return self._sync_wait

  @property
  def sync_warmup(self) -> dt.timedelta:
    return self._sync_warmup

  @property
  def measurement_method(self) -> MeasurementMethod:
    return self._measurement_method

  @property
  def viewport(self) -> Optional[vp.Viewport]:
    return self._viewport

  @property
  def shuffle_seed(self) -> ShuffleSeedT:
    return self._shuffle_seed

  @property
  def url_params(self) -> Dict[str, str]:
    url_params = super().url_params
    if sync_wait := self.sync_wait:
      url_params["waitBeforeSync"] = str(to_ms(sync_wait))
    if sync_warmup := self.sync_warmup:
      url_params["warmupBeforeSync"] = str(to_ms(sync_warmup))
    if self.measurement_method != MeasurementMethod.RAF:
      url_params["measurementMethod"] = str(self.measurement_method)
    if viewport := self.viewport:
      url_params["viewport"] = f"{viewport.width}x{viewport.height}"
    if self.shuffle_seed is not None:
      url_params["shuffleSeed"] = str(self.shuffle_seed)
    return url_params

  def log_run_test_url(self, run: Run) -> None:
    del run
    params = self.url_params
    params["suites"] = ",".join(self.substories)
    params["developerMode"] = "true"
    params["startAutomatically"] = "true"
    official_test_url = helper.update_url_query(self.URL, params)
    logging.info("STORY PUBLIC TEST URL: %s", official_test_url)


class Speedometer3BenchmarkStoryFilter(SpeedometerBenchmarkStoryFilter):
  __doc__ = SpeedometerBenchmarkStoryFilter.__doc__

  @classmethod
  def add_cli_parser(
      cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
    parser = super().add_cli_parser(parser)
    parser.add_argument(
        "--sync-wait",
        default=dt.timedelta(0),
        type=DurationParser.positive_or_zero_duration,
        help="Add a custom wait timeout before each sync step.")
    parser.add_argument(
        "--sync-warmup",
        default=dt.timedelta(0),
        type=DurationParser.positive_or_zero_duration,
        help="Run a warmup loop for the given duration before each sync step.")

    measurement_method_group = parser.add_argument_group(
        "Measurement Method Option")
    measurement_method_group = parser.add_mutually_exclusive_group()
    measurement_method_group.add_argument(
        "--raf",
        dest="measurement_method",
        default=MeasurementMethod.RAF,
        const=MeasurementMethod.RAF,
        action="store_const",
        help=("Use the default requestAnimationFrame-based approach "
              "for async time measurement."))
    measurement_method_group.add_argument(
        "--timer",
        dest="measurement_method",
        const=MeasurementMethod.TIMER,
        action="store_const",
        help=("Use the 'classical' setTimeout-based approach "
              "for async time measurement. "
              "This might omit measuring some async work."))

    parser.add_argument(
        "--story-viewport",
        type=vp.Viewport.parse_sized,
        help="Specify the speedometer workload viewport size.")
    parser.add_argument(
        "--shuffle-seed",
        type=parse_shuffle_seed,
        help=("Set a shuffle seed to run the stories in a"
              "non-default order."))
    return parser

  @classmethod
  def kwargs_from_cli(cls, args: argparse.Namespace) -> Dict[str, Any]:
    kwargs = super().kwargs_from_cli(args)
    kwargs["iterations"] = args.iterations
    kwargs["measurement_method"] = args.measurement_method
    kwargs["sync_wait"] = args.sync_wait
    kwargs["sync_warmup"] = args.sync_warmup
    kwargs["viewport"] = args.story_viewport
    kwargs["shuffle_seed"] = args.shuffle_seed
    return kwargs

  def __init__(self,
               story_cls: Type[SpeedometerStory],
               patterns: Sequence[str],
               separate: bool = False,
               url: Optional[str] = None,
               iterations: Optional[int] = None,
               measurement_method: Optional[MeasurementMethod] = None,
               sync_wait: Optional[dt.timedelta] = None,
               sync_warmup: Optional[dt.timedelta] = None,
               viewport: Optional[vp.Viewport] = None,
               shuffle_seed: ShuffleSeedT = None):
    self.measurement_method = measurement_method
    self.sync_wait = sync_wait
    self.sync_warmup = sync_warmup
    self.viewport = viewport
    self.shuffle_seed: ShuffleSeedT = shuffle_seed
    assert issubclass(story_cls, Speedometer30Story)
    super().__init__(story_cls, patterns, separate, url, iterations=iterations)

  def create_stories_from_names(self, names: List[str],
                                separate: bool) -> Sequence[SpeedometerStory]:
    return self.story_cls.from_names(
        names,
        separate=separate,
        url=self.url,
        iterations=self.iterations,
        measurement_method=self.measurement_method,
        sync_wait=self.sync_wait,
        sync_warmup=self.sync_warmup,
        viewport=self.viewport,
        shuffle_seed=self.shuffle_seed)


class Speedometer30Benchmark(SpeedometerBenchmark):
  """
  Benchmark runner for Speedometer 3.0
  """
  NAME: str = "speedometer_3.0"
  DEFAULT_STORY_CLS = Speedometer30Story
  STORY_FILTER_CLS = Speedometer3BenchmarkStoryFilter
  PROBES: ProbeClsTupleT = (Speedometer30Probe,)

  @classmethod
  def version(cls) -> Tuple[int, ...]:
    return (3, 0)

  @classmethod
  def aliases(cls) -> Tuple[str, ...]:
    return ("sp3", "speedometer_3") + super().aliases()

  @classmethod
  def add_cli_parser(
      cls, subparsers: argparse.ArgumentParser, aliases: Sequence[str] = ()
  ) -> CrossBenchArgumentParser:
    parser = super().add_cli_parser(subparsers, aliases)
    parser.add_argument(
        "--detailed-metrics",
        "--details",
        default=False,
        action="store_true",
        help="Report more detailed internal metrics.")
    return parser

  @classmethod
  def kwargs_from_cli(cls, args: argparse.Namespace) -> Dict[str, Any]:
    kwargs = super().kwargs_from_cli(args)
    kwargs["detailed_metrics"] = args.detailed_metrics
    return kwargs

  def __init__(self,
               stories: Sequence[Story],
               custom_url: Optional[str] = None,
               detailed_metrics: bool = False):
    self._detailed_metrics = detailed_metrics
    super().__init__(stories, custom_url)

  @property
  def detailed_metrics(self) -> bool:
    return self._detailed_metrics