from typing import Callable, Any, List, Union
from ding.data.buffer import BufferedData
from ding.utils import fastcopy


def clone_object():
    """
    Overview:
        This middleware freezes the objects saved in memory buffer and return copies during sampling,
        try this middleware when you need to keep the object unchanged in buffer, and modify\
        the object after sampling it (usually in multiple threads)
    """

    def push(chain: Callable, data: Any, *args, **kwargs) -> BufferedData:
        data = fastcopy.copy(data)
        return chain(data, *args, **kwargs)

    def sample(chain: Callable, *args, **kwargs) -> Union[List[BufferedData], List[List[BufferedData]]]:
        data = chain(*args, **kwargs)
        return fastcopy.copy(data)

    def _clone_object(action: str, chain: Callable, *args, **kwargs):
        if action == "push":
            return push(chain, *args, **kwargs)
        elif action == "sample":
            return sample(chain, *args, **kwargs)
        return chain(*args, **kwargs)

    return _clone_object