# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import torch
from deepspeed.utils.torch import required_torch_version

backward_inputs = []

enabled_patched_func = False
original_grad_fn = None
base_meta = type(torch.autograd.Function)

if required_torch_version(min_version=2.7):

    class FunctionMeta(base_meta):

        def __new__(cls, name, bases, dct):
            if name == "CompiledFunction":
                original_backward_impl = dct.get("_backward_impl")

                def wrapped_backward_impl(ctx, all_args):
                    assert original_backward_impl is not None

                    if enabled_patched_func:
                        backward_inputs.append(all_args)
                        wrapped_backward_impl.owner_class.compiled_bw = None

                    return original_backward_impl(ctx, all_args)

                wrapped_backward_impl.owner_class = None
                dct["_backward_impl"] = staticmethod(wrapped_backward_impl)
                new_class = super().__new__(cls, name, bases, dct)
                wrapped_backward_impl.owner_class = new_class

                return new_class

            return super().__new__(cls, name, bases, dct)

elif required_torch_version(min_version=2.6):

    class FunctionMeta(base_meta):

        def __new__(cls, name, bases, dct):
            if name == "CompiledFunction":
                original_backward_prologue = dct.get("_backward_prologue")

                def wrapped_backward_prologue(ctx, *grad_outputs):
                    assert original_backward_prologue is not None

                    all_args = original_backward_prologue(ctx, *grad_outputs)
                    if enabled_patched_func:
                        backward_inputs.append(all_args)
                        wrapped_backward_prologue.owner_class.compiled_bw = None

                    return all_args

                wrapped_backward_prologue.owner_class = None
                dct["_backward_prologue"] = staticmethod(wrapped_backward_prologue)
                new_class = super().__new__(cls, name, bases, dct)
                wrapped_backward_prologue.owner_class = new_class

                return new_class

            return super().__new__(cls, name, bases, dct)


def patch_compiled_func():

    global enabled_patched_func
    enabled_patched_func = True

    class PatchedFunction(torch.autograd.Function, metaclass=FunctionMeta):
        pass

    global original_grad_fn
    original_grad_fn = torch.autograd.Function
    torch.autograd.Function = PatchedFunction

    return backward_inputs


def unpatch_compiled_func():
    global enabled_patched_func
    enabled_patched_func = False

    global original_grad_fn
    torch.autograd.Function = original_grad_fn


def get_backward_inputs():
    return backward_inputs