"""PyTorch-compatible distributed communication shim (ADR-0023 D11).

Provides a ``torch.distributed``-like API whose public surface matches
real PyTorch so that bench code looks identical to a DDP training script.

Only the ``ahbm`` backend is implemented. It:

1. Reads ``ccl.yaml`` to decide which collective algorithm to run.
2. Derives world_size from the algorithm entry, the defaults section, or
   from the topology spec (``system.sips.count × sip.cube_mesh × pe_layout``).
3. At ``init_process_group`` time, eagerly installs the IPCQ neighbor
   table once (one-time comm setup — mirrors NCCL communicator creation).
4. On each ``all_reduce(tensor)`` call, reads per-shard metadata from the
   tensor handle and dispatches ``torch.launch`` with the registered
   kernel. The kernel performs intra-PE ring/tree/mesh CCL via IPCQ,
   and Phase 2 DataExecutor replays math + copies from op_log so
   MemoryStore is correct when ``all_reduce`` returns.

Host bench code uses only real-PyTorch names:
    dist.init_process_group, dist.is_initialized, dist.get_world_size,
    dist.get_rank, dist.get_backend, dist.all_reduce, dist.barrier
"""
from __future__ import annotations

import importlib
import math
from typing import Any


class AhbmCCLBackend:
    """Ahbm CCL backend — drives kernel-level collectives via IPCQ."""

    def __init__(self, torch_ctx: Any) -> None:
        from kernbench.ccl.install import (
            load_ccl_config,
            resolve_algorithm_config,
        )

        self.ctx = torch_ctx
        self._cfg_all = load_ccl_config()
        self._merged = resolve_algorithm_config(self._cfg_all)
        self._algo_module = importlib.import_module(self._merged["module"])
        self._world_size = self._resolve_world_size()
        self._pending_collective_handles: list = []
        self._dist_ctx: Any = None

        spec = self.ctx.spec or {}
        self._n_sips = int(spec.get("system", {}).get("sips", {}).get("count", 1))
        self._sip_topo = str(
            spec.get("system", {}).get("sips", {}).get("topology", "ring_1d")
        )
        cm = spec.get("sip", {}).get("cube_mesh", {})
        self._cube_w = int(cm.get("w", 4))
        self._cube_h = int(cm.get("h", 4))

        # Resolve SIP topology dims for the kernel
        topo_map = getattr(self._algo_module, "TOPO_NAME_TO_KIND", None)
        if topo_map is not None:
            self._sip_topo_kind = topo_map.get(self._sip_topo, 0)
        else:
            self._sip_topo_kind = 0
        sips = spec.get("system", {}).get("sips", {})
        if self._sip_topo == "ring_1d":
            self._sip_topo_w, self._sip_topo_h = 0, 0
        elif sips.get("w") is not None and sips.get("h") is not None:
            w, h = int(sips["w"]), int(sips["h"])
            if w * h != self._n_sips:
                raise ValueError(
                    f"sip layout {w}x{h} != sips.count ({self._n_sips})"
                )
            self._sip_topo_w, self._sip_topo_h = w, h
        else:
            side = int(round(math.sqrt(self._n_sips)))
            if side * side != self._n_sips:
                raise ValueError(
                    f"SIP topology '{self._sip_topo}' requires square "
                    f"sips.count or explicit sips.w/h, got {self._n_sips}"
                )
            self._sip_topo_w, self._sip_topo_h = side, side

        # IPCQ install: wire all pe0s across all cubes and SIPs
        engine = getattr(self.ctx, "engine", None)
        if engine is not None:
            from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
            configure_sfr_intercube_multisip(engine, spec, self._merged)

    def _resolve_world_size(self) -> int:
        """Derive world_size (priority: algorithm override > defaults > topology).

        ADR-0024 D1: topology fallback is SIP count. Each rank represents one
        SIP (TP dimension). Intra-SIP parallelism is expressed via DPPolicy
        inside each worker and is independent of world_size.
        Explicit ``ccl.yaml`` override still respected — legacy "rank = flat
        PE index" tests use this path.
        """
        if "world_size" in self._merged:
            return int(self._merged["world_size"])
        defaults = self._cfg_all.get("defaults", {})
        if "world_size" in defaults:
            return int(defaults["world_size"])
        spec = self.ctx.spec or {}
        return int(spec.get("system", {}).get("sips", {}).get("count", 1))

    @property
    def world_size(self) -> int:
        return self._world_size

    def all_reduce(self, tensor: Any, op: str = "sum") -> None:
        """Dispatch the configured CCL algorithm as a single kernel launch.

        Raises if ``op != "sum"`` (current kernels only implement add
        reduction) or if the tensor's shard count disagrees with the
        world_size that was installed into PE_IPCQ.
        """
        if op != "sum":
            raise NotImplementedError(f"all_reduce op={op!r} not supported")
        if tensor._handle is None:
            raise RuntimeError(
                f"Tensor '{tensor.name}' is not deployed (call torch.zeros "
                "with a DPPolicy first)"
            )
        shards = tensor._handle.shards
        if not shards:
            raise RuntimeError(
                f"all_reduce tensor '{tensor.name}' has no shards"
            )
        n_elem = shards[0].nbytes // tensor.itemsize
        kernel_fn = self._algo_module.kernel
        # Derive effective cube dims from tensor's actual shard placement
        # (may differ from topology mesh when TP uses fewer cubes).
        sip0_cubes = sorted({s.cube for s in shards if s.sip == shards[0].sip})
        eff_n_cubes = len(sip0_cubes) if sip0_cubes else 1
        if eff_n_cubes == 1:
            eff_cube_w, eff_cube_h = 1, 1
        else:
            eff_cube_w, eff_cube_h = self._cube_w, self._cube_h
        kernel_args = self._algo_module.kernel_args(
            self._world_size, n_elem,
            cube_w=eff_cube_w, cube_h=eff_cube_h,
        )

        # Resolve sip_rank from the current greenlet's bound rank
        from greenlet import getcurrent as _gc
        g = _gc()
        dist_ctx = getattr(self, "_dist_ctx", None)
        if dist_ctx is not None:
            sip_rank = int(dist_ctx._rank_by_greenlet.get(g, 0))
        else:
            sip_rank = 0

        extra_args = (
            sip_rank,
            self._sip_topo_kind,
            self._sip_topo_w,
            self._sip_topo_h,
        )

        pending = self.ctx.launch(
            self._merged["algorithm"], kernel_fn, tensor,
            *kernel_args, *extra_args,
            _defer_wait=True,
        )
        from greenlet import getcurrent
        g = getcurrent()
        if g.parent is not None and not g.parent.dead:
            # Multi-greenlet mode: hand pending to the backend-level queue so
            # the main scheduler drains. Worker just yields.
            self._pending_collective_handles.extend(pending)
            g.parent.switch()
            # On resume, all pending handles have been drained by main.
        else:
            # Single-driver (no bench scheduler): drain inline.
            for h, _sip_id, meta in pending:
                self.ctx.wait(h, _meta=meta)

    def barrier(self) -> None:
        # Single-driver model → no cross-process sync needed. Keeping the
        # method so ``dist.barrier()`` is callable (pytorch-compat surface).
        return None


