"""IPCQ install plan for AhbmCCLBackend (ADR-0023 D10/D11/D12).

Given a ccl.yaml config, the topology, and the engine, this module:

1. Loads ccl.yaml and resolves the chosen algorithm.
2. Maps each rank to a (sip, cube, pe) PE address using a linear scheme.
3. Allocates per-rank IPCQ ring buffer base addresses (synthetic but
   unique-per-PE; see notes below).
4. Builds neighbor tables via the algorithm's ``topology`` field plus the
   optional ``neighbors()`` override hook from the algorithm module.
5. Wires bidirectional credit-return SimPy Stores between every (PE, peer)
   pair.
6. Installs each PE_IPCQ component's neighbor table directly via its
   ``_install_neighbors`` sideband call (equivalent to fan-out IpcqInitMsg
   without going through fabric).

Address scheme
--------------
For the first implementation we use a synthetic address scheme that
guarantees uniqueness per (sip, cube, pe, direction) without going
through ``PEMemAllocator``. The address is encoded as:

    base = IPCQ_BASE | (sip << 40) | (cube << 32) | (pe << 24)
    rx_base[direction_idx] = base + direction_idx * (n_slots * slot_size)

The ``buffer_kind`` (tcm/hbm/sram) selects the *MemoryStore space* into
which data is written. Within a space, addresses are unique per PE so
the existing MemoryStore (``{space: {addr: ndarray}}``) handles them
naturally.

This bypasses the topology's address resolver / PhysAddr encoding and
treats IPCQ buffers as a separate, parallel address namespace. Real PA
encoding can be plugged in later without changing the rest of the design.
"""
from __future__ import annotations

from pathlib import Path
from typing import Any

import simpy
import yaml

from kernbench.ccl.topologies import resolve_topology
from kernbench.common.ipcq_types import (
    IpcqEndpoint,
    IpcqInitEntry,
)
from kernbench.runtime_api.kernel import IpcqInitMsg


# IPCQ synthetic address space top bit
_IPCQ_BASE = 1 << 60


def _ipcq_base_for_pe(sip: int, cube: int, pe: int) -> int:
    return _IPCQ_BASE | (sip << 40) | (cube << 32) | (pe << 24)


# ── ccl.yaml loading ─────────────────────────────────────────────────


def load_ccl_config(path: str | Path | None = None) -> dict:
    """Load and validate ccl.yaml. Searches cwd and project root."""
    if path is None:
        candidates = [
            Path.cwd() / "ccl.yaml",
            Path(__file__).resolve().parents[3] / "ccl.yaml",
        ]
        for p in candidates:
            if p.exists():
                path = p
                break
    if path is None:
        raise FileNotFoundError(
            "ccl.yaml not found. Place it at project root or cwd."
        )
    with open(path) as f:
        cfg = yaml.safe_load(f)
    if "defaults" not in cfg:
        raise ValueError("ccl.yaml missing 'defaults' section")
    if "algorithms" not in cfg:
        raise ValueError("ccl.yaml missing 'algorithms' section")
    return cfg


def resolve_algorithm_config(cfg: dict, name: str | None = None) -> dict:
    """Merge defaults with the chosen algorithm's overrides.

    Returns a flat dict with at minimum: module, topology, buffer_kind,
    backpressure, n_slots, slot_size, ipcq_credit_size_bytes, world_size.
    """
    defaults = dict(cfg.get("defaults", {}))
    algo_name = name or defaults.get("algorithm")
    if algo_name is None:
        raise ValueError("ccl.yaml: defaults.algorithm not set")
    algos = cfg.get("algorithms", {})
    if algo_name not in algos:
        raise ValueError(
            f"ccl.yaml: algorithm '{algo_name}' not in algorithms section"
        )
    merged = defaults.copy()
    merged.update(algos[algo_name])
    merged["algorithm"] = algo_name
    return merged


# ── rank → PE mapping ────────────────────────────────────────────────


