#
# 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 re
import unittest

import hamcrest as hc
import pytest

import apache_beam as beam
from apache_beam import metrics
from apache_beam.metrics.cells import DistributionData
from apache_beam.metrics.execution import MetricKey
from apache_beam.metrics.execution import MetricsContainer
from apache_beam.metrics.execution import MetricsEnvironment
from apache_beam.metrics.metric import Lineage
from apache_beam.metrics.metric import MetricResults
from apache_beam.metrics.metric import Metrics
from apache_beam.metrics.metric import MetricsFilter
from apache_beam.metrics.metricbase import MetricName
from apache_beam.runners.direct.direct_runner import BundleBasedDirectRunner
from apache_beam.runners.worker import statesampler
from apache_beam.testing.metric_result_matchers import DistributionMatcher
from apache_beam.testing.metric_result_matchers import MetricResultMatcher
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
from apache_beam.utils import counters
from apache_beam.utils.histogram import LinearBucket


class NameTest(unittest.TestCase):
  def test_basic_metric_name(self):
    name = MetricName('namespace1', 'name1')
    self.assertEqual(name.namespace, 'namespace1')
    self.assertEqual(name.name, 'name1')
    self.assertEqual(name, MetricName('namespace1', 'name1'))

    key = MetricKey('step1', name)
    self.assertEqual(key.step, 'step1')
    self.assertEqual(key.metric.namespace, 'namespace1')
    self.assertEqual(key.metric.name, 'name1')
    self.assertEqual(key, MetricKey('step1', MetricName('namespace1', 'name1')))


class MetricResultsTest(unittest.TestCase):
  def test_metric_filter_namespace_matching(self):
    filter = MetricsFilter().with_namespace('ns1')
    name = MetricName('ns1', 'name1')
    key = MetricKey('step1', name)
    self.assertTrue(MetricResults.matches(filter, key))

  def test_metric_filter_name_matching(self):
    filter = MetricsFilter().with_name('name1').with_namespace('ns1')
    name = MetricName('ns1', 'name1')
    key = MetricKey('step1', name)
    self.assertTrue(MetricResults.matches(filter, key))

    filter = MetricsFilter().with_name('name1')
    name = MetricName('ns1', 'name1')
    key = MetricKey('step1', name)
    self.assertTrue(MetricResults.matches(filter, key))

  def test_metric_filter_step_matching(self):
    name = MetricName('ns1', 'name1')
    filter = MetricsFilter().with_step('Step1')

    key = MetricKey('Step1', name)
    self.assertTrue(MetricResults.matches(filter, key))

    key = MetricKey('Step10', name)
    self.assertFalse(MetricResults.matches(filter, key))

    key = MetricKey('Step10/Step1', name)
    self.assertTrue(MetricResults.matches(filter, key))

    key = MetricKey('Top1/Outer1/Inner1', name)

    filter = MetricsFilter().with_step('Top1/Outer1/Inner1')
    self.assertTrue(MetricResults.matches(filter, key))

    filter = MetricsFilter().with_step('Top1/Outer1')
    self.assertTrue(MetricResults.matches(filter, key))

    filter = MetricsFilter().with_step('Outer1/Inner1')
    self.assertTrue(MetricResults.matches(filter, key))

    filter = MetricsFilter().with_step('Top1/Inner1')
    self.assertFalse(MetricResults.matches(filter, key))


