# Copyright (c) 2022 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

import itertools
from abc import abstractmethod
from datetime import datetime
from decimal import Decimal
from enum import Enum
from random import randint
from typing import (Any, Callable, Iterable, List, Mapping, Optional, Tuple,
                    Type, Union)

from dateutil.tz import tzutc
from faker import Faker
from toolz.dicttoolz import itemfilter, keyfilter

from czt.datagenerator.exceptions import DoesNotApply
from czt.datagenerator.keys import KeyDefinition, KeyGenerator


class TestDataType(Enum):
    SINGLE_VALID_INSTANCE = 1
    VALID_DATA_TEST_CASES = 2
    INVALID_DATA_TEST_CASES = 3


class DataGenerator:
    '''
    Base class for all data generators

    valid_values (optional) - If specified, all valid values will be randomly selected from the list. If not specified
        all valid values will be generated
    invalid_values (optional) - If specified, all invalid values will be randomly selected from the list. If not specified
        all invalid values will be generated
    can_be_empty - If set to true, an empty value (i.e. '') will be returned as one of the valid test cases
        otherwise it will be returned as an invalid test case.
    '''
    _fake: Faker
    _can_be_empty: bool
    _can_be_none: bool = False
    _valid_values: Union[List[Any], Mapping[Any, List[Any]]] = []
    _invalid_values: Union[List[Any], Mapping[Any, List[Any]]] = []

    @property
    def can_be_none(self):
        return self._can_be_none

    @can_be_none.setter
    def can_be_none(self, value):
        self._can_be_none = value

    def __init__(self, *, valid_values: Optional[List[Any]] = None, invalid_values: Optional[List[Any]] = None, can_be_empty=False) -> None:
        self._can_be_empty = can_be_empty
        if valid_values or invalid_values:
            self._valid_values = valid_values or []
            self._invalid_values = invalid_values or []
            self._fake = Faker()
        else:
            self._fake = Faker()

    def __call__(self, test_data_type: TestDataType = TestDataType.SINGLE_VALID_INSTANCE, skip_none_test: bool = True) -> Union[List, Any]:
        if test_data_type == TestDataType.SINGLE_VALID_INSTANCE:
            return self._valid_values[randint(0, len(self._valid_values) - 1)] if self._valid_values else self.generate_valid()
        elif test_data_type == TestDataType.VALID_DATA_TEST_CASES:
            return [
                *self._get_nonetype_test_case(test_data_type, skip_none_test),
                *self._get_empty_test_case(test_data_type),
                *self._get_all_valid_values()
            ]
        elif test_data_type == TestDataType.INVALID_DATA_TEST_CASES:
            return [
                *self._get_nonetype_test_case(test_data_type, skip_none_test),
                *self._get_empty_test_case(test_data_type),
                *self._get_all_invalid_values()
            ]
        else:
            raise Exception(f'{test_data_type} is not a valid TestDataType')

    def _get_nonetype_test_case(self, test_data_type: TestDataType, skip_none_test: bool):
        return [] if skip_none_test or ((test_data_type == TestDataType.INVALID_DATA_TEST_CASES) == self.can_be_none) else [None]

    def _get_all_valid_values(self):
        if not self._valid_values:
            return [self.generate_valid()]

        return self._valid_values

    def _get_all_invalid_values(self):
        try:
            if not self._invalid_values:
                return [self.generate_invalid()]

            return self._invalid_values
        except DoesNotApply:
            return []

    def _get_empty_test_case(self, test_data_type: TestDataType):
        try:
            if (self._can_be_empty and test_data_type == TestDataType.INVALID_DATA_TEST_CASES) or \
               (not self._can_be_empty and test_data_type == TestDataType.VALID_DATA_TEST_CASES):
                return []

            return [self.get_empty_value()]
        except DoesNotApply:
            return []

    @abstractmethod
    def generate_valid(self):
        '''
        This method must be overridden to generate valid values unless the 'valid_values' list was passed.
        '''
        ...

    def generate_invalid(self):
        '''
        This method can be overridden to generate invalid values unless the 'invalid_values' list was passed.
        If it is not overridden and no 'invalid_values' are passed, then invalid data tests are bypassed for this item.
        '''
        raise DoesNotApply()

    def get_empty_value(self):
        '''
        This method can be overridden to generate the 'empty' test item.
        If it is not overridden empty item tests are bypassed for this item.
        '''
        raise DoesNotApply()


