import orjson

from docarray.typing.tensor.abstract_tensor import AbstractTensor


def _default_orjson(obj):
    """
    default option for orjson dumps.
    :param obj:
    :return: return a json compatible object
    """

    if isinstance(obj, AbstractTensor):
        return obj._docarray_to_json_compatible()
    else:
        return obj


def orjson_dumps(v, *, default=None) -> bytes:
    # dumps to bytes using orjson
    return orjson.dumps(v, default=_default_orjson, option=orjson.OPT_SERIALIZE_NUMPY)


def orjson_dumps_and_decode(v, *, default=None) -> str:
    # dumps to bytes using orjson
    return orjson_dumps(v, default=default).decode()

import orjson


def _default_orjson(obj):
    """
    default option for orjson dumps. It will call _to_json_compatible
    from docarray typing object that expose such method.
    :param obj:
    :return: return a json compatible object
    """

    if getattr(obj, '_to_json_compatible'):
        return obj._to_json_compatible()
    else:
        return obj


def orjson_dumps(v, *, default=None):
    # orjson.dumps returns bytes, to match standard json.dumps we need to decode
    return orjson.dumps(
        v, default=_default_orjson, option=orjson.OPT_SERIALIZE_NUMPY
    ).decode()