class MetricsTest(unittest.TestCase):
  def test_get_namespace_class(self):
    class MyClass(object):
      pass

    self.assertEqual(
        '{}.{}'.format(MyClass.__module__, MyClass.__name__),
        Metrics.get_namespace(MyClass))

  def test_get_namespace_string(self):
    namespace = 'MyNamespace'
    self.assertEqual(namespace, Metrics.get_namespace(namespace))

  def test_get_namespace_error(self):
    with self.assertRaises(ValueError):
      Metrics.get_namespace(object())

  def test_counter_empty_name(self):
    with self.assertRaises(ValueError):
      Metrics.counter("namespace", "")

  def test_counter_empty_namespace(self):
    with self.assertRaises(ValueError):
      Metrics.counter("", "names")

  def test_distribution_empty_name(self):
    with self.assertRaises(ValueError):
      Metrics.distribution("namespace", "")

  def test_distribution_empty_namespace(self):
    with self.assertRaises(ValueError):
      Metrics.distribution("", "names")

  # Do not change the behaviour of str(), do tno update/delete this test case
  # if the behaviour of str() is changed. Doing so will
  # break end user beam code which depends on the str() behaviour.
  def test_user_metric_name_str(self):
    mn = MetricName("my_namespace", "my_name")
    expected_str = 'MetricName(namespace=my_namespace, name=my_name)'
    self.assertEqual(str(mn), expected_str)

  def test_general_urn_metric_name_str(self):
    mn = MetricName(
        "my_namespace", "my_name", urn='my_urn', labels={'key': 'value'})
    expected_str = (
        "MetricName(namespace=my_namespace, name=my_name, "
        "urn=my_urn, labels={'key': 'value'})")
    self.assertEqual(str(mn), expected_str)

  @pytest.mark.it_validatesrunner
  def test_user_counter_using_pardo(self):
    class SomeDoFn(beam.DoFn):
      """A custom dummy DoFn using yield."""
      static_counter_elements = metrics.Metrics.counter(
          "SomeDoFn", 'metrics_static_counter_element')

      def __init__(self):
        self.user_counter_elements = metrics.Metrics.counter(
            self.__class__, 'metrics_user_counter_element')

      def process(self, element):
        self.static_counter_elements.inc(2)
        self.user_counter_elements.inc()
        distro = Metrics.distribution(self.__class__, 'element_dist')
        distro.update(element)
        yield element

    pipeline = TestPipeline()
    nums = pipeline | 'Input' >> beam.Create([1, 2, 3, 4])
    results = nums | 'ApplyPardo' >> beam.ParDo(SomeDoFn())
    assert_that(results, equal_to([1, 2, 3, 4]))

    res = pipeline.run()
    res.wait_until_finish()

    # Verify static counter.
    metric_results = (
        res.metrics().query(
            MetricsFilter().with_metric(SomeDoFn.static_counter_elements)))
    outputs_static_counter = metric_results['counters'][0]

    self.assertEqual(
        outputs_static_counter.key.metric.name,
        'metrics_static_counter_element')
    self.assertEqual(outputs_static_counter.committed, 8)

    # Verify user counter.
    metric_results = (
        res.metrics().query(
            MetricsFilter().with_name('metrics_user_counter_element')))
    outputs_user_counter = metric_results['counters'][0]

    self.assertEqual(
        outputs_user_counter.key.metric.name, 'metrics_user_counter_element')
    self.assertEqual(outputs_user_counter.committed, 4)

    # Verify user distribution counter.
    metric_results = res.metrics().query()
    matcher = MetricResultMatcher(
        step=hc.contains_string('ApplyPardo'),
        namespace=hc.contains_string('SomeDoFn'),
        name='element_dist',
        committed=DistributionMatcher(
            sum_value=hc.greater_than_or_equal_to(0),
            count_value=hc.greater_than_or_equal_to(0),
            min_value=hc.greater_than_or_equal_to(0),
            max_value=hc.greater_than_or_equal_to(0)))
    hc.assert_that(metric_results['distributions'], hc.has_item(matcher))

  def test_create_counter_distribution(self):
    sampler = statesampler.StateSampler('', counters.CounterFactory())
    statesampler.set_current_tracker(sampler)
    state1 = sampler.scoped_state(
        'mystep', 'myState', metrics_container=MetricsContainer('mystep'))

    try:
      sampler.start()
      with state1:
        counter_ns = 'aCounterNamespace'
        distro_ns = 'aDistributionNamespace'
        name = 'a_name'
        counter = Metrics.counter(counter_ns, name)
        distro = Metrics.distribution(distro_ns, name)
        counter.inc(10)
        counter.dec(3)
        distro.update(10)
        distro.update(2)
        self.assertTrue(isinstance(counter, Metrics.DelegatingCounter))
        self.assertTrue(isinstance(distro, Metrics.DelegatingDistribution))

        del distro
        del counter

        container = MetricsEnvironment.current_container()
        self.assertEqual(
            container.get_counter(MetricName(counter_ns,
                                             name)).get_cumulative(),
            7)
        self.assertEqual(
            container.get_distribution(MetricName(distro_ns,
                                                  name)).get_cumulative(),
            DistributionData(12, 2, 2, 10))
    finally:
      sampler.stop()


class LineageTest(unittest.TestCase):
  def test_fq_name(self):
    test_cases = {
        "apache-beam": "apache-beam",
        "`apache-beam`": "`apache-beam`",
        "apache.beam": "`apache.beam`",
        "apache:beam": "`apache:beam`",
        "apache beam": "`apache beam`",
        "`apache beam`": "`apache beam`",
        "apache\tbeam": "`apache\tbeam`",
        "apache\nbeam": "`apache\nbeam`"
    }
    for k, v in test_cases.items():
      self.assertEqual("apache:" + v, Lineage.get_fq_name("apache", k))
      self.assertEqual(
          "apache:beam:" + v, Lineage.get_fq_name("apache", k, subtype="beam"))
      self.assertEqual(
          "apache:beam:" + v + '.' + v,
          Lineage.get_fq_name("apache", k, k, subtype="beam"))

  def test_add(self):
    lineage = Lineage(Lineage.SOURCE)
    added = set()
    # override
    lineage.metric = added
    lineage.add("s", "1", "2")
    lineage.add("s:3.4")
    lineage.add("s", "5", "6.7")
    lineage.add("s", "1", "2", subtype="t")
    lineage.add("sys", "seg1", "seg2", "seg3/part2/part3", last_segment_sep='/')
    self.assertSetEqual(
        added,
        {('s:', '1.', '2'), ('s:3.4:', ), ('s:', '5.', '6.7'),
         ('s:', 't:', '1.', '2'),
         ('sys:', 'seg1.', 'seg2.', 'seg3/', 'part2/', 'part3')})


class HistogramTest(unittest.TestCase):
  def test_histogram(self):
    class WordExtractingDoFn(beam.DoFn):
      def __init__(self):
        self.word_lengths_dist = Metrics.histogram(
            self.__class__,
            'latency_histogram_ms',
            LinearBucket(0, 1, num_buckets=10))

      def process(self, element):
        text_line = element.strip()
        words = re.findall(r'[\w\']+', text_line, re.UNICODE)
        for w in words:
          self.word_lengths_dist.update(len(w))
        return words

    with beam.Pipeline(runner=BundleBasedDirectRunner()) as p:
      lines = p | 'read' >> beam.Create(["x x x yyyyyy yyyyyy yyyyyy"])
      _ = (
          lines
          | 'split' >>
          (beam.ParDo(WordExtractingDoFn()).with_output_types(str)))

    result = p.result

    filter = MetricsFilter().with_name('latency_histogram_ms')
    query_result = result.metrics().query(filter)
    histogram = query_result['histograms'][0].committed.histogram
    assert histogram._buckets == {1: 3, 6: 3}
    assert histogram.total_count() == 6
    assert 1 < histogram.get_linear_interpolation(0.50) < 3
    assert histogram.get_linear_interpolation(0.99) > 3


if __name__ == '__main__':
  unittest.main()