# coding: utf-8
#
# Copyright 2021 The Oppia Authors. All Rights Reserved.
#
# Licensed 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.

"""Utilities for running Apache Beam tests."""

from __future__ import annotations

import ast
import contextlib
import datetime
import re

from core.jobs import base_jobs
from core.jobs import job_options
from core.jobs.types import job_run_result
from core.platform import models
from core.tests import test_utils

import apache_beam as beam
from apache_beam import runners
from apache_beam.testing import test_pipeline
from apache_beam.testing import util as beam_testing_util

from typing import Any, Iterator, Optional, Sequence, Type

MYPY = False
if MYPY:  # pragma: no cover
    from mypy_imports import base_models
    from mypy_imports import datastore_services

(base_models,) = models.Registry.import_models([models.Names.BASE_MODEL])

datastore_services = models.Registry.import_datastore_services()


class PipelinedTestBase(test_utils.AppEngineTestBase):
    """Base class that runs tests within the context of a TestPipeline."""

    # TODO(#11464): Find a newer version of Apache Beam that fixes
    # GroupIntoBatches() to provide correct type info, so we don't have to
    # provide this hook for tests to override.
    RUNTIME_TYPE_CHECK = True

    # Helpful constants used by tests to create models.
    NOW = datetime.datetime.utcnow()
    YEAR_AGO = NOW - datetime.timedelta(weeks=52)
    YEAR_LATER = NOW + datetime.timedelta(weeks=52)

    # Here we use type Any because we need to match the behavior of super
    # class's constructor and super class's constructor can accept arbitrary
    # number of arguments with different types of values.
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.pipeline = test_pipeline.TestPipeline(
            runner=runners.DirectRunner(),
            options=job_options.JobOptions(namespace=self.namespace))
        self._pipeline_context_stack: Optional[contextlib.ExitStack] = None

    def setUp(self) -> None:
        super().setUp()
        with contextlib.ExitStack() as pipeline_context_stack:
            pipeline_context_stack.enter_context(decorate_beam_errors())
            pipeline_context_stack.enter_context(self.pipeline)
            self._pipeline_context_stack = pipeline_context_stack.pop_all()

    def tearDown(self) -> None:
        try:
            self._exit_pipeline_context()
        finally:
            super().tearDown()

    def assert_pcoll_equal(
        self, actual: beam.PCollection, expected: beam.PCollection
    ) -> None:
        """Asserts that the given PCollections are equal.
        NOTE: At most one PCollection assertion can be called in a test. This is
        because running assertions on pipelines requires flushing it and waiting
        for it to run to completion. If another assertion needs to be run, then
        the pipeline must be populated with values all over again (which is
        equivalent to writing a new test case anyway).

        Args:
            actual: PCollection. The PCollection generated by the test.
            expected: PCollection. A PCollection with the expected values.

        Raises:
            RuntimeError. A PCollection assertion has already been called.
        """
        self._assert_pipeline_context_is_acquired()
        beam_testing_util.assert_that(
            actual, beam_testing_util.equal_to(expected))
        self._exit_pipeline_context()

    def assert_pcoll_empty(self, actual: beam.PCollection) -> None:
        """Asserts that the given PCollection is empty.
        NOTE: At most one PCollection assertion can be called in a test. This is
        because running assertions on pipelines requires flushing it and waiting
        for it to run to completion. If another assertion needs to be run, then
        the pipeline must be populated with values all over again (which is
        equivalent to writing a new test case anyway).

        Args:
            actual: PCollection. The PCollection generated by the test.

        Raises:
            RuntimeError. A PCollection assertion has already been called.
        """
        self._assert_pipeline_context_is_acquired()
        beam_testing_util.assert_that(actual, beam_testing_util.is_empty())
        self._exit_pipeline_context()

    # Here we use type Any because this method can accept different properties
    # of models and those properties can be of type str, int, bool, Dict and
    # other types too. So, to allow every type of property we used Any here.
    def create_model(
        self,
        model_class: Type[base_models.SELF_BASE_MODEL],
        **properties: Any
    ) -> base_models.SELF_BASE_MODEL:
        """Helper method for creating valid models with common default values.

        Args:
            model_class: *. A subclass of BaseModel.
            **properties: dict(str: *). Properties to assign to the model. By
                default, this method will try to fill the required properties
                with default values.

        Returns:
            *. A new instance of the given model type.

        Raises:
            ValueError. A required property's default value is invalid.
        """
        property_values = {
            p._name: p._default for p in model_class._properties.values() # pylint: disable=protected-access
            if p._required # pylint: disable=protected-access
        }
        property_values['created_on'] = self.YEAR_AGO
        property_values['last_updated'] = self.YEAR_AGO
        property_values.update(properties)
        return model_class(**property_values)

    def _assert_pipeline_context_is_acquired(self) -> None:
        """Raises a RuntimeError when the pipeline context hasn't been entered.

        Raises:
            RuntimeError. The error.
        """
        if self._pipeline_context_stack is None:
            raise RuntimeError(
                'PCollection assertions must be run in the pipeline context.\n'
                '\n'
                'NOTE: This error most likely means you have called more than '
                'one PCollection assertion, which is forbidden. This is '
                'because running assertions on pipelines require us to wait '
                'for it to finish processing all of its data, after which '
                'there is nothing left to inspect. If you need to make '
                'multiple assertions, then split them into separate test '
                'cases.')

    def _exit_pipeline_context(self) -> None:
        """Flushes the pipeline and waits for it to finish running."""
        if self._pipeline_context_stack is not None:
            self._pipeline_context_stack.close()
            self._pipeline_context_stack = None


