# 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 dataclasses
from datetime import date, datetime
from decimal import Decimal
from enum import Enum
from functools import singledispatch
from typing import Any
from uuid import UUID

import simplejson as json
from pydantic import BaseModel


class ExtendedEncoder(json.JSONEncoder):
    def __init__(self, use_serialize_method=False, **kwargs):
        self.use_serialize_method = use_serialize_method
        super().__init__(**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):
            return dataclasses.asdict(o)
        if self.use_serialize_method and hasattr(o, 'serializable'):
            return o.serializable()
        return super().default(o)


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


@dumps.register
def _(obj: BaseModel, **kwargs):
    return json.dumps(
        obj.model_dump(
            by_alias=True,
            **kwargs,
        ),
        cls=ExtendedEncoder,
        use_decimal=False,
        iterable_as_array=True,
        use_serialize_method=True,
    )


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


def serializable(blob: Any) -> Any:
    return loads(dumps(blob, use_serialize_method=True))
