# 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 collections.abc
import dataclasses
import inspect
from datetime import date, datetime
from decimal import Decimal
from enum import Enum
from typing import Any
from uuid import UUID

import simplejson as json
from pydantic import BaseModel


def _get_allowed_kwargs(func: collections.abc.Callable) -> set[str]:
    allowed_kwargs = set()
    argspec = inspect.getfullargspec(func)
    allowed_kwargs.update(argspec.args)
    allowed_kwargs.update(argspec.kwonlyargs)
    return allowed_kwargs


allowed_json_encoder_kwargs = _get_allowed_kwargs(json.JSONEncoder.__init__)


class ExtendedEncoder(json.JSONEncoder):
    def __init__(self, use_serialize_method=False, **kwargs):
        self.use_serialize_method = use_serialize_method
        self._pydantic_kwargs = {
            'by_alias': True,
            **{key: value for key, value in kwargs.items() if key not in allowed_json_encoder_kwargs},
        }
        super_kwargs = {key: value for key, value in kwargs.items() if key in allowed_json_encoder_kwargs}

        super().__init__(**super_kwargs)

    def default(self, o):
        if isinstance(o, Decimal):
            return str(o)
        if isinstance(o, Enum):
            return o.value
        if isinstance(o, datetime) or isinstance(o, date):
            return o.isoformat()
        if isinstance(o, UUID):
            return str(o)
        if dataclasses.is_dataclass(o) and not isinstance(o, type):
            return dataclasses.asdict(o)
        if self.use_serialize_method and hasattr(o, 'serializable'):
            return o.serializable()
        if isinstance(o, BaseModel):
            return o.model_dump(**self._pydantic_kwargs)
        if isinstance(o, collections.abc.Iterable):
            return list(o)

        return super().default(o)


def dumps(obj, **kwargs):
    return json.dumps(
        obj,
        **{**dict(cls=ExtendedEncoder, use_decimal=False, use_serialize_method=True), **kwargs},
    )


def loads(s, **kwargs):
    return json.loads(s, **kwargs)


def load(fp, **kwargs):
    return json.load(fp, **kwargs)


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