#
# 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.
#

import logging
import os
import sys
import tempfile
import unittest

import mock
import yaml

import apache_beam as beam
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
from apache_beam.yaml import yaml_provider
from apache_beam.yaml.yaml_provider import YamlProviders
from apache_beam.yaml.yaml_transform import SafeLineLoader
from apache_beam.yaml.yaml_transform import YamlTransform


class WindowIntoTest(unittest.TestCase):
  def __init__(self, methodName="runWindowIntoTest"):
    unittest.TestCase.__init__(self, methodName)
    self.parse_duration = YamlProviders.WindowInto._parse_duration

  def test_parse_duration_ms(self):
    value = self.parse_duration('1000ms', 'size')
    self.assertEqual(1, value)

  def test_parse_duration_sec(self):
    value = self.parse_duration('1s', 'size')
    self.assertEqual(1, value)

  def test_parse_duration_min(self):
    value = self.parse_duration('1m', 'size')
    self.assertEqual(60, value)

  def test_parse_duration_hour(self):
    value = self.parse_duration('1h', 'size')
    self.assertEqual(3600, value)

  def test_parse_duration_from_decimal(self):
    value = self.parse_duration('1.5m', 'size')
    self.assertEqual(90, value)

  def test_parse_duration_to_decimal(self):
    value = self.parse_duration('1ms', 'size')
    self.assertEqual(0.001, value)

  def test_parse_duration_with_missing_suffix(self):
    value = self.parse_duration('1', 'size')
    self.assertEqual(1, value)

  def test_parse_duration_with_invalid_suffix(self):
    with self.assertRaises(ValueError):
      self.parse_duration('1x', 'size')

  def test_parse_duration_with_missing_value(self):
    with self.assertRaises(ValueError):
      self.parse_duration('s', 'size')


class ProviderParsingTest(unittest.TestCase):

  INLINE_PROVIDER = {'type': 'TEST', 'name': 'INLINED'}
  INCLUDED_PROVIDER = {'type': 'TEST', 'name': 'INCLUDED'}
  EXTRA_PROVIDER = {'type': 'TEST', 'name': 'EXTRA'}

  @classmethod
  def setUpClass(cls):
    cls.tempdir = tempfile.TemporaryDirectory()
    cls.to_include = os.path.join(cls.tempdir.name, 'providers.yaml')
    with open(cls.to_include, 'w') as fout:
      yaml.dump([cls.INCLUDED_PROVIDER], fout)
    cls.to_include_nested = os.path.join(
        cls.tempdir.name, 'nested_providers.yaml')
    with open(cls.to_include_nested, 'w') as fout:
      yaml.dump([{'include': './providers.yaml'}, cls.EXTRA_PROVIDER], fout)

  @classmethod
  def tearDownClass(cls):
    cls.tempdir.cleanup()

  @mock.patch(
      'apache_beam.yaml.yaml_provider.ExternalProvider.provider_from_spec',
      lambda _, x: x)
  def test_include_file(self):
    flattened = [
        SafeLineLoader.strip_metadata(spec)
        for spec in yaml_provider.parse_providers(
            '', [
                self.INLINE_PROVIDER,
                {
                    'include': self.to_include
                }, ])
    ]

    self.assertEqual([
        self.INLINE_PROVIDER,
        self.INCLUDED_PROVIDER,
    ],
                     flattened)

  @mock.patch(
      'apache_beam.yaml.yaml_provider.ExternalProvider.provider_from_spec',
      lambda _, x: x)
  def test_include_url(self):
    flattened = [
        SafeLineLoader.strip_metadata(spec)
        for spec in yaml_provider.parse_providers(
            '', [
                self.INLINE_PROVIDER,
                {
                    'include': 'file:///' + self.to_include
                }, ])
    ]

    self.assertEqual([
        self.INLINE_PROVIDER,
        self.INCLUDED_PROVIDER,
    ],
                     flattened)

  @mock.patch(
      'apache_beam.yaml.yaml_provider.ExternalProvider.provider_from_spec',
      lambda _, x: x)
  def test_nested_include(self):
    flattened = [
        SafeLineLoader.strip_metadata(spec)
        for spec in yaml_provider.parse_providers(
            '', [
                self.INLINE_PROVIDER,
                {
                    'include': self.to_include_nested
                }, ])
    ]

    self.assertEqual([
        self.INLINE_PROVIDER,
        self.INCLUDED_PROVIDER,
        self.EXTRA_PROVIDER,
    ],
                     flattened)


