from __future__ import annotations

from collections.abc import Iterable

from ..knob_descriptor import KnobConstraints, KnobDescriptor


def default_vllm_descriptors(*, tunable_flags: Iterable[str]) -> tuple[KnobDescriptor, ...]:
    tunable = set(tunable_flags)
    descriptors: list[KnobDescriptor] = []
    if "max-num-seqs" in tunable:
        descriptors.append(
            KnobDescriptor(
                name="max-num-seqs",
                location="flag",
                value_type="int",
                mechanisms=("admission_capacity", "kv_memory_pressure"),
                search_geometry="positive_capacity",
                operators=("coordinate_line_search", "frontier_delta_projection"),
                constraints=KnobConstraints(min_value=1, integer=True, multiple_of=8),
                directional_effects={
                    "increase": ("admission_capacity",),
                    "decrease": ("kv_memory_pressure",),
                },
                risk_effects={
                    "increase": ("kv_memory_pressure", "decode_tail_latency"),
                },
            )
        )
    if "max-num-batched-tokens" in tunable:
        descriptors.append(
            KnobDescriptor(
                name="max-num-batched-tokens",
                location="flag",
                value_type="int",
                mechanisms=("prefill_scheduling", "decode_batching"),
                search_geometry="positive_capacity",
                operators=("coordinate_line_search", "frontier_delta_projection"),
                constraints=KnobConstraints(min_value=1, integer=True, multiple_of=128),
                directional_effects={
                    "increase": ("prefill_scheduling", "decode_batching"),
                    "decrease": ("prefill_tail_latency",),
                },
                risk_effects={
                    "increase": ("prefill_tail_latency", "kv_memory_pressure"),
                },
            )
        )
    if "gpu-memory-utilization" in tunable:
        descriptors.append(
            KnobDescriptor(
                name="gpu-memory-utilization",
                location="flag",
                value_type="float",
                mechanisms=("kv_memory_capacity", "launch_feasibility"),
                search_geometry="bounded_fraction",
                operators=("coordinate_line_search", "frontier_delta_projection"),
                constraints=KnobConstraints(min_value=0.0, max_value=0.97),
                directional_effects={
                    "increase": ("kv_memory_capacity",),
                    "decrease": ("launch_feasibility",),
                },
                risk_effects={
                    "increase": ("launch_feasibility",),
                },
            )
        )
    if "enable-chunked-prefill" in tunable:
        descriptors.append(
            KnobDescriptor(
                name="enable-chunked-prefill",
                location="flag",
                value_type="bool",
                mechanisms=("prefill_scheduling",),
                search_geometry="toggle",
                operators=("coordinate_line_search",),
                directional_effects={
                    "toggle": ("prefill_scheduling",),
                },
            )
        )
    return tuple(descriptors)
