import os
import traceback
import torch
import torch.distributed as dist
import torch.multiprocessing as mp


class DistributedConfig:
    def __init__(
        self,
        gpus=None,
        backend: str = "auto",
        master_addr: str = "localhost",
        master_port: str = "29500",
        start_method: str = "auto",
    ):
        detected = torch.cuda.device_count()
        self.gpus = gpus if gpus is not None else (detected if detected > 0 else 1)
        self.backend = backend
        self.master_addr = master_addr
        self.master_port = master_port
        self.start_method = start_method

    def resolved_backend(self) -> str:
        if self.backend != "auto":
            return self.backend
        enough_cuda = torch.cuda.is_available() and torch.cuda.device_count() >= self.gpus
        return "nccl" if enough_cuda else "gloo"

    def resolved_start_method(self) -> str:
        if self.start_method != "auto":
            return self.start_method
        # Default to spawn when CUDA is present to avoid forked CUDA init issues.
        return "spawn" if torch.cuda.is_available() else "fork"


def _setup_env(rank, world_size, cfg: DistributedConfig, backend: str):
    os.environ["MASTER_ADDR"] = cfg.master_addr
    os.environ["MASTER_PORT"] = cfg.master_port
    os.environ["RANK"] = str(rank)
    os.environ["WORLD_SIZE"] = str(world_size)
    os.environ["LOCAL_RANK"] = str(rank)

    if backend == "nccl" and torch.cuda.is_available():
        if rank >= torch.cuda.device_count():
            raise RuntimeError(f"Rank {rank} requests GPU but only {torch.cuda.device_count()} devices are visible.")
        torch.cuda.set_device(rank)

    dist.init_process_group(
        backend=backend,
        rank=rank,
        world_size=world_size,
    )


def _worker(rank, world_size, cfg, backend, train_fn, args, kwargs, error_queue):
    _setup_env(rank, world_size, cfg, backend)
    try:
        train_fn(rank, world_size, *args, **kwargs)
    except Exception:
        if error_queue is not None:
            error_queue.put(traceback.format_exc())
        raise
    finally:
        if dist.is_initialized():
            dist.destroy_process_group()


def launch(train_fn, cfg: DistributedConfig, *args, **kwargs):
    """
    train_fn(rank: int, world_size: int, *args, **kwargs)
    """
    backend = cfg.resolved_backend()
    start_method = cfg.resolved_start_method()

    if cfg.gpus < 1:
        raise ValueError("world_size/gpus must be at least 1")

    world_size = cfg.gpus
    ctx = mp.get_context(start_method)
    error_queue = ctx.SimpleQueue()

    mp.start_processes(
        _worker,
        args=(world_size, cfg, backend, train_fn, args, kwargs, error_queue),
        nprocs=world_size,
        join=True,
        start_method=start_method,
    )

    if not error_queue.empty():
        raise RuntimeError(f"Worker failure:\n{error_queue.get()}")