class YamlDefinedProider(unittest.TestCase):
  def test_yaml_define_provider(self):
    providers = '''
    - type: yaml
      transforms:
        Range:
          config_schema:
            properties:
              end: {type: integer}
          requires_inputs: false
          body: |
            type: Create
            config:
              elements:
                {% for ix in range(end) %}
                - {{ix}}
                {% endfor %}
        Power:
          config_schema:
            properties:
              n: {type: integer}
          body:
            type: MapToFields
            config:
              language: python
              append: true
              fields:
                power: "element**{{n}}"
    '''

    pipeline = '''
    type: chain
    transforms:
      - type: Range
        config:
          end: 4
      - type: Power
        config:
          n: 2
    '''

    with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
        pickle_library='cloudpickle')) as p:
      result = p | YamlTransform(
          pipeline,
          providers=yaml_provider.parse_providers(
              '', yaml.load(providers, Loader=SafeLineLoader)))
      assert_that(
          result | beam.Map(lambda x: (x.element, x.power)),
          equal_to([(0, 0), (1, 1), (2, 4), (3, 9)]))

  def test_recursive(self):
    providers = '''
    - type: yaml
      transforms:
        Factorial:
          config_schema:
            properties:
              n: {type: integer}
          requires_inputs: false
          body: |
            {% if n <= 1 %}
              type: Create
              config:
                elements:
                  - {value: 1}
            {% else %}
              type: chain
              transforms:
                - type: Factorial
                  config:
                    n: {{n-1}}
                - type: MapToFields
                  name: Multiply
                  config:
                    language: python
                    fields:
                      value: value * {{n}}
            {% endif %}
    '''

    pipeline = '''
    type: Factorial
    config:
      n: 5
    '''

    with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
        pickle_library='cloudpickle')) as p:
      result = p | YamlTransform(
          pipeline,
          providers=yaml_provider.parse_providers(
              '', yaml.load(providers, Loader=SafeLineLoader)))
      assert_that(result | beam.Map(lambda x: x.value), equal_to([120]))


class PythonProviderDepsTest(unittest.TestCase):
  def test_env_package_sensitive(self):
    self.assertNotEqual(
        yaml_provider.PypiExpansionService._key('base', ['pkg1']),
        yaml_provider.PypiExpansionService._key('base', ['pkg2']))

  def test_env_base_sensitive(self):
    self.assertNotEqual(
        yaml_provider.PypiExpansionService._key('base1', ['pkg']),
        yaml_provider.PypiExpansionService._key('base2', ['pkg']))

  def test_env_order_invariant(self):
    self.assertEqual(
        yaml_provider.PypiExpansionService._key('base', ['pkg1', 'pkg2']),
        yaml_provider.PypiExpansionService._key('base', ['pkg2', 'pkg1']))

  def test_env_path_invariant(self):
    with tempfile.TemporaryDirectory() as tmpdir:
      os.mkdir(os.path.join(tmpdir, 'a'))
      pkgA = os.path.join(tmpdir, 'a', 'pkg.tgz')
      os.mkdir(os.path.join(tmpdir, 'b'))
      pkgB = os.path.join(tmpdir, 'b', 'pkg.tgz')
      with open(pkgA, 'w') as fout:
        fout.write('content')
      with open(pkgB, 'w') as fout:
        fout.write('content')
      self.assertEqual(
          yaml_provider.PypiExpansionService._key('base', [pkgA]),
          yaml_provider.PypiExpansionService._key('base', [pkgB]))

  def test_env_content_sensitive(self):
    with tempfile.TemporaryDirectory() as tmpdir:
      pkg = os.path.join(tmpdir, 'pkg.tgz')
      with open(pkg, 'w') as fout:
        fout.write('content')
      before = yaml_provider.PypiExpansionService._key('base', [pkg])
      with open(pkg, 'w') as fout:
        fout.write('new content')
      after = yaml_provider.PypiExpansionService._key('base', [pkg])
      self.assertNotEqual(before, after)