class JobTestBase(PipelinedTestBase):
    """Base class with helpful methods for testing Oppia's jobs.
    Subclasses must add the class constant JOB_CLASS to use the helper methods.
    """

    # NOTE: run() raises a NotImplementedError.
    JOB_CLASS: Type[base_jobs.JobBase] = base_jobs.JobBase

    # Here we use type Any because we need to match the behavior of super
    # class's constructor and super class's constructor can accept arbitrary
    # number of arguments with different types of values.
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.job = self.JOB_CLASS(self.pipeline)

    def run_job(self) -> beam.PCollection[job_run_result.JobRunResult]:
        """Runs a new instance of self.JOB_CLASS and returns its output.
        Test authors should override this method if their jobs need arguments
        for their run() method, for example:
            class FooJob(JobBase):
                def run(self, model_kind):
                    pass
        Should override this method to provide a value for `model_kind`.

        Returns:
            PCollection. The output of the job.
        """
        job_results = self.job.run()
        # NDB operations in Beam do not properly update the context cache
        # (this cache is separate for every application thread), thus we clear
        # it ourselves.
        with datastore_services.get_ndb_context() as ndb_context:
            ndb_context.clear_cache()
        return job_results

    def put_multi(self, model_list: Sequence[base_models.BaseModel]) -> None:
        """Puts the input models into the datastore.

        Args:
            model_list: list(Model). The NDB models to put into the datastore.
        """
        datastore_services.update_timestamps_multi(
            model_list, update_last_updated_time=False)
        datastore_services.put_multi(model_list)

    def assert_job_output_is(self, expected: beam.PCollection) -> None:
        """Asserts the output of self.JOB_CLASS matches the given PCollection.

        Args:
            expected: PCollection. A PCollection with the expected values.
        """
        self.assert_pcoll_equal(self.run_job(), expected)

    def assert_job_output_is_empty(self) -> None:
        """Asserts that the output of self.JOB_CLASS is an empty PCollection."""
        self.assert_pcoll_empty(self.run_job())


@contextlib.contextmanager
def decorate_beam_errors() -> Iterator[None]:
    """Context manager to improve the readability of beam_testing_util errors.
    The beam_testing_util module raises exceptions with a single string of
    repr()'d lists as the message. The items end up appearing on one long line,
    making it difficult to read when the elements of the lists are very long
    (which they tend to be, especially for Oppia's audit errors).
    This context manager tries to split the list elements into lines so that
    it's easier to understand which errors occurred and why. If it cannot parse
    the message successfully, it will raise the error unchanged.

    Yields:
        None. Nothing.

    Raises:
        AssertionError. The decorated exception.
    """
    try:
        yield
    except beam_testing_util.BeamAssertException as exception:
        exception_message = str(exception)
        match = (
            re.match(
                r'.*'
                r', unexpected elements (?P<unexpected>.*)'
                r', missing elements (?P<missing>.*)'
                r' \[(?P<context>while running .*)\]',
                exception_message)
            or re.match(
                r'.*'
                r', unexpected elements (?P<unexpected>.*)'
                r' \[(?P<context>while running .*)\]',
                exception_message)
            or re.match(
                r'.*'
                r', missing elements (?P<missing>.*)'
                r' \[(?P<context>while running .*)\]',
                exception_message)
            or re.match(
                r'.*'
                r'\[\] == (?P<unexpected>.*)'
                r' \[(?P<context>while running .*)\]',
                exception_message)
        )

        if match:
            groupdict = match.groupdict()
        else:
            raise AssertionError(exception_message) from exception

        unexpected_elements = groupdict.get('unexpected', None)
        try:
            unexpected_elements = (
                ast.literal_eval(unexpected_elements)
                if unexpected_elements else None)
        except (SyntaxError, ValueError) as e:
            raise AssertionError(exception_message) from e

        missing_elements = groupdict.get('missing', None)
        try:
            missing_elements = (
                ast.literal_eval(missing_elements)
                if missing_elements else None)
        except (SyntaxError, ValueError) as e:
            raise AssertionError(exception_message) from e

        error_lines = [
            'failed %s' % match.group('context'),
            '',
        ]
        if unexpected_elements:
            error_lines.append('Unexpected:')
            error_lines.extend('    %r' % e for e in unexpected_elements)
        if unexpected_elements and missing_elements:
            error_lines.append('')
        if missing_elements:
            error_lines.append('Missing:')
            error_lines.extend('    %r' % e for e in missing_elements)
        error_lines.append('')
        raise AssertionError('\n'.join(error_lines)) from exception