class DistributedContext:
    """torch.distributed-compat facade.

    Public surface matches real PyTorch so bench code reads identically
    to a DDP training script. Single-driver semantics: ``get_rank()``
    always returns 0 because kernbench runs as one Python process;
    ``get_world_size()`` returns the CCL group size (number of PEs
    participating in the collective).
    """

    def __init__(self) -> None:
        self._backend: AhbmCCLBackend | None = None
        # ADR-0024 D2: greenlet-local rank registry. Bench launcher calls
        # _bind_rank(g, rank) when spawning workers; get_rank() resolves the
        # current greenlet to its rank. Unbound greenlets fall back to 0 for
        # single-driver test compat.
        self._rank_by_greenlet: dict = {}

    def init_process_group(
        self,
        backend: str = "ahbm",
        world_size: int | None = None,
        rank: int | None = None,
        **kwargs: Any,
    ) -> None:
        """Create the default process group.

        ``world_size`` and ``rank`` are accepted for API parity with
        ``torch.distributed.init_process_group`` but ignored — the ahbm
        backend derives both from ``ccl.yaml`` + topology automatically
        (like reading ``RANK``/``WORLD_SIZE`` env vars in real DDP).
        """
        if backend != "ahbm":
            raise ValueError(
                f"Unsupported backend '{backend}'. Only 'ahbm' is supported."
            )
        ctx = getattr(self, "_ctx_ref", None)
        if ctx is None:
            raise RuntimeError(
                "DistributedContext not bound to a RuntimeContext"
            )
        self._backend = AhbmCCLBackend(torch_ctx=ctx)
        self._backend._dist_ctx = self

    def is_initialized(self) -> bool:
        return self._backend is not None

    def get_world_size(self) -> int:
        self._ensure_initialized()
        return self._backend.world_size

    def get_rank(self) -> int:
        """Return the rank bound to the current greenlet (default 0).

        ADR-0024 D2: workers spawned by the bench launcher each get a rank
        registered via ``_bind_rank``. Callers outside any bound greenlet
        fall back to rank 0 for single-driver test compat.
        """
        self._ensure_initialized()
        from greenlet import getcurrent
        g = getcurrent()
        return int(self._rank_by_greenlet.get(g, 0))

    def _bind_rank(self, g: Any, rank: int) -> None:
        """Bind a greenlet to a rank so ``get_rank()`` returns it (ADR-0024 D2)."""
        self._rank_by_greenlet[g] = int(rank)

    def get_backend(self) -> str:
        self._ensure_initialized()
        return "ahbm"

    def all_reduce(self, tensor: Any, op: str = "sum") -> None:
        self._ensure_initialized()
        self._backend.all_reduce(tensor, op=op)

    def barrier(self) -> None:
        self._ensure_initialized()
        self._backend.barrier()

    def _ensure_initialized(self) -> None:
        if self._backend is None:
            raise RuntimeError(
                "Default process group has not been initialized. "
                "Call init_process_group(backend='ahbm') first."
            )
