"""Deterministic memory, KV, and network planner for the GLM-5.2 Max alpha route.

Everything here is arithmetic over the exact pinned artifact bytes and the exact
pinned architecture. There is no measurement, no probing, and no heuristic tuned
to a result — the planner is written *before* the target runs so that a later
story cannot discover a topology that "works" and then rationalise it.

Three ideas do the real work.

**Unified memory is one pool.** On an integrated-GPU machine the "VRAM" the driver
reports is carved out of the same physical DRAM the OS is already counting. Adding
them produces a node that appears to hold twice what it holds, and the failure mode
is not a clean admission rejection — it is an OOM or a swap-thrash halfway through
a 200 GiB load. :class:`NodeMemory` therefore refuses to be constructed from an
additive claim about one shared pool.

**The reserve is not optional headroom.** Weights plus KV are not the whole
resident cost: backend workspaces, quantization scratch, the graph plan, the
process, and the OS all live outside them, and the largest of those scale with the
backend rather than with the shard. Alpha reserves ``max(20% of physically usable
memory, 8 GiB)`` per node, and the *remainder* is the placement budget.

**Equal layer counts are not equal bytes.** Embeddings and the output head are
endpoint-only; three layers are dense and 75 are MoE; shared experts, indexer
tensors, and quant block alignment all skew the per-node share. Until DGR-018/019
report measured per-tensor placement, the planner carries an explicit
:data:`PLACEMENT_IMBALANCE_FACTOR` and reports the arithmetic minimum and the
recommended count as two separate numbers. The arithmetic minimum is a fit probe;
it is admissible only with exact measured placement evidence behind it.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Literal

from .manifest import GIB, ArchitectureSnapshot, TargetManifest

# Q8_0 stores 32 int8 quants plus one fp16 scale per block: 34 bytes / 32 values.
Q8_0_BYTES_PER_VALUE = 34 / 32
F16_BYTES_PER_VALUE = 2.0

KV_DTYPES: dict[str, float] = {
    "Q8_0": Q8_0_BYTES_PER_VALUE,
    "F16": F16_BYTES_PER_VALUE,
}

# The alpha KV configuration, locked by the roadmap.
ALPHA_KV_DTYPE = "Q8_0"
ALPHA_CONTEXT_TOKENS = 16384
ALPHA_CONCURRENCY = 1

# The reserve every node holds outside its weight-plus-KV placement budget.
RESERVE_FRACTION = 0.20
RESERVE_FLOOR_GIB = 8.0

# The aggregate runtime-accessible memory at which the artifact *just* fits.
# This is an experimental hard-fit floor, not an operational envelope: it has no
# room for a backend that allocates more scratch than another, and none for the
# imbalance below.
AGGREGATE_HARD_FIT_FLOOR_GIB = 224.0

# How much more than an equal share the worst-placed node is expected to hold.
# 1.10 is the roadmap's recommended-topology column expressed as arithmetic: it
# reproduces 10 / 6 / 5 / 3 / 3 nodes for the 32 / 48 / 64 / 96 / 128 GiB tiers.
# DGR-019 must replace it with measured per-tensor placement.
PLACEMENT_IMBALANCE_FACTOR = 1.10

# Alpha network floor. A link rate is a bandwidth claim, never a speed claim.
MIN_LINK_RATE_GBPS = 2.5
RECOMMENDED_LINK_RATE_GBPS = 10.0

BF16_BYTES = 2
DSA_SIDEBAND_INT32_BYTES = 4

IndexerLayout = Literal["optimized", "conservative"]


class ResourcePlanError(ValueError):
    """Raised when a node or route cannot be accounted for honestly."""


@dataclass(frozen=True)
class NodeMemory:
    """One node's physically usable memory, counted once.

    ``physical_usable_gib`` is what the node can actually place bytes into after
    firmware and fixed carve-outs — not the marketing capacity, and not a sum of
    two views of the same DRAM.
    """

    name: str
    physical_usable_gib: float
    unified: bool

    def __post_init__(self) -> None:
        if not isinstance(self.name, str) or not self.name.strip():
            raise ResourcePlanError("node name must be a non-empty physical-host identity")
        if (
            isinstance(self.physical_usable_gib, bool)
            or not isinstance(self.physical_usable_gib, (int, float))
            or not math.isfinite(self.physical_usable_gib)
            or self.physical_usable_gib <= 0
        ):
            raise ResourcePlanError(
                f"node {self.name!r} must declare finite positive usable memory"
            )

    @classmethod
    def from_host(
        cls,
        name: str,
        *,
        system_ram_gib: float,
        gpu_memory_gib: float = 0.0,
        unified: bool,
    ) -> "NodeMemory":
        """Build a node from a host's reported RAM and GPU memory.

        On a unified machine the GPU memory *is* system RAM, so it is counted once
        and never added. Passing a non-zero ``gpu_memory_gib`` alongside
        ``unified=True`` is the double-count this project has already decided is a
        bug (RALPH-CONTEXT runtime decision 16), so it is rejected rather than
        silently discarded: a caller who believes an integrated GPU adds memory has
        a wrong model of the machine, and quietly ignoring the argument would let
        that belief survive.
        """
        if not isinstance(unified, bool):
            raise ResourcePlanError(f"node {name!r} unified flag must be boolean")
        for value, label, allow_zero in (
            (system_ram_gib, "system RAM", False),
            (gpu_memory_gib, "GPU memory", True),
        ):
            if (
                isinstance(value, bool)
                or not isinstance(value, (int, float))
                or not math.isfinite(value)
                or value < 0
                or (not allow_zero and value == 0)
            ):
                qualifier = "finite non-negative" if allow_zero else "finite positive"
                raise ResourcePlanError(f"node {name!r} must declare {qualifier} {label}")
        if unified:
            if gpu_memory_gib:
                raise ResourcePlanError(
                    f"node {name!r} declares unified memory and {gpu_memory_gib} GiB of "
                    "separate GPU memory. Integrated-GPU memory is carved out of the same "
                    "physical DRAM as system RAM; adding them double-counts one pool. "
                    "Pass unified=True with system_ram_gib only."
                )
            usable = system_ram_gib
        else:
            usable = system_ram_gib + gpu_memory_gib
        return cls(name=name, physical_usable_gib=usable, unified=unified)

    @property
    def reserve_gib(self) -> float:
        """``max(20% of physically usable memory, 8 GiB)``."""
        return max(RESERVE_FRACTION * self.physical_usable_gib, RESERVE_FLOOR_GIB)

    @property
    def placement_budget_gib(self) -> float:
        """What remains for weights plus KV after the reserve."""
        return self.physical_usable_gib - self.reserve_gib


def kv_bytes(
    snapshot: ArchitectureSnapshot,
    *,
    context_tokens: int = ALPHA_CONTEXT_TOKENS,
    concurrency: int = ALPHA_CONCURRENCY,
    dtype: str = ALPHA_KV_DTYPE,
    indexer_layout: IndexerLayout = "conservative",
    include_indexer: bool = True,
) -> int:
    """Bytes of MLA (and DSA indexer) KV cache for the whole model.

    ``indexer_layout`` is the honest part. Correct DSA only needs indexer keys for
    the Full producer layers, but the current experimental implementation may
    allocate them across every backbone layer. Alpha budgets ``conservative``
    (all 78) so that a route admitted by this planner cannot be surprised by the
    implementation it actually gets.
    """
    if (
        not isinstance(context_tokens, int)
        or isinstance(context_tokens, bool)
        or context_tokens <= 0
        or not isinstance(concurrency, int)
        or isinstance(concurrency, bool)
        or concurrency <= 0
    ):
        raise ResourcePlanError("context_tokens and concurrency must be positive integers")
    if dtype not in KV_DTYPES:
        raise ResourcePlanError(
            f"unsupported KV dtype {dtype!r}; alpha locks {ALPHA_KV_DTYPE} "
            f"(known: {', '.join(sorted(KV_DTYPES))})"
        )
    bytes_per_value = KV_DTYPES[dtype]

    layers = int(snapshot["num_hidden_layers"])
    mla_values = int(snapshot["mla_cached_values_per_token_per_layer"])
    total_values = mla_values * layers

    if include_indexer:
        if indexer_layout == "optimized":
            indexer_layers = int(snapshot["indexer_full_layers"])
        elif indexer_layout == "conservative":
            indexer_layers = layers
        else:  # pragma: no cover - Literal keeps this unreachable from typed callers
            raise ResourcePlanError(f"unknown indexer_layout {indexer_layout!r}")
        total_values += int(snapshot["index_head_dim"]) * indexer_layers

    return int(total_values * context_tokens * concurrency * bytes_per_value)


@dataclass(frozen=True)
class TopologyPlan:
    """The node count a homogeneous tier needs, and how it was reached."""

    physical_usable_gib: float
    reserve_gib: float
    placement_budget_gib: float
    weight_gib: float
    kv_gib: float
    total_placement_gib: float
    arithmetic_minimum_nodes: int
    recommended_nodes: int
    imbalance_factor: float

    @property
    def is_arithmetic_minimum_topology(self) -> bool:
        """True when the recommendation offers no imbalance headroom at all."""
        return self.recommended_nodes == self.arithmetic_minimum_nodes

    def to_dict(self) -> dict:
        return {
            "physical_usable_gib": round(self.physical_usable_gib, 3),
            "reserve_gib": round(self.reserve_gib, 3),
            "placement_budget_gib": round(self.placement_budget_gib, 3),
            "weight_gib": round(self.weight_gib, 3),
            "kv_gib": round(self.kv_gib, 3),
            "total_placement_gib": round(self.total_placement_gib, 3),
            "arithmetic_minimum_nodes": self.arithmetic_minimum_nodes,
            "recommended_nodes": self.recommended_nodes,
            "imbalance_factor": self.imbalance_factor,
        }


def plan_topology(
    manifest: TargetManifest,
    snapshot: ArchitectureSnapshot,
    *,
    physical_usable_gib: float,
    context_tokens: int = ALPHA_CONTEXT_TOKENS,
    concurrency: int = ALPHA_CONCURRENCY,
    kv_dtype: str = ALPHA_KV_DTYPE,
    indexer_layout: IndexerLayout = "conservative",
    imbalance_factor: float = PLACEMENT_IMBALANCE_FACTOR,
) -> TopologyPlan:
    """Minimum and recommended node count for a homogeneous tier of this size."""
    if (
        isinstance(imbalance_factor, bool)
        or not isinstance(imbalance_factor, (int, float))
        or not math.isfinite(imbalance_factor)
        or imbalance_factor < 1.0
    ):
        raise ResourcePlanError(
            "imbalance_factor must be finite and at least 1.0; a lower value would "
            "assume the worst-placed node holds less than an equal share"
        )

    node = NodeMemory(
        name=f"{physical_usable_gib:g}GiB-tier",
        physical_usable_gib=physical_usable_gib,
        unified=False,
    )
    budget = node.placement_budget_gib
    if budget <= 0:
        raise ResourcePlanError(
            f"a {physical_usable_gib:g} GiB node has no placement budget after its "
            f"{node.reserve_gib:.1f} GiB reserve"
        )

    weight_gib = manifest.total_bytes / GIB
    kv_gib = (
        kv_bytes(
            snapshot,
            context_tokens=context_tokens,
            concurrency=concurrency,
            dtype=kv_dtype,
            indexer_layout=indexer_layout,
        )
        / GIB
    )
    total = weight_gib + kv_gib

    return TopologyPlan(
        physical_usable_gib=physical_usable_gib,
        reserve_gib=node.reserve_gib,
        placement_budget_gib=budget,
        weight_gib=weight_gib,
        kv_gib=kv_gib,
        total_placement_gib=total,
        arithmetic_minimum_nodes=math.ceil(total / budget),
        recommended_nodes=math.ceil(total * imbalance_factor / budget),
        imbalance_factor=imbalance_factor,
    )


@dataclass(frozen=True)
class RouteFit:
    """Whether a concrete, possibly heterogeneous set of nodes can hold the target."""

    node_count: int
    aggregate_usable_gib: float
    aggregate_placement_budget_gib: float
    required_placement_gib: float
    fits: bool
    meets_hard_fit_floor: bool
    no_single_node_can_admit_target: bool
    headroom_gib: float
    reasons: tuple[str, ...]

    def to_dict(self) -> dict:
        return {
            "node_count": self.node_count,
            "aggregate_usable_gib": round(self.aggregate_usable_gib, 3),
            "aggregate_placement_budget_gib": round(self.aggregate_placement_budget_gib, 3),
            "required_placement_gib": round(self.required_placement_gib, 3),
            "fits": self.fits,
            "meets_hard_fit_floor": self.meets_hard_fit_floor,
            "no_single_node_can_admit_target": self.no_single_node_can_admit_target,
            "headroom_gib": round(self.headroom_gib, 3),
            "reasons": list(self.reasons),
        }


def plan_route(
    manifest: TargetManifest,
    snapshot: ArchitectureSnapshot,
    nodes: list[NodeMemory],
    *,
    context_tokens: int = ALPHA_CONTEXT_TOKENS,
    concurrency: int = ALPHA_CONCURRENCY,
    kv_dtype: str = ALPHA_KV_DTYPE,
    indexer_layout: IndexerLayout = "conservative",
) -> RouteFit:
    """Evaluate a concrete route. Every node's memory is already counted once."""
    if len(nodes) < 2:
        raise ResourcePlanError(
            "the alpha target is distributed by definition; a route needs at least two "
            "physical nodes"
        )
    names = [node.name for node in nodes]
    if len(set(names)) != len(names):
        raise ResourcePlanError(
            "duplicate node names in the route; one physical machine counted twice is "
            "the same double-count as adding integrated-GPU memory to system RAM"
        )

    weight_gib = manifest.total_bytes / GIB
    kv_gib = (
        kv_bytes(
            snapshot,
            context_tokens=context_tokens,
            concurrency=concurrency,
            dtype=kv_dtype,
            indexer_layout=indexer_layout,
        )
        / GIB
    )
    required = weight_gib + kv_gib

    aggregate_usable = sum(node.physical_usable_gib for node in nodes)
    aggregate_budget = sum(node.placement_budget_gib for node in nodes)
    fits = aggregate_budget >= required
    largest_budget = max(node.placement_budget_gib for node in nodes)
    no_single_node = largest_budget < required

    reasons: list[str] = []
    if not fits:
        reasons.append(
            f"aggregate placement budget {aggregate_budget:.1f} GiB is below the "
            f"{required:.1f} GiB the target needs after each node's reserve"
        )
    if not no_single_node:
        reasons.append(
            "at least one node could admit the complete target alone; that is a "
            "single-host run, not distributed alpha"
        )
    if aggregate_usable < AGGREGATE_HARD_FIT_FLOOR_GIB:
        reasons.append(
            f"aggregate usable memory {aggregate_usable:.1f} GiB is below the "
            f"{AGGREGATE_HARD_FIT_FLOOR_GIB:g} GiB experimental hard-fit floor"
        )

    return RouteFit(
        node_count=len(nodes),
        aggregate_usable_gib=aggregate_usable,
        aggregate_placement_budget_gib=aggregate_budget,
        required_placement_gib=required,
        fits=fits,
        meets_hard_fit_floor=aggregate_usable >= AGGREGATE_HARD_FIT_FLOOR_GIB,
        no_single_node_can_admit_target=no_single_node,
        headroom_gib=aggregate_budget - required,
        reasons=tuple(reasons),
    )