class _ConstantGenerator(DataGenerator):
    '''
    This class is mostly used internally to wrap constants specified in a test data definition.
    '''
    def __init__(self, value: Any) -> None:
        super().__init__(valid_values=[value])

    def generate_valid(self):
        return self._fake.valid_value()


class TextGenerator(DataGenerator):
    '''
    This generator will generate string values with an optional prefix value.
    By default, this generator will not
    '''
    prefix: Optional[str]

    def __init__(self, prefix=None, *, valid_values: Optional[List[str]] = None, invalid_values: Optional[List[str]] = None, can_be_empty=False) -> None:
        self.prefix = prefix

        super().__init__(valid_values=valid_values, invalid_values=invalid_values, can_be_empty=can_be_empty)

    def generate_valid(self):
        value = self._fake.pystr()
        return f'{self.prefix}-{value}' if self.prefix else value

    def get_empty_value(self):
        return ''


Number = Union[int, float, Decimal]


class NumberGenerator(DataGenerator):
    '''
    This class will generate either floats or Decimal numbers.
    Invalid values can be generated either randomly generating numbers outside a specified range,
    or from a list of invalid values passed in.
    '''
    min_value: Optional[Number]
    max_value: Optional[Number]
    left_digits: int
    decimal_places: int
    type: Type

    def __init__(self, *, type: Type = float, precision: int = 8, decimal_places: int = 2, min_value: Optional[Number] = None, max_value: Optional[Number] = None, additional_invalid_values=None) -> None:
        super().__init__(invalid_values=['', 'abcd', '00.00.00', *(additional_invalid_values or [])])
        self.min_value = min_value
        self.max_value = max_value
        min_or_max = min_value or max_value
        self.left_digits = int((precision - decimal_places)) if not min_or_max else len(str(min_or_max).split('.')[0])
        self.decimal_places = decimal_places
        self.type = type

    def generate_valid(self):
        value = self._fake.pydecimal(left_digits=self.left_digits, right_digits=self.decimal_places, min_value=self.min_value, max_value=self.max_value)
        return self.type(value)


class CurrencyGenerator(NumberGenerator):
    '''
    This class will generate either floats or Decimal numbers as currency (i.e. 2 decimal places).
    Invalid values can be generated either randomly generating numbers outside a specified range,
    or from a list of invalid values passed in.
    '''
    def __init__(self, *, type: Type = float, precision: int = 8, min_value: Optional[float] = None, max_value: Optional[float] = None, additional_invalid_values=None) -> None:
        super().__init__(type=type, precision=precision, decimal_places=2, min_value=min_value, max_value=max_value, additional_invalid_values=additional_invalid_values)


class IntegerGenerator(DataGenerator):
    '''
    This class will generate integers.
    Invalid values can be generated either randomly generating numbers outside a specified range,
    or from a list of invalid values passed in.
    '''

    min_value: int
    max_value: int

    def __init__(self, *, min_value: int = 0, max_value: int = 9999, additional_invalid_values=None) -> None:
        super().__init__(invalid_values=['', 'abcd', '5.2', *(additional_invalid_values or [])])
        self.min_value = min_value
        self.max_value = max_value

    def generate_valid(self):
        return self._fake.pyint(min_value=self.min_value, max_value=self.max_value)


class DateTimeType(Enum):
    datetime = 0
    isoformat = 1
    timestamp = 2


