"""Helpers for CCL algorithm authors (ADR-0023 D15).

These are pure utility functions usable from any kernel module:

    from kernbench.ccl.helpers import chunked, ring_step, tree_step

They keep algorithm code short and free of off-by-one bugs.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any


_DTYPE_BYTES = {
    "f16": 2, "fp16": 2, "float16": 2, "bf16": 2,
    "f32": 4, "fp32": 4, "float32": 4,
    "i8": 1, "int8": 1,
    "i16": 2, "int16": 2,
    "i32": 4, "int32": 4,
}


def _itemsize(dtype: str) -> int:
    if dtype not in _DTYPE_BYTES:
        raise ValueError(f"Unsupported dtype: {dtype}")
    return _DTYPE_BYTES[dtype]


# ── chunked ──────────────────────────────────────────────────────────


@dataclass(frozen=True)
class Chunk:
    """One chunk of a tensor used by collective algorithms."""

    addr: int
    n_elem: int
    nbytes: int


def chunked(
    base_addr: int,
    n_chunks: int,
    n_elem: int,
    dtype: str = "f16",
) -> list[Chunk]:
    """Slice a 1D buffer into ``n_chunks`` equal Chunks.

    Args:
        base_addr: starting address of the buffer.
        n_chunks: number of equal chunks to produce.
        n_elem: total number of elements (must be divisible by n_chunks).
        dtype: element type for byte-size calculation.

    Returns:
        List of ``Chunk`` objects whose addresses are consecutive.

    Raises:
        ValueError: if n_elem is not divisible by n_chunks.
    """
    if n_elem % n_chunks != 0:
        raise ValueError(
            f"chunked: n_elem ({n_elem}) not divisible by n_chunks ({n_chunks})"
        )
    per_chunk_elem = n_elem // n_chunks
    isize = _itemsize(dtype)
    per_chunk_bytes = per_chunk_elem * isize
    return [
        Chunk(
            addr=base_addr + i * per_chunk_bytes,
            n_elem=per_chunk_elem,
            nbytes=per_chunk_bytes,
        )
        for i in range(n_chunks)
    ]


# ── ring_step ────────────────────────────────────────────────────────


def ring_step(rank: int, step: int, world_size: int) -> tuple[int, int]:
    """Return ``(send_chunk_idx, recv_chunk_idx)`` for a ring algorithm step.

    Standard reduce-scatter / all-gather ring schedule:
        at step s, rank r sends chunk (r - s) and receives chunk (r - s - 1)
        modulo world_size.

    Used by ring all-reduce kernels:

        for step in range(world_size - 1):
            send_idx, recv_idx = ring_step(rank, step, world_size)
            tl.send(dir="E", src=chunks[send_idx])
            chunks[recv_idx] += tl.recv(dir="W").data
    """
    send_idx = (rank - step) % world_size
    recv_idx = (rank - step - 1) % world_size
    return send_idx, recv_idx


# ── tree_step ────────────────────────────────────────────────────────


def tree_step(rank: int, world_size: int) -> dict[str, Any]:
    """Return parent/children for binary tree rooted at rank 0.

    Returns:
        ``{"parent": int|None, "children": list[int]}``
    """
    parent = (rank - 1) // 2 if rank > 0 else None
    children: list[int] = []
    left = 2 * rank + 1
    right = 2 * rank + 2
    if left < world_size:
        children.append(left)
    if right < world_size:
        children.append(right)
    return {"parent": parent, "children": children}