@dataclass(frozen=True)
class SeamPlan:
    """Bytes and latency across the activation seams of a route.

    Bandwidth and latency are reported separately on purpose. Decode moves almost
    nothing — 12 KiB per token per seam — so a faster link barely helps it. What
    decode pays is *serial*: every generated token crosses every seam in order, so
    the cost that matters is ``seams x per-hop latency``. A route that claims to be
    fast because it is on 10 GbE has confused the two.
    """

    node_count: int
    seam_count: int
    hidden_size: int
    bytes_per_token_per_seam: int
    prefill_bytes_per_seam: int
    decode_bytes_per_seam_per_token: int
    dsa_sideband_bytes_per_query: int
    link_rate_gbps: float
    meets_alpha_minimum: bool
    is_recommended_link: bool
    decode_serialization_ms_per_token: float
    decode_latency_ms_per_token: float
    decode_bandwidth_share_ms_per_token: float
    prefill_serialization_ms: float

    def to_dict(self) -> dict:
        return {
            "node_count": self.node_count,
            "seam_count": self.seam_count,
            "hidden_size": self.hidden_size,
            "bytes_per_token_per_seam": self.bytes_per_token_per_seam,
            "prefill_bytes_per_seam": self.prefill_bytes_per_seam,
            "decode_bytes_per_seam_per_token": self.decode_bytes_per_seam_per_token,
            "dsa_sideband_bytes_per_query": self.dsa_sideband_bytes_per_query,
            "link_rate_gbps": self.link_rate_gbps,
            "meets_alpha_minimum": self.meets_alpha_minimum,
            "is_recommended_link": self.is_recommended_link,
            "decode_serialization_ms_per_token": round(self.decode_serialization_ms_per_token, 4),
            "decode_latency_ms_per_token": round(self.decode_latency_ms_per_token, 4),
            "decode_bandwidth_share_ms_per_token": round(
                self.decode_bandwidth_share_ms_per_token, 4
            ),
            "prefill_serialization_ms": round(self.prefill_serialization_ms, 3),
        }


