# Copyright 2025 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 abc
import datetime as dt
import enum
from typing import (TYPE_CHECKING, Any, ClassVar, MutableMapping, Optional,
                    Sequence, cast)

from typing_extensions import override

from crossbench.benchmarks.speedometer.speedometer import (
    SpeedometerBenchmark, SpeedometerBenchmarkStoryFilter, SpeedometerProbe,
    SpeedometerProbeContext, SpeedometerStory)
from crossbench.browsers import viewport as vp
from crossbench.helper import url_helper
from crossbench.parse import DurationParser, NumberParser, ObjectParser
from crossbench.str_enum_with_help import StrEnumWithHelp

if TYPE_CHECKING:
  import argparse

  from crossbench.action_runner.config import ActionRunnerConfig
  from crossbench.cli.parser import CrossBenchArgumentParser
  from crossbench.runner.actions import Actions
  from crossbench.stories.story import Story
  from crossbench.types import Json

  ShuffleSeedT = str | int | None


class Speedometer3Probe(SpeedometerProbe, metaclass=abc.ABCMeta):
  """
  Speedometer3-specific probe (compatible with v3 benchmarks).
  Extracts all speedometer times and scores.
  """

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

  @override
  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


class Speedometer3ProbeContext(SpeedometerProbeContext, metaclass=abc.ABCMeta):
  JS: ClassVar = "return JSON.stringify(window.benchmarkClient.metrics);"

  @override
  def to_json(self, actions: Actions) -> Json:
    json_data = super().to_json(actions)
    return ObjectParser.non_empty_dict(json_data, "speedometer metrics")

  @override
  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


@enum.unique
class MeasurementMethod(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 Speedometer3Story(SpeedometerStory, metaclass=abc.ABCMeta):
  __doc__ = SpeedometerStory.__doc__
  URL_LOCAL: ClassVar[str] = "http://127.0.0.1:8080"
  SUBSTORIES: ClassVar[tuple[str, ...]] = tuple(SPEEDOMETER_3_STORY_DATA.keys())

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

  @override
  def _wait_for_ready(self, actions: Actions) -> None:
    actions.wait_js_condition(
        "return !!window.benchmarkClient", 0.5, timeout=10)

  def _setup_substories(self, actions: Actions) -> None:
    # Handled via URL params
    pass

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

  @property
  @override
  def url_params(self) -> MutableMapping[str, str]:
    url_params: MutableMapping[str, str] = super().url_params
    if tuple(self.substories) != self.default_story_names():
      url_params["suites"] = ",".join(self.substories)
    return url_params

  @property
  def test_url(self) -> str:
    params: MutableMapping[str, str] = self.url_params
    params["developerMode"] = ""
    params["startAutomatically"] = ""
    official_test_url = url_helper.update_url_query(self.URL, params)
    return official_test_url


class Speedometer3BenchmarkStoryFilter(SpeedometerBenchmarkStoryFilter):
  __doc__ = SpeedometerBenchmarkStoryFilter.__doc__

  @classmethod
  @override
  def add_cli_arguments(
      cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
    parser = super().add_cli_arguments(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
  @override
  def url_params_from_cli(cls,
                          args: argparse.Namespace) -> MutableMapping[str, Any]:
    url_params: MutableMapping[str, str] = super().url_params_from_cli(args)
    if sync_wait := args.sync_wait:
      url_params["waitBeforeSync"] = str(to_ms(sync_wait))
    if sync_warmup := args.sync_warmup:
      url_params["warmupBeforeSync"] = str(to_ms(sync_warmup))
    if args.measurement_method != MeasurementMethod.RAF:
      url_params["measurementMethod"] = str(args.measurement_method)
    if viewport := args.story_viewport:
      url_params["viewport"] = f"{viewport.width}x{viewport.height}"
    if args.shuffle_seed is not None:
      url_params["shuffleSeed"] = str(args.shuffle_seed)
    return url_params


class Speedometer3Benchmark(SpeedometerBenchmark, metaclass=abc.ABCMeta):
  """
  Abstract benchmark runner for Speedometer 3.
  """
  STORY_FILTER_CLS: ClassVar = Speedometer3BenchmarkStoryFilter

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

  @classmethod
  @override
  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],
               action_runner_config: Optional[ActionRunnerConfig] = None,
               custom_url: Optional[str] = None,
               detailed_metrics: bool = False) -> None:
    self._detailed_metrics = detailed_metrics
    super().__init__(stories, action_runner_config, custom_url)

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