class DateTimeGenerator(DataGenerator):
    '''
    This class will generate random date/times between a set of limits. They can be returned as
    datetime classes, iso format, or timestamps.
    '''
    _invalid_values: Mapping[DateTimeType, List] = {
        DateTimeType.datetime: [],
        DateTimeType.isoformat: ['', 'Not a date time', '2021:01:02T00:00:00'],
        DateTimeType.timestamp: [0, -1]
    }

    _output_gen = {
        DateTimeType.datetime: lambda x: datetime.fromtimestamp(x, tz=tzutc()),
        DateTimeType.isoformat: lambda x: datetime.fromtimestamp(x, tz=tzutc()).isoformat(),
        DateTimeType.timestamp: lambda x: x
    }

    _output_type: DateTimeType
    _timestamp_generator: NumberGenerator

    def __init__(self, output_type: DateTimeType, *, min_datetime: datetime, max_datetime: Optional[datetime] = None, invalid_values: Optional[List[Any]] = None, add_out_of_range_values=False) -> None:
        super().__init__(invalid_values=[
            *self._invalid_values[output_type],
            *(invalid_values or []),
            *([self._output_gen[output_type](NumberGenerator(max_value=min_datetime.timestamp() - 100)()) for x in range(3)] if add_out_of_range_values and min_datetime is not None else []),
            *([self._output_gen[output_type](NumberGenerator(min_value=max_datetime.timestamp() + 100)()) for x in range(3)] if add_out_of_range_values and max_datetime is not None else [])
        ])
        self._output_type = output_type
        self._timestamp_generator = NumberGenerator(min_value=min_datetime.timestamp(), max_value=(max_datetime.timestamp() if max_datetime else None))

    def generate_valid(self):
        timestamp = self._timestamp_generator()
        return self._output_gen[self._output_type](timestamp)


class ListGenerator(DataGenerator):
    '''
    This class will generate a list of objects. Invalid data test cases will only be applied
    to a single random item in the list.
    '''
    min_elements: int
    max_elements: int
    item_definition: DataGenerator

    def __init__(self, item_definition, min_elements=0, max_elements=20, *, required=False) -> None:
        super().__init__()
        self.item_definition = _get_generator(item_definition, required)
        self.min_elements = min_elements
        self.max_elements = max_elements

    def _get_all_valid_values(self):
        return [self._gen_list(test_item) for test_item in self.item_definition(TestDataType.VALID_DATA_TEST_CASES)]

    def _get_all_invalid_values(self):
        return [self._gen_list(test_item) for test_item in self.item_definition(TestDataType.INVALID_DATA_TEST_CASES)]

    def generate_valid(self):
        return self._gen_list()

    def generate_invalid(self):
        raise Exception('This should never be called')

    def _gen_list(self, test_item=None):
        item_count = self._fake.pyint(min_value=self.min_elements, max_value=self.max_elements)
        data = [self.item_definition(TestDataType.SINGLE_VALID_INSTANCE) for i in range(item_count)]
        if test_item:
            if item_count:
                test_item_idx = self._fake.pyint(min_value=0, max_value=(item_count - 1))
                data[test_item_idx] = test_item
            else:
                data.append(test_item)

        return data


