"""Distributed process setup for data-parallel and FSDP training."""

import os
from datetime import timedelta

import torch
import torch.distributed as dist


def print_rank0(message: str) -> None:
    if int(os.getenv("RANK", "0")) == 0:
        print(message)


def init_distributed_environment(backend: str = "nccl") -> None:
    if dist.is_initialized():
        return
    global_rank = int(os.getenv("RANK", "0"))
    world_size = int(os.getenv("WORLD_SIZE", "1"))
    print_rank0(f"Initializing distributed environment with world_size={world_size}")
    dist.init_process_group(
        backend=backend, world_size=world_size, rank=global_rank,
        timeout=timedelta(minutes=30),
    )


def setup_distributed_training() -> tuple[int, int, int, torch.device]:
    local_rank = int(os.getenv("LOCAL_RANK", "0"))
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    init_distributed_environment()
    return local_rank, dist.get_rank(), dist.get_world_size(), device


def clean_dist_env() -> None:
    if dist.is_initialized():
        dist.destroy_process_group()