#   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 subprocess
import os
from parl.utils import SerializeError, DeserializeError

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

try:
    import pyarrow
    pyarrow_installed = True
except ImportError:
    pyarrow_installed = False

if pyarrow_installed:
    # 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)

    # if pyarrow is installed, parl will use pyarrow to serialize/deserialize objects.
    serialize = lambda data: pyarrow.serialize(data, context=context).to_buffer()
    deserialize = lambda data: pyarrow.deserialize(data, context=context)
else:
    # if pyarrow is not installed, parl will use cloudpickle to serialize/deserialize objects.
    serialize = lambda data: cloudpickle.dumps(data)
    deserialize = lambda data: cloudpickle.loads(data)


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 = serialize([args, kwargs])
    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 = deserialize(data)
    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 = serialize(data)
    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 = deserialize(data)
    except Exception as e:
        raise DeserializeError(e)

    return ret
