#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# pytype: skip-file

import logging
import re
import time
from datetime import datetime
from typing import Any
from typing import Optional

from google.cloud import monitoring_v3
from google.protobuf.duration_pb2 import Duration

import apache_beam.testing.load_tests.dataflow_cost_consts as costs
from apache_beam.metrics.execution import MetricResult
from apache_beam.runners.dataflow.dataflow_runner import DataflowPipelineResult
from apache_beam.runners.dataflow.internal.apiclient import DataflowApplicationClient
from apache_beam.runners.runner import PipelineState
from apache_beam.testing.load_tests.load_test import LoadTest


class DataflowCostBenchmark(LoadTest):
  """Base class for Dataflow performance tests which export metrics to
  external databases: BigQuery or/and InfluxDB. Calculates the expected cost
  for running the job on Dataflow in region us-central1.

  Refer to :class:`~apache_beam.testing.load_tests.LoadTestOptions` for more
  information on the required pipeline options.

  If using InfluxDB with Basic HTTP authentication enabled, provide the
  following environment options: `INFLUXDB_USER` and `INFLUXDB_USER_PASSWORD`.

  If the hardware configuration for the job includes use of a GPU, please
  specify the version in use with the Accelerator enumeration. This is used to
  calculate the cost of the job later, as different accelerators have different
  billing rates per hour of use.
  """

  WORKER_START_PATTERN = re.compile(
      r'All workers have finished the startup processes and '
      r'began to receive work requests')
  WORKER_STOP_PATTERN = re.compile(r'Stopping worker pool')

  def __init__(
      self,
      metrics_namespace: Optional[str] = None,
      is_streaming: bool = False,
      gpu: Optional[costs.Accelerator] = None,
      pcollection: str = 'ProcessOutput.out0',
      subscription: Optional[str] = None):
    self.is_streaming = is_streaming
    self.gpu = gpu
    self.pcollection = pcollection
    self.subscription = subscription
    super().__init__(metrics_namespace=metrics_namespace)
    self.dataflow_client = DataflowApplicationClient(
        self.pipeline.get_pipeline_options())
    self.monitoring_client = monitoring_v3.MetricServiceClient()

  def run(self) -> None:
    try:
      self.test()
      if not hasattr(self, 'result'):
        self.result = self.pipeline.run()
      state = self.result.wait_until_finish(duration=self.timeout_ms)
      assert state != PipelineState.FAILED

      logging.info(
          'Pipeline complete, sleeping for 4 minutes to allow resource '
          'metrics to populate.')
      time.sleep(240)

      self.extra_metrics = self._retrieve_cost_metrics(self.result)
      additional_metrics = self._get_additional_metrics(self.result)
      self.extra_metrics.update(additional_metrics)

      logging.info(self.extra_metrics)
      self._metrics_monitor.publish_metrics(self.result, self.extra_metrics)
    except Exception:
      raise
    finally:
      self.cleanup()

  def _retrieve_cost_metrics(self,
                             result: DataflowPipelineResult) -> dict[str, Any]:
    """Calculates estimated cost based on pipeline resource usage."""
    job_id = result.job_id()
    metrics = result.metrics().all_metrics(job_id)
    metrics_dict = self._process_metrics_list(metrics)

    cost = 0.0
    if self.is_streaming:
      cost += metrics_dict.get(
          "TotalVcpuTime", 0.0) / 3600 * costs.VCPU_PER_HR_STREAMING
      cost += metrics_dict.get(
          "TotalMemoryUsage", 0.0) / 1000 / 3600 * costs.MEM_PER_GB_HR_STREAMING
      cost += metrics_dict.get(
          "TotalStreamingDataProcessed", 0.0) * costs.SHUFFLE_PER_GB_STREAMING
    else:
      cost += metrics_dict.get(
          "TotalVcpuTime", 0.0) / 3600 * costs.VCPU_PER_HR_BATCH
      cost += metrics_dict.get(
          "TotalMemoryUsage", 0.0) / 1000 / 3600 * costs.MEM_PER_GB_HR_BATCH
      cost += metrics_dict.get(
          "TotalStreamingDataProcessed", 0.0) * costs.SHUFFLE_PER_GB_BATCH

    if self.gpu:
      rate = costs.ACCELERATOR_TO_COST[self.gpu]
      cost += metrics_dict.get("TotalGpuTime", 0.0) / 3600 * rate

    cost += metrics_dict.get("TotalPdUsage", 0.0) / 3600 * costs.PD_PER_GB_HR
    cost += metrics_dict.get(
        "TotalSsdUsage", 0.0) / 3600 * costs.PD_SSD_PER_GB_HR

    metrics_dict["EstimatedCost"] = cost
    return metrics_dict

  def _process_metrics_list(self,
                            metrics: list[MetricResult]) -> dict[str, Any]:
    system_metrics = {}
    for entry in metrics:
      metric_key = entry.key
      metric = metric_key.metric
      if metric_key.step == '' and metric.namespace == 'dataflow/v1b3':
        system_metrics[metric.name] = entry.committed or 0.0
    return system_metrics

  def _get_worker_time_interval(
      self, job_id: str) -> tuple[Optional[str], Optional[str]]:
    """Extracts worker start and stop times from job messages."""
    start_time, end_time = None, None
    page_token = None
    all_messages = []
    last_message_time = None
    while True:
      messages, page_token = self.dataflow_client.list_messages(
        job_id=job_id,
        start_time=None,
        end_time=None,
        page_token=page_token,
        minimum_importance='JOB_MESSAGE_DEBUG')
      for message in messages:
        text = message.messageText
        if getattr(message, 'time', None):
          last_message_time = message.time
        if text:
          all_messages.append(text)
          if self.WORKER_START_PATTERN.search(text):
            start_time = message.time
            logging.info('Matched WORKER_START_PATTERN: %r', text)
          if self.WORKER_STOP_PATTERN.search(text):
            end_time = message.time
            logging.info('Matched WORKER_STOP_PATTERN: %r', text)
      if not page_token or (start_time and end_time):
        break
    if start_time and not end_time and self.is_streaming and last_message_time:
      end_time = last_message_time
      logging.info(
          'Using last job message time as end_time for streaming job: %s',
          end_time)
    if not start_time or not end_time:
      logging.warning(
          'Could not determine worker time interval. '
          'start_time=%s, end_time=%s, total messages=%d',
          start_time,
          end_time,
          len(all_messages))
    return start_time, end_time

  def _get_throughput_metrics(
      self,
      project: str,
      job_id: str,
      start_time: str,
      end_time: str,
      pcollection_name: Optional[str] = None,
  ) -> dict[str, float]:
    """Query Cloud Monitoring for per-PCollection throughput."""
    name = (
        pcollection_name if pcollection_name is not None else self.pcollection)
    interval = monitoring_v3.TimeInterval(
        start_time=start_time, end_time=end_time)
    aggregation = monitoring_v3.Aggregation(
        alignment_period=Duration(seconds=60),
        per_series_aligner=monitoring_v3.Aggregation.Aligner.ALIGN_MEAN)

    requests = {
        "Bytes": monitoring_v3.ListTimeSeriesRequest(
            name=f"projects/{project}",
            filter=f'metric.type='
            f'"dataflow.googleapis.com/job/estimated_byte_count" '
            f'AND metric.labels.job_id='
            f'"{job_id}" AND metric.labels.pcollection="{name}"',
            interval=interval,
            aggregation=aggregation),
        "Elements": monitoring_v3.ListTimeSeriesRequest(
            name=f"projects/{project}",
            filter=f'metric.type="dataflow.googleapis.com/job/element_count" '
            f'AND metric.labels.job_id="{job_id}" '
            f'AND metric.labels.pcollection="{name}"',
            interval=interval,
            aggregation=aggregation)
    }

    metrics = {}
    for key, req in requests.items():
      time_series = self.monitoring_client.list_time_series(request=req)
      values = [
          point.value.double_value for series in time_series
          for point in series.points
      ]
      metrics[f"AvgThroughput{key}"] = sum(values) / len(
          values) if values else 0.0

    return metrics

  def _get_streaming_throughput_metrics(
      self, project: str, start_time: str, end_time: str) -> dict[str, float]:
    if not self.subscription:
      return {'AvgThroughputBytes': 0.0, 'AvgThroughputElements': 0.0}

    sub_parts = self.subscription.split('/')
    subscription_id = sub_parts[-1] if sub_parts else self.subscription

    interval = monitoring_v3.TimeInterval(
        start_time=start_time, end_time=end_time)
    aggregation = monitoring_v3.Aggregation(
        alignment_period=Duration(seconds=60),
        per_series_aligner=monitoring_v3.Aggregation.Aligner.ALIGN_RATE)

    requests = {
        "Bytes": monitoring_v3.ListTimeSeriesRequest(
            name=f"projects/{project}",
            filter=f'metric.type='
            f'"pubsub.googleapis.com/subscription/byte_cost" '
            f'AND resource.labels.subscription_id="{subscription_id}"',
            interval=interval,
            aggregation=aggregation),
        "Elements": monitoring_v3.ListTimeSeriesRequest(
            name=f"projects/{project}",
            filter=f'metric.type='
            f'"pubsub.googleapis.com/subscription/sent_message_count" '
            f'AND resource.labels.subscription_id="{subscription_id}"',
            interval=interval,
            aggregation=aggregation),
    }

    metrics = {}
    for key, req in requests.items():
      time_series = list(self.monitoring_client.list_time_series(request=req))
      values = [
          point.value.double_value for series in time_series
          for point in series.points
      ]
      avg_rate = sum(values) / len(values) if values else 0.0
      metrics[f"AvgThroughput{key}"] = avg_rate
    return metrics

  def _get_job_runtime(self, start_time: str, end_time: str) -> float:
    """Calculates the job runtime duration in seconds."""
    start_dt = datetime.fromisoformat(start_time[:-1])
    end_dt = datetime.fromisoformat(end_time[:-1])
    return (end_dt - start_dt).total_seconds()

  def _get_additional_metrics(self,
                              result: DataflowPipelineResult) -> dict[str, Any]:
    job_id = result.job_id()
    job = self.dataflow_client.get_job(job_id)
    project = job.projectId
    start_time, end_time = self._get_worker_time_interval(job_id)
    if not start_time or not end_time:
      logging.warning('Could not find valid worker start/end times.')
      return {}

    runtime_seconds = self._get_job_runtime(start_time, end_time)
    if self.is_streaming:
      throughput_metrics = self._get_streaming_throughput_metrics(
          project, start_time, end_time)
    else:
      throughput_metrics = self._get_throughput_metrics(
          project, job_id, start_time, end_time)
      if (throughput_metrics.get('AvgThroughputBytes', 0) == 0 and
          throughput_metrics.get('AvgThroughputElements', 0) == 0):
        logging.warning(
            'No throughput data for PCollection "%s". Check Dataflow job %s '
            'graph for actual PCollection names (Runner v2 may use different '
            'naming).',
            self.pcollection,
            job_id)
    return {
        **throughput_metrics,
        "JobRuntimeSeconds": runtime_seconds,
    }
