"""Typed AlgoKnob hierarchy.

Knobs are frozen dataclasses keyed by their own class — _index_knobs() lets
backends look up the value with ``knobs.get(KnobClass)`` rather than scanning
a list. Fleet-level knobs (set once at Fleet construction) and Handle-level
knobs (set per dispatch/combine) live in the same namespace; each backend's
constructor inspects the relevant set.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, FrozenSet, Sequence

if TYPE_CHECKING:
    import torch

from .config import QuantType


class AlgoKnob:
    """Base marker for both Fleet and Handle knobs."""


# ---------------------------------------------------------------- Fleet-level


@dataclass(frozen=True)
class FleetAlgoKnobQuantization(AlgoKnob):
    """Enable per-Fleet quantization features."""

    quants: FrozenSet[QuantType] = field(default_factory=frozenset)


@dataclass(frozen=True)
class FleetAlgoKnobNumChannelsPerRank(AlgoKnob):
    n: int


@dataclass(frozen=True)
class FleetAlgoKnobNumQpsPerRank(AlgoKnob):
    n: int


@dataclass(frozen=True)
class FleetAlgoKnobRdmaBufferSize(AlgoKnob):
    bytes_: int


@dataclass(frozen=True)
class FleetAlgoKnobTopologyCapacity(AlgoKnob):
    """Reserve transport state for up to this many ranks (grow/shrink later)."""

    n: int


@dataclass(frozen=True)
class FleetAlgoKnobFaultTolerance(AlgoKnob):
    """Enable rank masking so a dead/slow peer degrades instead of trapping.

    Without this knob a peer that stops responding during dispatch or combine
    trips a GPU ``trap()`` and takes the job down. With it, both transports
    mask the offending rank, skip it, and let the collective complete — its
    experts simply contribute nothing (see the note on dropped tokens below).

    LOW_LATENCY only, on both transports: nccl_ep leaves its mask buffer NULL
    under HIGH_THROUGHPUT (later mask calls then abort the process) and
    nixl_ep has no HT mask support at all. ``validate_fleet_params`` rejects
    the combination up front.

    The runtime API lives on the Fleet: :meth:`~flashinfer.moe_ep.core.comm.
    fleet.Fleet.query_fault`, ``query_active_mask``, ``set_active_mask``,
    ``reconcile_active_mask`` and ``clear_faults``. See
    ``docs/design_docs/moe_ep_runbook.md`` for the recovery state machine.

    Note: a masked rank's experts are dropped from combine with **no** topk
    re-normalization, so affected tokens come out scaled by the surviving
    weight sum. That is a deliberate choice — see the runbook.
    """

    enabled: bool = True
    # GPU wait-loop timeout. 0 = transport default (nccl_ep ~100 s, nixl_ep 30 s).
    # Too low marks merely-slow ranks dead; prefer the default unless testing.
    timeout_ms: int = 0
    # Wall-clock budget for reconcile_active_mask()'s store rendezvous.
    reconcile_timeout_s: float = 30.0
    # Extra budget for coordinator takeover when the elected coordinator
    # also turns out to be dead.
    coordinator_takeover_s: float = 10.0

    def __post_init__(self) -> None:
        if self.timeout_ms < 0:
            raise ValueError(
                f"FleetAlgoKnobFaultTolerance.timeout_ms must be >= 0 "
                f"(0 = transport default), got {self.timeout_ms}"
            )
        if self.reconcile_timeout_s <= 0 or self.coordinator_takeover_s <= 0:
            raise ValueError(
                "FleetAlgoKnobFaultTolerance: reconcile_timeout_s and "
                "coordinator_takeover_s must be positive, got "
                f"{self.reconcile_timeout_s} / {self.coordinator_takeover_s}"
            )


@dataclass(frozen=True)
class FleetAlgoKnobAllocator(AlgoKnob):
    """Route NCCL-EP device buffers through a custom allocator.

    Without this knob the group uses NCCL-EP's default ``cudaMalloc`` /
    ``cudaFree``. Two mutually-exclusive modes:

    * ``torch_caching=True`` — FlashInfer installs allocator trampolines
      backed by torch's CUDA caching allocator, so EP transport buffers draw
      from the same pool as the rest of the process (e.g. vLLM's memory
      budget) instead of a second, invisible cudaMalloc arena. FlashInfer
      owns the trampoline lifetime (anchored on the Fleet).
    * ``alloc_fn`` / ``free_fn`` addresses — use these C-callable addresses
      directly (``ncclEpAllocFn_t`` / ``ncclEpFreeFn_t``; see
      ``nccl.ep.allocator``). The CALLER owns their lifetime and MUST keep
      them alive at least as long as the Fleet.
    """

    torch_caching: bool = False
    alloc_fn: int = 0
    free_fn: int = 0
    context: int = 0

    def __post_init__(self) -> None:
        if self.torch_caching and (self.alloc_fn or self.free_fn):
            raise ValueError(
                "FleetAlgoKnobAllocator: set either torch_caching=True OR "
                "explicit alloc_fn/free_fn addresses, not both."
            )
        if (self.alloc_fn == 0) != (self.free_fn == 0):
            raise ValueError(
                "FleetAlgoKnobAllocator: alloc_fn and free_fn must be provided "
                "together (both non-zero) or both left 0."
            )


# --------------------------------------------------------------- Handle-level


@dataclass(frozen=True)
class HandleAlgoKnobUserStream(AlgoKnob):
    """Run the dispatch/combine on a non-default CUDA stream.

    The stream is passed as an int (cudaStream_t) since dataclasses with
    torch.cuda.Stream fields aren't hashable.
    """

    stream: int


@dataclass(frozen=True)
class HandleAlgoKnobSplitOperation(AlgoKnob):
    """Marker: enable send_only=1 + ncclEpComplete staging."""


@dataclass(frozen=True)
class HandleAlgoKnobTopKWeights(AlgoKnob):
    """Carries the topk_weights tensor for combine to reweight by."""

    weights: "torch.Tensor"


@dataclass(frozen=True)
class HandleAlgoKnobNumReceivedTokens(AlgoKnob):
    """Where to write the recv-count from ncclEpHandleGetNumRecvTokens."""

    target: "torch.Tensor"


def _index_knobs(knobs: Sequence[AlgoKnob]) -> "dict[type, AlgoKnob]":
    """Build a dict keyed by knob class so backends can ``.get(KnobClass)``.

    If the same knob class appears twice in the sequence, the later entry
    wins (mirrors dict() over key-value pairs). Knobs with class identity
    only (e.g. SplitOperation) act as flags: their presence in the dict
    signals "set"; the stored instance is the marker itself.
    """
    out: dict[type, AlgoKnob] = {}
    for k in knobs:
        if not isinstance(k, AlgoKnob):
            raise TypeError(f"expected AlgoKnob, got {type(k).__name__}")
        out[type(k)] = k
    return out
