"""Tests for CCL backend install (ADR-0023 D10/D11)."""
from __future__ import annotations

from kernbench.ccl.install import (
    install_ipcq,
    linear_rank_to_pe,
    load_ccl_config,
    resolve_algorithm_config,
)
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import resolve_topology


def _engine():
    topo = resolve_topology("topology.yaml").topology_obj
    return GraphEngine(topo, enable_data=True), topo


def test_load_ccl_config():
    cfg = load_ccl_config()
    assert "defaults" in cfg
    assert "algorithms" in cfg


def test_resolve_algorithm_config_default():
    cfg = load_ccl_config()
    merged = resolve_algorithm_config(cfg)
    assert merged["algorithm"] == cfg["defaults"]["algorithm"]
    # ccl.yaml no longer carries defaults.world_size — backend derives
    # it from topology.yaml at install time. Just check the field is
    # absent here (verified per-test where install_ipcq is called).
    assert "world_size" not in merged or merged["world_size"] >= 1


def test_resolve_algorithm_config_override():
    cfg = load_ccl_config()
    merged = resolve_algorithm_config(cfg, name="ring_allreduce_hbm")
    assert merged["algorithm"] == "ring_allreduce_hbm"
    assert merged["buffer_kind"] == "hbm"  # algo override
    # defaults still apply
    assert merged["n_slots"] == cfg["defaults"]["n_slots"]


def test_linear_rank_to_pe():
    engine, topo = _engine()
    spec = topo.spec
    # Cube 0 of SIP 0
    assert linear_rank_to_pe(0, spec) == (0, 0, 0)
    assert linear_rank_to_pe(7, spec) == (0, 0, 7)
    # Should not exceed total PE count
    pes_per_sip = (
        spec["sip"]["cube_mesh"]["w"] * spec["sip"]["cube_mesh"]["h"]
        * spec["cube"]["pe_layout"]["pe_per_corner"]
        * len(spec["cube"]["pe_layout"]["corners"])
    )
    sips = spec["system"]["sips"]["count"]
    total = sips * pes_per_sip
    assert total >= 8


def test_install_ipcq_neighbors_correct():
    engine, topo = _engine()
    cfg = load_ccl_config()
    merged = resolve_algorithm_config(cfg, name="ring_allreduce_tcm")
    # Force a single-cube 8-rank install for the assertions below.
    merged["world_size"] = 8
    plan = install_ipcq(engine, topo.spec, merged)

    assert plan["world_size"] == 8
    assert plan["buffer_kind"] == "tcm"

    # Each rank should have E and W entries
    for r, nbrs in plan["neighbor_table"].items():
        assert "E" in nbrs
        assert "W" in nbrs

    # Inspect installed PE_IPCQ for rank 0
    ipcq = engine._components["sip0.cube0.pe0.pe_ipcq"]
    qp_e = ipcq.queue_pairs["E"]
    qp_w = ipcq.queue_pairs["W"]
    assert qp_e["peer"].pe == 1   # rank 0's E neighbor is rank 1
    assert qp_w["peer"].pe == 7   # rank 0's W neighbor is rank 7
    # rx_base addresses should be unique
    assert qp_e["my_rx_base_pa"] != qp_w["my_rx_base_pa"]


def test_install_ipcq_credit_stores_wired():
    engine, topo = _engine()
    cfg = load_ccl_config()
    merged = resolve_algorithm_config(cfg, name="ring_allreduce_tcm")
    merged["world_size"] = 8
    install_ipcq(engine, topo.spec, merged)

    # rank 0 (pe0) sending E goes to rank 1 (pe1)
    # rank 0's peer_credit_store on E direction should equal rank 1's credit_inbox
    pe0 = engine._components["sip0.cube0.pe0.pe_ipcq"]
    pe1 = engine._components["sip0.cube0.pe1.pe_ipcq"]

    qp_e = pe0.queue_pairs["E"]
    assert qp_e["peer_credit_store"] is pe1.credit_inbox