class JoinUrlOrFilepathTest(unittest.TestCase):
  def test_join_url_relative_path(self):
    self.assertEqual(
        yaml_provider._join_url_or_filepath('http://example.com/a', 'b/c.yaml'),
        'http://example.com/b/c.yaml')
    self.assertEqual(
        yaml_provider._join_url_or_filepath(
            'http://example.com/a/', 'b/c.yaml'),
        'http://example.com/a/b/c.yaml')

    # use os.path.join to mock gcs filesystem split and join.
    with mock.patch('apache_beam.io.filesystems.FileSystems.split',
                    new=lambda x:
                    ("gs://bucket", x.removeprefix("gs://bucket/"))):
      with mock.patch('apache_beam.io.filesystems.FileSystems.join',
                      new=lambda *args: '/'.join(args)):
        self.assertEqual(
            yaml_provider._join_url_or_filepath('gs://bucket', 'b/c.yaml'),
            'gs://bucket/b/c.yaml')
        self.assertEqual(
            yaml_provider._join_url_or_filepath('gs://bucket/', 'b/c.yaml'),
            'gs://bucket/b/c.yaml')
        self.assertEqual(
            yaml_provider._join_url_or_filepath('gs://bucket/a', 'b/c.yaml'),
            'gs://bucket/b/c.yaml')

  def test_join_filepath_relative_path(self):
    if sys.platform != 'win32':
      self.assertEqual(
          yaml_provider._join_url_or_filepath('/a/b/', 'c/d.yaml'),
          '/a/b/c/d.yaml')
      self.assertEqual(
          yaml_provider._join_url_or_filepath('/a/b', 'c/d.yaml'),
          '/a/c/d.yaml')
    else:
      self.assertEqual(
          yaml_provider._join_url_or_filepath('C:\\a\\b\\', 'c\\d.yaml'),
          'C:\\a\\b\\c\\d.yaml')
      self.assertEqual(
          yaml_provider._join_url_or_filepath('C:\\a\\b', 'c\\d.yaml'),
          'C:\\a\\c\\d.yaml')

  def test_absolute_path(self):
    self.assertEqual(
        yaml_provider._join_url_or_filepath(
            'gs://bucket/a', 'gs://bucket/b/c.yaml'),
        'gs://bucket/b/c.yaml')

    if sys.platform != 'win32':
      self.assertEqual(
          yaml_provider._join_url_or_filepath('/a/b', '/c/d.yaml'), '/c/d.yaml')

  def test_different_scheme(self):
    self.assertEqual(
        yaml_provider._join_url_or_filepath(
            'http://example.com/a', 'gs://bucket/b/c.yaml'),
        'gs://bucket/b/c.yaml')

  def test_empty_base(self):
    self.assertEqual(
        yaml_provider._join_url_or_filepath('', 'a/b.yaml'), 'a/b.yaml')
    self.assertEqual(
        yaml_provider._join_url_or_filepath(None, 'a/b.yaml'), 'a/b.yaml')


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  unittest.main()


class YamlProvidersCreateTest(unittest.TestCase):
  def test_create_mixed_types(self):
    with beam.Pipeline() as p:
      # A mix of a primitive (Row(element=1)) and a dict (Row(a=2))
      result = p | YamlProviders.create([1, {"a": 2}])
      assert_that(
          result | beam.Map(lambda x: sorted(x._asdict().items())),
          equal_to([
              [('a', None), ('element', 1)],
              [('a', 2), ('element', None)],
          ]))