class DictGenerator(DataGenerator):
    '''
    This class will generate a dictionary where each key/value can be generated. When retreiving different test case data,
    the test case will be applied for each field as appropriate. For example, if a single field indicates a set of valid values,
    then the dictionary will return at least that many test cases with each value (selected in random order).

    For valid test cases, the maximum number of test cases will be equal to the largest number of valid test cases for any field
    within that dictionary (or within subdictionaries). For example, if all fields have only 1 test value, then the dictionary will only return
    one test case data. If one field has 3 possible valid test values while another has 2 and the rest have 1, then a total of 3 test cases will
    be returned.

    For invalid test cases, the minimum number of test cases will be based on the number of fields in the dictionary (and any lists or dictionaries
    within that dictionary). At least one test case per field where the field is invalid, but all remaining fields are valid.
    '''
    _raw_definitions: dict
    _field_definitions: dict
    _fields_required: bool

    def __init__(self, field_definitions: dict, *, required=False) -> None:
        super().__init__()

        def _get_key_modifier(key_value):
            return key_value if isinstance(key_value, KeyDefinition) else KeyDefinition(key_value, required)

        key_generators = keyfilter(lambda k: isinstance(k, KeyGenerator), field_definitions)

        self._raw_definitions = field_definitions
        self._field_definitions = {
            **{_get_key_modifier(key): value for key, value in field_definitions.items() if not isinstance(key, KeyGenerator)},
            **{_get_key_modifier(key): value for key_gen, value in key_generators.items() for key in key_gen()}
        }

        self._fields_required = required

    def extend(self, field_definitions: dict, *, required=None):
        if required is None:
            required = self._fields_required

        return DictGenerator({
            **self._raw_definitions,
            **(field_definitions or {})
        }, required=required)

    def _get_special_test_cases(self, test_data_type: TestDataType):
        for field_to_test, _ in itemfilter(lambda item: self._is_testable(item[0], item[1]), self._field_definitions).items():
            field_to_test = field_to_test if isinstance(field_to_test, KeyDefinition) else KeyDefinition(field_to_test, is_required=False)

            if (not field_to_test.is_required and test_data_type == TestDataType.VALID_DATA_TEST_CASES) or \
               (field_to_test.is_required and test_data_type == TestDataType.INVALID_DATA_TEST_CASES):
                yield keyfilter(lambda k: k != field_to_test.value, self._generate_instance())

    def _get_all_valid_values(self):
        return [
            *self._get_special_test_cases(TestDataType.VALID_DATA_TEST_CASES),
            *self._generate_valid_test_cases()
        ]

    def _get_all_invalid_values(self):
        return [
            *self._get_special_test_cases(TestDataType.INVALID_DATA_TEST_CASES),
            *self._generate_invalid_test_cases()
        ]

    def _generate_invalid_test_cases(self):
        for field_to_test, field_definition in itemfilter(lambda item: self._is_testable(item[0], item[1]), self._field_definitions).items():
            try:
                test_values = self._call_generator(field_definition, TestDataType.INVALID_DATA_TEST_CASES)
                for test_value in test_values:
                    yield self._generate_instance({field_to_test.value: test_value})
            except DoesNotApply:
                pass

    def _call_generator(self, definition, test_data_type: TestDataType = TestDataType.SINGLE_VALID_INSTANCE) -> Iterable:
        return _get_generator(definition, self._fields_required)(test_data_type, skip_none_test=False)

    def _generate_valid_test_cases(self) -> Iterable:
        test_case_data = {
            k.value: iter(self._call_generator(v, TestDataType.VALID_DATA_TEST_CASES)) for k, v in itemfilter(lambda item: self._is_testable(item[0], item[1]), self._field_definitions).items()
        }

        def get_next_test_value_iteration(key, test_value_iterator):
            try:
                return {key: next(test_value_iterator)}
            except StopIteration:
                return {}

        while(True):
            test_data = {
                **{k: v for field, definition in test_case_data.items() for k, v in get_next_test_value_iteration(field, definition).items()}
            }

            if len(test_data) == 0:
                return

            yield self._generate_instance(test_data)

    def generate_valid(self):
        return self._generate_instance()

    def generate_invalid(self):
        raise Exception('This should never be called')

    def _generate_instance(self, test_data=None):
        test_data = test_data or {}

        return {
            **{field.value: self._call_generator(field_definition) for field, field_definition in self._field_definitions.items()},
            **(test_data)
        }

    def _is_testable(self, key_definition, field_definition):
        return not key_definition.skip and (isinstance(field_definition, Callable) or isinstance(field_definition, dict))


class All():
    _sub_generators: Tuple[DataGenerator]

    def __init__(self, *sub_generators) -> None:
        super().__init__()
        self._sub_generators = sub_generators  # type: ignore

    def __call__(self, test_data_type: TestDataType = TestDataType.SINGLE_VALID_INSTANCE, skip_none_test: bool = True):
        if test_data_type == TestDataType.SINGLE_VALID_INSTANCE:
            idx = randint(0, len(self._sub_generators) - 1)
            return self._sub_generators[idx](test_data_type, skip_none_test=skip_none_test)

        return itertools.chain.from_iterable([sub_generator(test_data_type, skip_none_test) for sub_generator in self._sub_generators])


def _get_generator(field_definition, required=False) -> DataGenerator:
    '''
    This will wrap various data types in an appropriate generator class.
    '''
    if isinstance(field_definition, DataGenerator):
        return field_definition
    elif isinstance(field_definition, dict):
        return DictGenerator(field_definition, required=required)
    else:
        return _ConstantGenerator(field_definition)


def Maybe(data_generator: DataGenerator):
    data_generator.can_be_none = True
    return data_generator
