#  Copyright (C) 2021
#  ABM, Moscow
#
#  UNPUBLISHED PROPRIETARY MATERIAL.
#  ALL RIGHTS RESERVED.
#
#  Authors: Mike Orlov <m.orlov@abm-jsc.ru>
#
import base64
import datetime
import json
from dataclasses import is_dataclass
from decimal import Decimal
from functools import partial
from typing import Union, TypeAlias, Callable, Protocol, Any

from .dataclass_protocol import DataclassProtocol
from .dict_to_dataclass import get_dataclass_field_name_to_field


JsonLoaded: TypeAlias = Union[list['JsonLoaded'], dict[str, 'JsonLoaded'], str, int, float, None]
Jsonable: TypeAlias = Union[JsonLoaded, list['Jsonable'], tuple['Jsonable', ...], dict[str | int, 'Jsonable']]


def read_json_file_by_path(path: str) -> JsonLoaded:
    with open(path) as file:
        return json.load(file)


class ReprInDumps:
    def __repr_in_dumps__(self):
        return repr(self)


CustomJsonable: TypeAlias = Union[
    Jsonable, list['CustomJsonable'], tuple['CustomJsonable', ...], dict[str | int, 'CustomJsonable'],
    set['CustomJsonable'], frozenset['CustomJsonable'],
    bytes, ReprInDumps, type, DataclassProtocol, datetime.datetime, datetime.date, datetime.time, Decimal, Exception
]


class EnhancedJSONEncoder(json.JSONEncoder):
    def default(self, o):
        if isinstance(o, ReprInDumps):
            return o.__repr_in_dumps__()
        if isinstance(o, type):
            return {'_t': 'PY::class', 'key': o.__qualname__}
        if is_dataclass(o):
            return {
                key: getattr(o, key)
                for key, field in get_dataclass_field_name_to_field(type(o), with_init_vars=False).items()
                if field.repr
            }
        if isinstance(o, datetime.datetime):
            if o.tzinfo is None or o.tzinfo.utcoffset(o) is None:  # PROHIBIT naive datetime serialisation
                raise TypeError("TypeError: datetime.datetime WITHOUT tzinfo is not JSON serializable")
        if isinstance(o, datetime.date | datetime.time):  # date or time or datetime
            return o.isoformat()
        if isinstance(o, set | frozenset):
            return tuple(o)
        if isinstance(o, bytes):
            return f'data:application/octet-stream;base64,{base64.b64encode(o)}'
        if isinstance(o, Decimal):
            return str(o)
        if isinstance(o, Exception):
            return {'_t': 'PY::Exception', 'key': type(o).__qualname__, 'args': o.args}
        return super().default(o)


class JsonDumper(Protocol):
    @staticmethod
    def __call__(
            obj: CustomJsonable, *, skipkeys: bool = False, ensure_ascii: bool = True, check_circular: bool = True,
            allow_nan: bool = True, indent: int | None = None, separators: tuple[str, str] | None = None,
            default: Callable[[Any], Jsonable] = None, sort_keys: bool = False, **kw
    ) -> str:
        ...


custom_dumps: JsonDumper = partial(json.dumps, cls=EnhancedJSONEncoder)
