# Copyright (c) 2020 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 os
from collections import namedtuple
from functools import reduce
from random import randint
from typing import Callable, NamedTuple

import hypothesis.strategies as st
import pytest
import simplejson as json
from voluptuous import Any, Coerce, Match, Number, Optional, Required, Schema


class StreamingBody:
    def __init__(self, body):
        self.body = body

    def read(self):
        return json.dumps(self.body).encode()


# Voluptuous Hypothesis
def _from_schema(draw, s):
    if isinstance(s, dict):
        return reduce(lambda acc, i: {**acc, _from_schema(draw, i[0]): _from_schema(draw, i[1])}, s.items(), {})
    if isinstance(s, list):
        return reduce(lambda acc, v: acc + [_from_schema(draw, v), ], s, [])
    if isinstance(s, set):
        return reduce(lambda acc, v: acc | {_from_schema(draw, v), }, s, set())
    if isinstance(s, tuple):
        return reduce(lambda acc, v: acc + (_from_schema(draw, v),), s, ())
    if isinstance(s, Schema):
        return _from_schema(draw, s.__dict__['schema'])
    if isinstance(s, Required):
        return s.schema
    if isinstance(s, Optional):
        return s.schema
    if isinstance(s, Any):
        return _from_schema(draw, s.validators[randint(0, len(s.validators))])
    if isinstance(s, Coerce):
        return _from_schema(draw, s.type)
    if isinstance(s, Number):
        kwargs = {'allow_nan': False, 'allow_infinity': False}
        return draw(st.decimals(**kwargs)) if s.yield_decimal else draw(st.floats(**kwargs))
    if isinstance(s, Match):
        return draw(st.from_regex(s.pattern, fullmatch=True))
    if callable(s) and s.__name__ == 'Email':
        user = 'foo'
        domain = 'bar'
        return f'{user}@{domain}.com'
    if isinstance(s, str):
        return s
    if s is None:
        return None
    return draw(st.from_type(s))


def schema_to_strategy(schema: Schema) -> Callable:
    @st.composite
    def f(draw):
        return _from_schema(draw, schema.__dict__['schema'])

    return f


class Context:
    def __init__(self, mocker, module_under_test, modules_to_mock=()):
        self._mocker = mocker
        self._orig_env = os.environ.copy()
        self._module_under_test = module_under_test
        self._modules_to_mock = {f"mock_{m.replace('.', '_')}": m for m in modules_to_mock}

    def __enter__(self):
        context = namedtuple('context', ['prefix', 'os', 'mocker'] + list(self._modules_to_mock.keys()))

        context.prefix = self._module_under_test.__name__
        for m, target in self._modules_to_mock.items():
            setattr(context, m, self._mocker.patch(f'{context.prefix}.{target}'))

        context.os = {'environ': os.environ}
        context.mocker = self._mocker
        return context

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._mocker.stopall()
        os.environ = self._orig_env


def _extract_call_elements(mock, call_number=1):
    if mock is None or mock.call_count == 0:
        raise TypeError('Mock must have calls')
    if call_number > len(mock.call_args_list) or call_number < 1:
        raise TypeError(f'Invalid call number, should be between 1 and {len(mock.mock_calls)}, inclusive')
    args, kwargs = mock.call_args_list[call_number - 1]
    return args, kwargs


def extract_call_args(mock, call_number=1):
    args, _ = _extract_call_elements(mock, call_number)
    return args


def extract_call_kwargs(mock, call_number=1):
    _, kwargs = _extract_call_elements(mock, call_number)
    return kwargs


