# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import functools
import os
from typing import Optional, Callable

import torch
import torch.distributed as dist

import lipforcing.utils.logging_utils as logger


def world_size():
    """Get the world size."""
    if dist.is_initialized() and torch.cuda.is_available():
        return dist.get_world_size()
    return 1


def get_rank(group: Optional[dist.ProcessGroup] = None) -> int:
    """Get the rank (GPU device) of the worker.

    Returns:
        rank (int): The rank of the worker.
    """
    rank = 0
    if dist.is_available() and dist.is_initialized():
        rank = dist.get_rank(group)
    return rank


def is_rank0() -> bool:
    """Return True if this is rank 0 (the primary loading rank)."""
    return get_rank() == 0


def synchronize():
    """
    Synchronize all devices.

    This method checks if the current running environment
    is distributed with a world-size greater than 1.
    If so, we use `dist.barrier` to synchronize
    all processes.
    """
    if not dist.is_available():
        return
    if not dist.is_initialized():
        return

    world_size = dist.get_world_size()
    if world_size == 1:
        return
    logger.debug(f"Synchronizing all devices with world size {world_size}")
    dist.barrier(device_ids=[int(os.environ.get("LOCAL_RANK", "0"))])
    logger.debug(f"Synchronized all devices with world size {world_size}")


def rank0_only(func: Callable) -> Callable:
    """Apply this function only to the master GPU.

    Example usage:
        @rank0_only
        def func(x):
            return x + 1

    Args:
        func (Callable): any function.

    Returns:
        (Callable): A function wrapper executing the function only on the master GPU.
    """

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if is_rank0():
            return func(*args, **kwargs)
        else:
            return None

    return wrapper


def clean_up():
    if dist.is_available() and dist.is_initialized():
        try:
            logger.info("Distributed clean up.")
            dist.destroy_process_group()
        except ValueError as e:
            logger.error(f"Error destroying default process group: {e}")


def sync_any(local_any: bool, device: torch.device) -> bool:
    """Synchronize local any across distributed ranks.

    Args:
        local_any: any() in each rank
        device: Device for tensor operations

    Returns:
        global_any
    """
    global_any = torch.tensor([local_any], dtype=torch.uint8, device=device)

    if world_size() > 1:
        # MAX reduction: global_any is True if any rank has any samples in second stage
        torch.distributed.all_reduce(global_any, op=torch.distributed.ReduceOp.MAX)

    return global_any.to(torch.bool).item()