# 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 collections.abc import Callable
from dataclasses import is_dataclass
from datetime import date, datetime
from functools import wraps
from inspect import isclass
from typing import Any, Union

from dateutil.parser import isoparse
from pydantic import BaseModel
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, 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]))
            if _issubclass(origin, list):
                return [create_schema(type_args[0])]
            if _issubclass(origin, dict):
                return {create_schema(type_args[0]): create_schema(type_args[1])}
            if _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: to_type.model_validate(x).model_dump(by_alias=True)  # type: ignore
    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)
