# Copyright (c) 2021 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 dataclasses
import typing
from dataclasses import is_dataclass
from datetime import datetime, date
from functools import wraps
from inspect import isclass
from typing import Any, Callable, Type, Union

from dateutil.parser import isoparse
from pydantic import BaseModel, parse_obj_as
from voluptuous import (PREVENT_EXTRA, REMOVE_EXTRA, All, Coerce, Maybe,
                        Optional, Required, Schema)

from czc import json

T = typing.TypeVar('T')


def deserializable(schema: typing.Optional[Schema] = None) -> Callable[[Type[T]], Type[T]]:
    def _deserializable(cls: Type[T]) -> Type[T]:
        if not _supports_deserializable(cls):
            raise TypeError(f'@deserializable is only supported on typed NamedTuple, @dataclass, '
                            f'or a Pydantic BaseModel, not {cls}')
        if hasattr(cls, '__cz_has_schema__'):
            raise TypeError(f'@deserializable cannot be applied multiple times to {cls}')

        orig_new = cls.__new__
        orig_init = cls.__init__
        resolved_schema = schema or _create_schema_from_type(cls)

        @wraps(orig_new)
        def new_wrapper(inner_cls, *args, **kwargs):
            if len(args) == 1 and len(kwargs) == 0 and isinstance(args[0], dict):
                return inner_cls(**resolved_schema(args[0]))
            return orig_new(inner_cls, *args, **kwargs)

        @wraps(orig_init)
        def init_wrapper(inner_self, *args, **kwargs):
            if orig_init == object.__init__:
                return orig_init(inner_self)
            if len(args) == 1 and len(kwargs) == 0 and isinstance(args[0], dict):
                return orig_init(inner_self, **resolved_schema(args[0]))
            return orig_init(inner_self, *args, **kwargs)

        setattr(cls, '__cz_has_schema__', True)

        if '__new__' in cls.__dict__:
            setattr(cls, '__new__', new_wrapper)
        setattr(cls, '__init__', init_wrapper)

        return cls

    return _deserializable


def coerce(to_type: Type[T], value: Any) -> T:
    return Coerce(to_type)(value)


def _supports_deserializable(cls):
    return (_issubclass(cls, tuple) or _issubclass(cls, BaseModel) or is_dataclass(cls)) and typing.get_type_hints(cls)


def _is_optional(origin, type_args):
    return origin is Union and len(type_args) == 2 and isinstance(None, type_args[1])


def _schema_key(key_name: str, to_type: Any):
    def has_tuple_default():
        return _issubclass(to_type, tuple) and hasattr(to_type, '_field_defaults') and key_name in to_type._field_defaults

    def has_dataclass_default():
        if not is_dataclass(to_type):
            return False
        field = next((x for x in dataclasses.fields(to_type) if x.name == key_name), None)
        if not field:
            return False
        return field.default != dataclasses.MISSING or field.default_factory != dataclasses.MISSING

    has_default = has_tuple_default() or has_dataclass_default()
    return Optional(key_name) if has_default else Required(key_name)


def _issubclass(cls, from_type):
    return isclass(cls) and issubclass(cls, from_type)


def _create_voluptuous_schema_from_type(to_type: Type[T], prevent_extra=False) -> Callable[[Any], dict]:
    def create_schema(t, rt=None):
        origin = typing.get_origin(t)
        if origin:
            type_args = typing.get_args(t)
            if _is_optional(origin, type_args):
                return Maybe(create_schema(type_args[0]))
            elif _issubclass(origin, list):
                return [create_schema(type_args[0])]
            elif _issubclass(origin, dict):
                return {create_schema(type_args[0]): create_schema(type_args[1])}
            elif _issubclass(origin, set):
                return All([create_schema(type_args[0])], Coerce(set))
        elif _supports_deserializable(t) and not getattr(t, '__cz_has_schema__', False):
            hints = typing.get_type_hints(t)
            if hints:
                return All(Schema({_schema_key(k, t): create_schema(v) for k, v in hints.items()},
                                  extra=PREVENT_EXTRA if prevent_extra else REMOVE_EXTRA),
                           lambda x: (rt or t)(**x))
        elif _issubclass(t, datetime):
            return isoparse
        elif _issubclass(t, date):
            return lambda x: isoparse(x).date()

        return Coerce(t)

    return create_schema(to_type, dict)


def _create_schema_from_type(to_type: Type[T]) -> Callable[[Any], dict]:
    if _issubclass(to_type, BaseModel):
        return lambda x: parse_obj_as(typing.cast(Type[BaseModel], to_type), x).dict(by_alias=True)
    else:
        return _create_voluptuous_schema_from_type(to_type)


def serializable(blob: Any) -> Any:
    return json.loads(json.dumps(blob))


def require(value: typing.Optional[T], error_message: str = '') -> T:
    if value is None:
        raise ValueError(error_message)
    return typing.cast(T, value)
