"""DDP-safe utilities for distributed training."""

import os
from typing import Optional

# Lazy import torch to avoid requiring it for server-only usage
_torch = None


def _get_torch():
    """Lazy load torch."""
    global _torch
    if _torch is None:
        import torch
        _torch = torch
    return _torch


def is_distributed() -> bool:
    """Check if running in distributed mode."""
    torch = _get_torch()
    return torch.distributed.is_available() and torch.distributed.is_initialized()


def get_rank() -> int:
    """Get current process rank. Returns 0 if not distributed."""
    if not is_distributed():
        return 0
    torch = _get_torch()
    return torch.distributed.get_rank()


def get_world_size() -> int:
    """Get world size. Returns 1 if not distributed."""
    if not is_distributed():
        return 1
    torch = _get_torch()
    return torch.distributed.get_world_size()


def is_rank0() -> bool:
    """Check if this is rank 0 (main process)."""
    return get_rank() == 0


def safe_barrier():
    """
    Execute a barrier if in distributed mode.
    Safe to call even if not in distributed mode (no-op).
    """
    if is_distributed():
        torch = _get_torch()
        torch.distributed.barrier()


class DDPContext:
    """
    Context manager for DDP-safe evaluation.
    
    Usage:
        with DDPContext() as ctx:
            if ctx.is_main:
                # Generate outputs
                # Send to server
    
    Barriers are automatically called on enter and exit.
    """
    
    def __init__(self):
        self.is_main = False
    
    def __enter__(self):
        safe_barrier()
        self.is_main = is_rank0()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        safe_barrier()
        return False
