#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import cloudpickle
import pyarrow
from parl.utils import SerializeError, DeserializeError

__all__ = ['dumps_argument', 'loads_argument', 'dumps_return', 'loads_return']


# Reference: https://github.com/apache/arrow/blob/f88474c84e7f02e226eb4cc32afef5e2bbc6e5b4/python/pyarrow/tests/test_serialization.py#L658-L682
def _serialize_serializable(obj):
    return {"type": type(obj), "data": obj.__dict__}


def _deserialize_serializable(obj):
    val = obj["type"].__new__(obj["type"])
    val.__dict__.update(obj["data"])
    return val


context = pyarrow.default_serialization_context()

# support deserialize in another environment
context.set_pickle(cloudpickle.dumps, cloudpickle.loads)

# support serialize and deserialize custom class
context.register_type(
    object,
    "object",
    custom_serializer=_serialize_serializable,
    custom_deserializer=_deserialize_serializable)


def dumps_argument(*args, **kwargs):
    """

    Serialize arguments passed to a function.

    args: 
        *args, **kwargs are general a commonly used representation of arguments in python.

    Returns:
        Implementation-dependent object in bytes.
    """
    try:
        ret = pyarrow.serialize([args, kwargs], context=context).to_buffer()
    except Exception as e:
        raise SerializeError(e)

    return ret


def loads_argument(data):
    """
    Restore bytes data to their initial data formats.

    Args:
        data: the output of `dumps_argument`.

    Returns:
        deserialized arguments [args, kwargs]
        like the input of `dumps_argument`, args is a tuple, and kwargs is a dict 
    """
    try:
        ret = pyarrow.deserialize(data, context=context)
    except Exception as e:
        raise DeserializeError(e)

    return ret


def dumps_return(data):
    """
    Serialize the return data of a function.

    Args:
        data: the output of a function.

    Returns:
        Implementation-dependent object in bytes.
    """
    try:
        ret = pyarrow.serialize(data, context=context).to_buffer()
    except Exception as e:
        raise SerializeError(e)

    return ret


def loads_return(data):
    """
    Deserialize the data generated by `dumps_return`.

    Args:
        data: the output of `dumps_return`

    Returns:
        deserialized data
    """
    try:
        ret = pyarrow.deserialize(data, context=context)
    except Exception as e:
        raise DeserializeError(e)

    return ret
