# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import os
import subprocess
import sys
from functools import wraps
from types import FunctionType
from typing import Any, Callable, Literal, Optional

import pytest
import torch
import torch.distributed as dist
from compressed_tensors.offload.utils import send_tensors


def assert_device_equal(
    device_a: torch.device | Literal["disk"],
    device_b: torch.device | Literal["disk"],
):
    if device_a == "disk":
        device_a = torch.device("meta")
    if device_b == "disk":
        device_b = torch.device("meta")

    cur_index = torch.cuda.current_device()
    a_index = cur_index if device_a.index is None else device_a.index
    b_index = cur_index if device_b.index is None else device_b.index
    assert device_a.type == device_b.type and a_index == b_index


def assert_tensor_equal(
    tensor_a: torch.Tensor,
    tensor_b: torch.Tensor,
    device: Optional[torch.device | str] = None,
):
    if device is not None:
        tensor_b = send_tensors(tensor_b, "meta" if device == "disk" else device)

    assert tensor_a.__class__ == tensor_b.__class__
    assert tensor_a.requires_grad == tensor_b.requires_grad
    assert tensor_a.__dict__ == tensor_b.__dict__

    if tensor_a.is_meta or tensor_b.is_meta:
        assert (
            tensor_a.device == tensor_b.device
            and tensor_a.shape == tensor_b.shape
            and tensor_a.dtype == tensor_b.dtype
        )
    else:
        assert torch.equal(tensor_a, tensor_b)


def torchrun(world_size: int = 1) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """
    Test a function within parallel `torchrun` subprocesses, each running with `pytest`
    ```
    # (main) -> (torchrun) -
    #              \\-- (rank 0)
    #              \\-- (rank 1)
    ```

    :param world_size: number of ranks to spawn
    """

    def decorator(func: FunctionType):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # We're running in a torchrun subprocess:
            # init distributed and run test func
            if "TORCHELASTIC_RUN_ID" in os.environ:
                rank = int(os.environ["RANK"])
                local_rank = int(os.environ["LOCAL_RANK"])

                torch.cuda.set_device(local_rank)
                dist.init_process_group(
                    backend="nccl",
                    init_method="env://",
                    rank=rank,
                    world_size=world_size,
                    device_id=local_rank,
                )
                dist.barrier()

                return func(*args, **kwargs)

            # First time calling in the main process:
            # trigger torchrun with this function as the pytest target
            else:
                file_path = sys.modules.get(func.__module__).__file__
                func_name = func.__name__

                cmd = (
                    f"{sys.executable} "
                    f"-m torch.distributed.run --nproc_per_node {world_size} "
                    "--log-dir /tmp/torchrun-logs --tee 3 --role torchrun "
                    f"-m pytest {file_path}::{func_name} -sx"
                )

                proc = subprocess.run(cmd.split(" "))
                assert proc.returncode == 0

        return wrapper

    return decorator


@pytest.fixture()
def cuda_device():
    return (
        torch.device("cuda")
        if "TORCHELASTIC_RUN_ID" in os.environ
        else torch.device("cuda:0")
    )