def linear_rank_to_pe(rank: int, spec: dict) -> tuple[int, int, int]:
    """Map a rank to (sip, cube, pe) using linear topology order."""
    sips = spec["system"]["sips"]["count"]
    cubes_per_sip = spec["sip"]["cube_mesh"]["w"] * spec["sip"]["cube_mesh"]["h"]
    pe_layout = spec["cube"]["pe_layout"]
    pes_per_cube = pe_layout["pe_per_corner"] * len(pe_layout["corners"])

    pes_per_sip = cubes_per_sip * pes_per_cube
    if rank >= sips * pes_per_sip:
        raise ValueError(
            f"rank {rank} exceeds total PE count {sips * pes_per_sip}"
        )
    sip = rank // pes_per_sip
    rem = rank % pes_per_sip
    cube = rem // pes_per_cube
    pe = rem % pes_per_cube
    return sip, cube, pe


# ── Install plan ─────────────────────────────────────────────────────


def install_ipcq(
    engine: Any,
    spec: dict,
    cfg: dict,
    algo_module: Any | None = None,
    rank_to_pe: list[tuple[int, int, int]] | None = None,
) -> dict[str, Any]:
    """Build neighbor tables and install them in every participating PE_IPCQ.

    Args:
        engine: GraphEngine with ``_components`` dict
        spec: topology spec dict
        cfg: merged algorithm config (from ``resolve_algorithm_config``)
        algo_module: optional algorithm Python module (for neighbors override)
        rank_to_pe: optional explicit rank → (sip, cube, pe) mapping. If
                    None, the default linear mapping is used.

    Returns:
        A diagnostics dict with the install plan (rank → PE map, neighbor table).
    """
    if "world_size" in cfg:
        world_size = int(cfg["world_size"])
    else:
        # Topology-derived fallback (mirrors AhbmCCLBackend / RuntimeContext).
        sips = int(spec.get("system", {}).get("sips", {}).get("count", 1))
        cm = spec.get("sip", {}).get("cube_mesh", {})
        cubes_per_sip = int(cm.get("w", 1)) * int(cm.get("h", 1))
        pl = spec.get("cube", {}).get("pe_layout", {})
        corners = pl.get("corners", [])
        pe_per_corner = int(pl.get("pe_per_corner", 1))
        pes_per_cube = pe_per_corner * max(len(corners), 1)
        world_size = sips * cubes_per_sip * pes_per_cube
    buffer_kind = cfg["buffer_kind"]
    n_slots = int(cfg["n_slots"])
    slot_size = int(cfg["slot_size"])
    backpressure = cfg["backpressure"]
    credit_size_bytes = int(cfg.get("ipcq_credit_size_bytes", 16))

    # Step 1: rank → (sip, cube, pe)
    if rank_to_pe is not None:
        if len(rank_to_pe) != world_size:
            raise ValueError(
                f"rank_to_pe has {len(rank_to_pe)} entries but world_size={world_size}"
            )
        rank_pe = list(rank_to_pe)
    else:
        rank_pe: list[tuple[int, int, int]] = [
            linear_rank_to_pe(r, spec) for r in range(world_size)
        ]
    pe_to_rank = {(s, c, p): r for r, (s, c, p) in enumerate(rank_pe)}

    # Step 2: resolve topology fn (with optional override)
    topo_fn = resolve_topology(cfg["topology"], algo_module=algo_module)

    # Build per-rank neighbor map
    neighbor_table: dict[int, dict[str, int]] = {}
    for r in range(world_size):
        neighbor_table[r] = topo_fn(r, world_size)

    # Step 3: pull the live engine reference for each PE_IPCQ
    components = engine._components
    pe_ipcq_id = lambda s, c, p: f"sip{s}.cube{c}.pe{p}.pe_ipcq"

    # Step 4: per-PE rx_base address and per-PE credit_inbox
    direction_keys = sorted({d for nt in neighbor_table.values() for d in nt})
    direction_idx = {d: i for i, d in enumerate(direction_keys)}
    bytes_per_direction = n_slots * slot_size

    def rx_base(s: int, c: int, p: int, d: str) -> int:
        return _ipcq_base_for_pe(s, c, p) + direction_idx[d] * bytes_per_direction

    # Wire bidirectional credit stores: backend creates the SimPy Stores
    # by reading each rank's PE_IPCQ.credit_inbox property.
    rank_to_credit_inbox: dict[int, simpy.Store] = {}
    for r, (s, c, p) in enumerate(rank_pe):
        comp = components[pe_ipcq_id(s, c, p)]
        # Trigger lazy creation of credit_inbox if not yet started.
        # PE_IPCQ.start() creates it; we ensure it exists.
        if comp._credit_inbox is None:
            comp._credit_inbox = simpy.Store(engine._env)
        rank_to_credit_inbox[r] = comp.credit_inbox

    # Step 5: build IpcqInitMsg per rank and call _install_neighbors directly
    plan: dict[str, Any] = {
        "world_size": world_size,
        "rank_to_pe": rank_pe,
        "buffer_kind": buffer_kind,
        "neighbor_table": neighbor_table,
    }

    _OPPOSITE_DIR = {
        "E": "W", "W": "E", "N": "S", "S": "N",
        "intra_E": "intra_W", "intra_W": "intra_E",
        "intra_N": "intra_S", "intra_S": "intra_N",
        "global_E": "global_W", "global_W": "global_E",
        "global_N": "global_S", "global_S": "global_N",
    }

    def reverse_direction(my_rank: int, peer_rank: int, my_dir: str) -> str | None:
        """Find peer's direction that reciprocates my_dir→peer_rank.

        Prefer the OPPOSITE direction (E↔W, N↔S) when the peer has it
        pointing back to us (ADR-0025 D1). This matters in 2-rank
        bidirectional rings where both E and W on one side point to the
        same peer — without the preference, dict-order first-match would
        route data into the wrong rx slot. Falls back to any direction
        pointing back for topologies without an opposite convention
        (e.g. tree_binary's parent/child).
        """
        nt = neighbor_table[peer_rank]
        opp = _OPPOSITE_DIR.get(my_dir)
        if opp is not None and nt.get(opp) == my_rank:
            return opp
        for d, target in nt.items():
            if target == my_rank:
                return d
        return None

    for r, (s, c, p) in enumerate(rank_pe):
        my_pe_ipcq = components[pe_ipcq_id(s, c, p)]
        nbrs = neighbor_table[r]
        entries: list[IpcqInitEntry] = []
        for d, peer_rank in nbrs.items():
            if peer_rank is None:
                continue
            peer_s, peer_c, peer_p = rank_pe[peer_rank]
            peer_dir = reverse_direction(r, peer_rank, d)
            if peer_dir is None:
                # Peer doesn't have a reverse entry — skip (asymmetric topology)
                continue
            peer_endpoint = IpcqEndpoint(
                sip=peer_s, cube=peer_c, pe=peer_p,
                buffer_kind=buffer_kind,
                rx_base_pa=rx_base(peer_s, peer_c, peer_p, peer_dir),
                rx_base_va=0,
                n_slots=n_slots, slot_size=slot_size,
            )
            entries.append(IpcqInitEntry(
                direction=d,
                peer=peer_endpoint,
                my_rx_base_pa=rx_base(s, c, p, d),
                my_rx_base_va=0,
                n_slots=n_slots, slot_size=slot_size,
                peer_credit_store=rank_to_credit_inbox[peer_rank],
            ))
        msg = IpcqInitMsg(
            correlation_id="ccl_init", request_id=f"init_r{r}",
            target_sips=(s,), target_cubes=(c,), target_pe=p,
            entries=tuple(entries),
            backpressure_mode=backpressure,
            buffer_kind=buffer_kind,
            credit_size_bytes=credit_size_bytes,
        )
        my_pe_ipcq._install_neighbors(msg)

    return plan