class CallAssertion(NamedTuple):
    args: tuple
    expected_args: tuple
    kwargs: dict
    expected_kwargs: dict

    def __bool__(self):
        return self.expected_args == self.args and self.expected_kwargs == self.kwargs

    def __eq__(self, other):
        return bool(self) == other

    def __repr__(self):
        if self.expected_args != self.args:
            if len(self.args) != len(self.expected_args):
                return f'(ARG COUNT: {len(self.args)} != {len(self.expected_args)})'
            for i, (a, b) in enumerate(zip(self.args, self.expected_args)):
                if a != b:
                    return f'(ARG({i}): {repr(a)} != {repr(b)})'

        if self.expected_kwargs != self.kwargs:
            if self.kwargs.keys() != self.expected_kwargs.keys():
                return f'(KWARGS: {repr(self.kwargs.keys())} != {repr(self.expected_kwargs.keys())})'
            for k in self.expected_kwargs.keys():
                kwarg = self.kwargs.get(k)
                expected_kwarg = self.expected_kwargs.get(k)
                if expected_kwarg != kwarg:
                    return f'(KWARG({k}): {repr(kwarg)} != {repr(expected_kwarg)})'

        return 'CallAssertion:MATCH'


class MockEx:
    def __init__(self, mock):
        self._mock = mock

    def must_call_with(self, *expected_args, **expected_kwargs):
        return self.if_called_with(*expected_args, **expected_kwargs).verify()

    def if_called_with_at_least(self, *expected_args, **expected_kwargs):
        return self.if_called_with(*expected_args, **expected_kwargs).allow_extra()

    def if_called_with_anything(self):
        return self.if_called_with_at_least()

    def if_called_with(self, *expected_args, **expected_kwargs):
        icw = _IfCalledWith(*expected_args, **expected_kwargs)
        icw._next_call = self._mock.side_effect
        self._mock.side_effect = icw
        return icw

    def was_called_once_with(self, *expected_args, **expected_kwargs):
        return self.was_called_with(*expected_args, **expected_kwargs) and self._mock.call_count == 1

    def was_called_with(self, *expected_args, **expected_kwargs):
        return _MockCall(self._mock, 1).was_called_with(*expected_args, **expected_kwargs)

    def call(self, num):
        return _MockCall(self._mock, num)


class _MockCall:
    def __init__(self, mock, call_num):
        self._mock = mock
        self._call_num = call_num

    def was_called_with(self, *expected_args, **expected_kwargs):
        args, kwargs = _extract_call_elements(self._mock, self._call_num)
        return CallAssertion(
            args=args,
            expected_args=expected_args,
            kwargs=kwargs,
            expected_kwargs=expected_kwargs
        )


class _IfCalledWith:
    def __init__(self, *expected_args, **expected_kwargs):
        self._expected_args = expected_args
        self._expected_kwargs = expected_kwargs
        self._return_value = None
        self._use_return_value = False
        self._verify = False
        self._allow_extra = False
        self._next_call = None

    def return_value(self, value):
        self._return_value = value
        self._use_return_value = True
        return self

    def verify(self, val=True):
        self._verify = val
        return self

    def allow_extra(self, val=True):
        self._allow_extra = val
        return self

    def _execute(self, args, kwargs):
        return self._return_value

    def __call__(self, *args, **kwargs):
        expected_args, expected_kwargs = (
            (*self._expected_args, *args[len(self._expected_args):]),
            {**kwargs, **self._expected_kwargs}
        ) if self._allow_extra else (
            self._expected_args,
            self._expected_kwargs
        )
        if self._verify:
            assert expected_args == args and expected_kwargs == kwargs
        if expected_args == args and expected_kwargs == kwargs:
            return self._execute(args, kwargs)
        elif self._next_call is not None:
            return self._next_call(*args, **kwargs)


class MATCH:
    def __init__(self, func):
        self._match_function = func

    def __eq__(self, other):
        return self._match_function(other)

    def __ne__(self, other):
        return not self._match_function(other)


def nested_approx(val, rel=None, abs=None, nan_ok=False):
    return json.loads(json.dumps(val), parse_float=lambda x: pytest.approx(float(x), rel, abs, nan_ok))


def maybe_expect_exception(func, should_raise, expected_exception=Exception):
    if should_raise:
        with pytest.raises(expected_exception) as err:
            func()
        return err
    else:
        return func()
