"""Optional DDP/FSDP construction without hard-coded GPU counts."""

from __future__ import annotations

import os

import torch
from torch import nn


def wrap_distributed(
    model: nn.Module, *, use_fsdp: bool = False, backend: str = "nccl"
) -> nn.Module:
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    if world_size <= 1:
        return model
    if backend == "nccl" and not torch.cuda.is_available():
        backend = "gloo"
    if not torch.distributed.is_initialized():
        torch.distributed.init_process_group(backend=backend)
    if use_fsdp:
        from torch.distributed.fsdp import FullyShardedDataParallel

        return FullyShardedDataParallel(model)
    device_id = torch.cuda.current_device() if torch.cuda.is_available() else None
    return torch.nn.parallel.DistributedDataParallel(
        model, device_ids=[device_id] if device_id is not None else None
    )