def plan_seams(
    snapshot: ArchitectureSnapshot,
    *,
    node_count: int,
    context_tokens: int = ALPHA_CONTEXT_TOKENS,
    link_rate_gbps: float = MIN_LINK_RATE_GBPS,
    per_hop_latency_ms: float = 0.5,
) -> SeamPlan:
    """Model seam bytes, wire serialization, and serial per-hop latency separately."""
    if not isinstance(node_count, int) or isinstance(node_count, bool) or node_count < 2:
        raise ResourcePlanError("a seam exists only between two nodes")
    if not isinstance(context_tokens, int) or isinstance(context_tokens, bool) or context_tokens <= 0:
        raise ResourcePlanError("context_tokens must be a positive integer")
    if (
        isinstance(link_rate_gbps, bool)
        or not isinstance(link_rate_gbps, (int, float))
        or not math.isfinite(link_rate_gbps)
        or link_rate_gbps <= 0
    ):
        raise ResourcePlanError("link_rate_gbps must be finite and positive")
    if (
        isinstance(per_hop_latency_ms, bool)
        or not isinstance(per_hop_latency_ms, (int, float))
        or not math.isfinite(per_hop_latency_ms)
        or per_hop_latency_ms < 0
    ):
        raise ResourcePlanError("per_hop_latency_ms must be finite and non-negative")

    hidden = int(snapshot["hidden_size"])
    bytes_per_token = hidden * BF16_BYTES
    seams = node_count - 1

    bits_per_ms = link_rate_gbps * 1e9 / 1e3
    decode_serialization_ms = (bytes_per_token * 8) / bits_per_ms
    prefill_serialization_ms = (bytes_per_token * context_tokens * 8) / bits_per_ms

    return SeamPlan(
        node_count=node_count,
        seam_count=seams,
        hidden_size=hidden,
        bytes_per_token_per_seam=bytes_per_token,
        prefill_bytes_per_seam=bytes_per_token * context_tokens,
        decode_bytes_per_seam_per_token=bytes_per_token,
        dsa_sideband_bytes_per_query=int(snapshot["index_topk"]) * DSA_SIDEBAND_INT32_BYTES,
        link_rate_gbps=link_rate_gbps,
        meets_alpha_minimum=link_rate_gbps >= MIN_LINK_RATE_GBPS,
        is_recommended_link=link_rate_gbps >= RECOMMENDED_LINK_RATE_GBPS,
        decode_serialization_ms_per_token=decode_serialization_ms * seams,
        decode_latency_ms_per_token=per_hop_latency_ms * seams,
        decode_bandwidth_share_ms_per_token=decode_serialization_ms * seams,
        prefill_serialization_ms=prefill_serialization_ms * seams,
    )


ALPHA_TIERS_GIB: tuple[float, ...] = (32.0, 48.0, 64.0, 96.0, 128.0)


def plan_all_tiers(
    manifest: TargetManifest, snapshot: ArchitectureSnapshot
) -> dict[str, TopologyPlan]:
    """The alpha tier table, recomputed from the pinned artifact and architecture."""
    return {
        f"{tier:g}": plan_topology(manifest, snapshot, physical_usable_gib=tier)
        for tier in ALPHA_TIERS_GIB
    }