"""Tests for CCL algorithm-author helpers (ADR-0023 D15)."""
from __future__ import annotations

import pytest

from kernbench.ccl.helpers import (
    Chunk,
    chunked,
    ring_step,
    tree_step,
)


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


def test_chunked_basic():
    chunks = chunked(base_addr=0x1000, n_chunks=4, n_elem=64, dtype="f16")
    assert len(chunks) == 4
    # Each chunk has 16 elements (64 / 4)
    assert chunks[0] == Chunk(addr=0x1000, n_elem=16, nbytes=32)
    assert chunks[1] == Chunk(addr=0x1020, n_elem=16, nbytes=32)
    assert chunks[2] == Chunk(addr=0x1040, n_elem=16, nbytes=32)
    assert chunks[3] == Chunk(addr=0x1060, n_elem=16, nbytes=32)


def test_chunked_f32():
    chunks = chunked(base_addr=0x100, n_chunks=2, n_elem=8, dtype="f32")
    assert chunks[0].nbytes == 16  # 4 elem × 4 bytes
    assert chunks[1].addr == 0x100 + 16


def test_chunked_uneven_raises():
    with pytest.raises(ValueError):
        chunked(base_addr=0x100, n_chunks=3, n_elem=10, dtype="f16")


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


def test_ring_step_4_ranks():
    # Standard reduce-scatter ring step:
    # at step s, rank r sends chunk (r-s) and receives chunk (r-s-1) (mod ws)
    assert ring_step(rank=0, step=0, world_size=4) == (0, 3)
    assert ring_step(rank=0, step=1, world_size=4) == (3, 2)
    assert ring_step(rank=1, step=0, world_size=4) == (1, 0)
    assert ring_step(rank=2, step=0, world_size=4) == (2, 1)


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


def test_tree_step_root():
    info = tree_step(rank=0, world_size=7)
    assert info["parent"] is None
    assert info["children"] == [1, 2]


def test_tree_step_internal():
    info = tree_step(rank=1, world_size=7)
    assert info["parent"] == 0
    assert info["children"] == [3, 4]


def test_tree_step_leaf():
    info = tree_step(rank=4, world_size=7)
    assert info["parent"] == 1
    assert info["children"] == []