# -*- coding: UTF-8 -*-

import dataclasses
import typing
from collections import UserString

import orjson
from starlette.responses import Response, StreamingResponse

type BaseTypeAnnotation = str | int | float | bool | None


@dataclasses.dataclass(slots=True, kw_only=True)
class _CustomDataType:
    type_class: type
    encoder: typing.Callable[[typing.Any], BaseTypeAnnotation]


_CUSTOM_DATA_TYPE_LIST = [_CustomDataType(type_class=UserString, encoder=str)]


def register_custom_data_types(
        *, type_class: type,
        encoder: typing.Callable[[typing.Any], BaseTypeAnnotation],
) -> None:
    for custom_data_type in _CUSTOM_DATA_TYPE_LIST:
        if custom_data_type.type_class == type_class:
            raise ValueError(f'Type class {type_class} already registered')
    _CUSTOM_DATA_TYPE_LIST.append(
        _CustomDataType(type_class=type_class, encoder=encoder))


def orjson_dumps_default(obj):
    for custom_data_type in _CUSTOM_DATA_TYPE_LIST:
        if isinstance(obj, custom_data_type.type_class):
            return custom_data_type.encoder(obj)
    raise TypeError(f'Object of type {type(obj)} is not JSON serializable')


class ORJSONResponse(Response):
    media_type = 'application/json'

    def render(self, content: typing.Any) -> bytes:
        return orjson.dumps(
            content, option=orjson.OPT_STRICT_INTEGER,
            default=orjson_dumps_default)


class StreamResponse(StreamingResponse):
    pass
