"""Process-group context and rank-local device selection."""

from __future__ import annotations

import os
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import timedelta
from typing import Iterator

import torch
import torch.distributed as dist


@dataclass(frozen=True)
class DistributedContext:
    enabled: bool = False
    rank: int = 0
    local_rank: int = 0
    world_size: int = 1
    device: torch.device = torch.device("cpu")

    @property
    def is_main_process(self) -> bool:
        return self.rank == 0

    @classmethod
    def single_process(cls, device: torch.device | str = "cpu") -> "DistributedContext":
        return cls(device=torch.device(device))


def _launched() -> bool:
    return "RANK" in os.environ or "WORLD_SIZE" in os.environ


def initialize(
    *,
    enabled: bool | None = None,
    backend: str | None = None,
    device: torch.device | str | None = None,
    timeout_seconds: int = 1800,
) -> DistributedContext:
    """Initialize ``torchrun``'s ``env://`` process group when requested."""

    launched = _launched()
    if enabled is False and launched and int(os.environ.get("WORLD_SIZE", "1")) > 1:
        raise ValueError(
            "distributed launcher variables are present but distributed.enabled=false"
        )
    active = launched if enabled is None else enabled
    if not active:
        requested = torch.device(
            device or ("cuda" if torch.cuda.is_available() else "cpu")
        )
        if requested.type == "cuda" and not torch.cuda.is_available():
            requested = torch.device("cpu")
        return DistributedContext.single_process(requested)

    try:
        rank = int(os.environ["RANK"])
        local_rank = int(os.environ.get("LOCAL_RANK", rank))
        world_size = int(os.environ["WORLD_SIZE"])
    except KeyError as error:
        raise ValueError(
            "distributed execution requires torchrun RANK and WORLD_SIZE"
        ) from error
    if world_size < 2:
        return DistributedContext.single_process(device or "cpu")
    requested = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
    if requested.type == "cuda":
        if not torch.cuda.is_available():
            raise RuntimeError(
                "CUDA was requested for distributed execution but is unavailable"
            )
        torch.cuda.set_device(local_rank)
        assigned = torch.device("cuda", local_rank)
        selected_backend = backend or "nccl"
    else:
        assigned = torch.device("cpu")
        selected_backend = backend or "gloo"
    if not dist.is_initialized():
        dist.init_process_group(
            backend=selected_backend,
            init_method="env://",
            timeout=timedelta(seconds=timeout_seconds),
        )
    return DistributedContext(True, rank, local_rank, world_size, assigned)


def finalize(context: DistributedContext) -> None:
    """Destroy a process group created for ``context``."""

    if context.enabled and dist.is_available() and dist.is_initialized():
        dist.destroy_process_group()


@contextmanager
def process_group(**kwargs: object) -> Iterator[DistributedContext]:
    context = initialize(**kwargs)
    try:
        yield context
    finally:
        finalize(context)