"""TP group state (ADR-0027 D3).

Single global TP group. Initial scope: TP size == world_size (pure TP;
mixed DP+TP is future work).
"""
from __future__ import annotations

_TP_WORLD_SIZE: int | None = None


def initialize_model_parallel(tensor_model_parallel_size: int) -> None:
    """Initialize the TP process group.

    Must be called after ``torch.distributed.init_process_group``.
    Only ``tensor_model_parallel_size == world_size`` is supported in the
    initial scope.
    """
    global _TP_WORLD_SIZE
    # Import here to avoid cycle when tp is imported before a ctx exists.
    _ws = _current_world_size()
    if tensor_model_parallel_size != _ws:
        raise NotImplementedError(
            f"Only TP == world_size supported; got TP={tensor_model_parallel_size}, "
            f"world_size={_ws}"
        )
    _TP_WORLD_SIZE = tensor_model_parallel_size


def get_tensor_model_parallel_world_size() -> int:
    """Return the TP group's world size.

    Raises if not initialised — callers must call
    :func:`initialize_model_parallel` first.
    """
    if _TP_WORLD_SIZE is None:
        raise RuntimeError(
            "TP group not initialised; call initialize_model_parallel() first"
        )
    return _TP_WORLD_SIZE


def get_tensor_model_parallel_rank() -> int:
    """Return this worker's rank within the TP group.

    Delegates to the greenlet-local rank registered by the spawn launcher
    (ADR-0024 D2 via ``torch.distributed.get_rank``).
    """
    # Resolve via the global torch.distributed facade on the active ctx.
    return _current_rank()


def _reset_for_tests() -> None:
    """Clear _TP_WORLD_SIZE so ordering-sensitive tests can re-init."""
    global _TP_WORLD_SIZE
    _TP_WORLD_SIZE = None


# ── helpers (resolve current ctx) ────────────────────────────────────


def _current_ctx():
    """Best-effort resolution of the currently-active RuntimeContext.

    In KernBench, the ``ctx`` is passed as the ``torch`` positional in
    bench/worker code. Since parallel_state is a module-global helper,
    we look it up via a weak registry maintained by RuntimeContext.
    """
    from kernbench.runtime_api.context import _get_active_context
    ctx = _get_active_context()
    if ctx is None:
        raise RuntimeError(
            "No active RuntimeContext; kernbench.tp requires one "
            "(call init_process_group / spawn under a live ctx)"
        )
    return ctx


def _current_world_size() -> int:
    return _current_ctx().distributed.get_world_size()


def _current_rank() -> int:
    return _current_ctx().distributed.get_rank()
