from __future__ import annotations

import argparse
import asyncio
from pathlib import Path

from agentic_pd_hybrid.benchmark import BenchmarkConfig, run_live_benchmark
from agentic_pd_hybrid.launcher import build_launch_plan
from agentic_pd_hybrid.microbench import SmallAppendTraceConfig, write_small_append_trace
from agentic_pd_hybrid.profile import ProfileConfig, print_profile_summary, write_profile
from agentic_pd_hybrid.replay import ReplayConfig, replay_trace
from agentic_pd_hybrid.sampling import SessionSampleConfig, sample_trace_sessions
from agentic_pd_hybrid.trace_profiles import (
    NormalizeTraceLengthsConfig,
    normalize_trace_lengths,
)
from agentic_pd_hybrid.topology import build_single_node_topology


def _normalize_mechanism_name(name: str) -> str:
    normalized = name.strip().lower()
    aliases = {
        "pd-disagg": "pd-disaggregation",
        "pd-disaggregation": "pd-disaggregation",
        "pd-hybrid": "pd-disaggregation",
        "baseline-pd-disagg": "pd-disaggregation",
        "pd-colo": "pd-colo",
        "direct-d": "pd-colo",
        "colocation": "pd-colo",
        "kvcache-centric": "kvcache-centric",
        "turn2+-direct-to-d": "kvcache-centric",
        "pd-with-d-append": "kvcache-centric",
    }
    if normalized not in aliases:
        raise ValueError(f"Unsupported mechanism: {name}")
    return aliases[normalized]


def _parse_gpu_id_list(value: str | None) -> tuple[int, ...] | None:
    if value is None:
        return None
    items = [item.strip() for item in value.split(",") if item.strip()]
    if not items:
        return tuple()
    return tuple(int(item) for item in items)


def main() -> None:
    parser = argparse.ArgumentParser(description="Agentic PD hybrid prototype")
    subparsers = parser.add_subparsers(dest="command", required=True)

    print_launch = subparsers.add_parser(
        "print-launch",
        help="Print one-node SGLang PD launch commands",
    )
    _add_topology_arguments(print_launch)
    print_launch.add_argument("--prefill-policy", default="round_robin")
    print_launch.add_argument("--decode-policy", default="manual")

    replay = subparsers.add_parser(
        "replay",
        help="Replay trace and log request-level metrics",
    )
    _add_topology_arguments(replay)
    replay.add_argument("--trace", type=Path, required=True)
    replay.add_argument("--output", type=Path, required=True)
    replay.add_argument(
        "--policy",
        choices=["default", "sticky", "kv-aware"],
        default="sticky",
    )
    replay.add_argument(
        "--mechanism",
        choices=[
            "pd-disaggregation",
            "pd-hybrid",
            "pd-disagg",
            "pd-colo",
            "direct-d",
            "kvcache-centric",
            "turn2+-direct-to-d",
            "pd-with-d-append",
        ],
        default="pd-disaggregation",
    )
    replay.add_argument("--router-url")
    replay.add_argument("--model")
    replay.add_argument(
        "--header-mode",
        choices=["auto", "none", "routing-key", "target-worker"],
        default="auto",
    )
    replay.add_argument(
        "--request-limit",
        type=int,
        default=None,
        help="Replay at most this many requests",
    )
    replay.add_argument(
        "--no-pace",
        action="store_true",
        help="Disable wall-clock pacing from trace timestamps",
    )
    replay.add_argument(
        "--time-scale",
        type=float,
        default=1.0,
        help="Scale trace timing by this factor when pacing is enabled",
    )
    replay.add_argument(
        "--concurrency-limit",
        type=int,
        default=32,
    )
    replay.add_argument(
        "--timeout-s",
        type=float,
        default=600.0,
    )
    replay.add_argument(
        "--no-stream",
        action="store_true",
        help="Use non-streaming OpenAI responses for more robust E2E-only replay.",
    )
    replay.add_argument(
        "--stream-idle-timeout-s",
        type=float,
        default=900.0,
        help="Abort a streaming request if no SSE line arrives within this many seconds.",
    )
    replay.add_argument(
        "--kvcache-direct-max-uncached-tokens",
        type=int,
        default=2048,
        help="For kvcache-centric routing, bypass P when the uncached suffix is at most this many tokens.",
    )
    replay.add_argument(
        "--kvcache-admission-mode",
        choices=["router", "worker"],
        default="router",
        help=(
            "For kvcache-centric routing, use router shadow-state admission "
            "or query the decode worker on the critical path."
        ),
    )
    replay.add_argument(
        "--kvcache-seed-max-resident-tokens",
        type=int,
        default=None,
        help=(
            "For kvcache-centric routing, do not seed/reseed a decode session "
            "when input+output tokens exceed this value."
        ),
    )
    replay.add_argument(
        "--kvcache-seed-max-output-tokens",
        type=int,
        default=None,
        help=(
            "For kvcache-centric routing, do not seed/reseed a decode session "
            "when output tokens exceed this value."
        ),
    )
    replay.add_argument(
        "--kvcache-seed-min-turn-id",
        type=int,
        default=1,
        help=(
            "For kvcache-centric routing, do not seed/reseed a decode session "
            "before this turn id."
        ),
    )
    replay.add_argument(
        "--kvcache-seed-only-multiturn-sessions",
        action="store_true",
        help=(
            "Oracle ablation for kvcache-centric routing: only seed sessions "
            "that have more than one turn in the replay trace."
        ),
    )
    replay.add_argument(
        "--kvcache-prefill-backup-policy",
        choices=["release-after-transfer", "capacity-backup"],
        default="release-after-transfer",
        help=(
            "For kvcache-centric seed/reseed, release the P-side session after "
            "P->D transfer or keep a capacity-limited P-side backup."
        ),
    )
    replay.add_argument(
        "--kvcache-seed-max-inflight-decode",
        type=int,
        default=3,
        help=(
            "For kvcache-centric routing, skip seed/reseed when the router "
            "shadow inflight decode load assigned to the target D exceeds this value. "
            "Use a negative value to disable this filter."
        ),
    )
    replay.add_argument(
        "--kvcache-seed-max-decode-transfer-queue-reqs",
        type=int,
        default=None,
        help=(
            "For kvcache-centric worker admission, skip seed/reseed when the "
            "target D reports more transfer-queue requests than this value. "
            "Use a negative value to disable this filter."
        ),
    )
    replay.add_argument(
        "--kvcache-direct-max-decode-transfer-queue-reqs",
        type=int,
        default=None,
        help=(
            "For kvcache-centric worker admission, skip direct-to-D append when "
            "the target D reports more transfer-queue requests than this value. "
            "Use a negative value to disable this filter."
        ),
    )
    replay.add_argument(
        "--kvcache-prefill-priority-eviction",
        action="store_true",
        help=(
            "For kvcache-centric routing, mark P-side prefixes predicted to move "
            "direct-to-D as lower priority. Requires P workers to use "
            "--radix-eviction-policy priority."
        ),
    )
    replay.add_argument("--kvcache-prefill-direct-priority", type=int, default=-100)
    replay.add_argument("--kvcache-prefill-normal-priority", type=int, default=100)
    replay.add_argument(
        "--pool-poll-interval-s",
        type=float,
        default=0.0,
        help=(
            "Poll each P/D worker's /server_info every N seconds and write a "
            "time-series snapshot to <run_dir>/d-pool-timeseries.jsonl. "
            "0 disables polling."
        ),
    )
    replay.add_argument(
        "--pool-poll-no-sessions",
        action="store_true",
        help=(
            "Disable per-session detail in the pool timeseries (smaller files)."
        ),
    )
    replay.add_argument(
        "--enable-backpressure",
        action="store_true",
        help=(
            "Honor recommended_pause_ms hints from D's admission endpoint. "
            "When set, replay sleeps before issuing requests to a saturated D. "
            "Default off — preserves baseline behavior."
        ),
    )
    replay.add_argument(
        "--backpressure-max-pause-s",
        type=float,
        default=2.0,
        help="Cap on per-request backpressure sleep, regardless of D hint.",
    )
    replay.add_argument(
        "--kvcache-migration-reject-threshold",
        type=int,
        default=3,
        help=(
            "Per-(session, D) admission-reject count after which KvAwarePolicy "
            "skips that D for the session (forces migration). 0 disables. "
            "See REFACTOR_PLAN_V1 §6.2 / TEAM_REPORT §2.1."
        ),
    )

    sample = subparsers.add_parser(
        "sample-sessions",
        help="Sample a session-granularity trace shard for live benchmarking",
    )
    sample.add_argument("--trace", type=Path, required=True)
    sample.add_argument("--output", type=Path, required=True)
    sample.add_argument("--target-duration-s", type=float, default=600.0)
    sample.add_argument("--start-time-s", type=float, default=0.0)
    sample.add_argument("--session-sample-rate", type=float, default=0.01)
    sample.add_argument("--min-turns", type=int, default=1)
    sample.add_argument("--max-requests", type=int, default=None)
    sample.add_argument(
        "--profile",
        choices=["default", "small-append"],
        default="default",
        help="Optional workload-shape filter for live benchmarks.",
    )
    sample.add_argument("--min-initial-input-tokens", type=int, default=None)
    sample.add_argument("--max-initial-input-tokens", type=int, default=None)
    sample.add_argument("--max-append-input-tokens", type=int, default=None)
    sample.add_argument("--max-output-tokens", type=int, default=None)
    sample.add_argument("--min-overlap-ratio", type=float, default=None)

    normalize = subparsers.add_parser(
        "normalize-trace-lengths",
        help="Rewrite a trace to a fixed turn1/append/output length profile",
    )
    normalize.add_argument("--trace", type=Path, required=True)
    normalize.add_argument("--output", type=Path, required=True)
    normalize.add_argument("--initial-input-length", type=int, default=10_000)
    normalize.add_argument("--append-input-length", type=int, default=1_000)
    normalize.add_argument("--output-length", type=int, default=1_000)
    normalize.add_argument("--max-requests", type=int, default=None)

    profile = subparsers.add_parser(
        "profile",
        help="Profile a trace and optional request metrics for routing analysis",
    )
    profile.add_argument("--trace", type=Path, required=True)
    profile.add_argument("--output", type=Path, default=None)
    profile.add_argument("--metrics", type=Path, default=None)
    profile.add_argument("--baseline-metrics", type=Path, default=None)
    profile.add_argument("--candidate-metrics", type=Path, default=None)
    profile.add_argument("--direct-max-uncached-tokens", type=int, default=2048)

    micro = subparsers.add_parser(
        "make-small-append-trace",
        help="Generate a synthetic multi-turn trace with small turn2+ appends",
    )
    micro.add_argument("--output", type=Path, required=True)
    micro.add_argument("--session-count", type=int, default=8)
    micro.add_argument("--turns-per-session", type=int, default=3)
    micro.add_argument("--initial-input-length", type=int, default=10_000)
    micro.add_argument("--append-input-length", type=int, default=1_000)
    micro.add_argument("--output-length", type=int, default=1_000)
    micro.add_argument("--inter-turn-gap-s", type=float, default=1.0)
    micro.add_argument("--session-stagger-s", type=float, default=0.1)

    benchmark = subparsers.add_parser(
        "benchmark-live",
        help="Launch a real PD stack, sample sessions, and collect live E2E numbers",
    )
    _add_topology_arguments(benchmark)
    benchmark.add_argument("--trace", type=Path, required=True)
    benchmark.add_argument(
        "--policy",
        choices=["default", "sticky", "kv-aware"],
        default="sticky",
    )
    benchmark.add_argument(
        "--mechanism",
        choices=[
            "pd-disaggregation",
            "pd-hybrid",
            "pd-disagg",
            "pd-colo",
            "direct-d",
            "kvcache-centric",
            "turn2+-direct-to-d",
            "pd-with-d-append",
        ],
        default="pd-disaggregation",
    )
    benchmark.add_argument("--output-root", type=Path, default=Path("outputs/live"))
    benchmark.add_argument("--target-duration-s", type=float, default=600.0)
    benchmark.add_argument("--start-time-s", type=float, default=0.0)
    benchmark.add_argument("--session-sample-rate", type=float, default=0.01)
    benchmark.add_argument("--min-turns", type=int, default=1)
    benchmark.add_argument("--time-scale", type=float, default=1.0)
    benchmark.add_argument("--concurrency-limit", type=int, default=32)
    benchmark.add_argument("--timeout-s", type=float, default=1200.0)
    benchmark.add_argument(
        "--request-timeout-s",
        type=float,
        default=None,
        help=(
            "Per-request replay/router timeout. If unset, --timeout-s is used. "
            "--timeout-s still controls stack startup."
        ),
    )
    benchmark.add_argument(
        "--no-stream",
        action="store_true",
        help="Use non-streaming OpenAI responses for E2E-only live benchmarking.",
    )
    benchmark.add_argument(
        "--stream-idle-timeout-s",
        type=float,
        default=900.0,
        help="Abort a streaming request if no SSE line arrives within this many seconds.",
    )
    benchmark.add_argument(
        "--kvcache-direct-max-uncached-tokens",
        type=int,
        default=2048,
        help="For kvcache-centric routing, bypass P when the uncached suffix is at most this many tokens.",
    )
    benchmark.add_argument(
        "--kvcache-admission-mode",
        choices=["router", "worker"],
        default="router",
        help=(
            "For kvcache-centric routing, use router shadow-state admission "
            "or query the decode worker on the critical path."
        ),
    )
    benchmark.add_argument(
        "--kvcache-seed-max-resident-tokens",
        type=int,
        default=None,
        help=(
            "For kvcache-centric routing, do not seed/reseed a decode session "
            "when input+output tokens exceed this value."
        ),
    )
    benchmark.add_argument(
        "--kvcache-seed-max-output-tokens",
        type=int,
        default=None,
        help=(
            "For kvcache-centric routing, do not seed/reseed a decode session "
            "when output tokens exceed this value."
        ),
    )
    benchmark.add_argument(
        "--kvcache-seed-min-turn-id",
        type=int,
        default=1,
        help=(
            "For kvcache-centric routing, do not seed/reseed a decode session "
            "before this turn id."
        ),
    )
    benchmark.add_argument(
        "--kvcache-seed-only-multiturn-sessions",
        action="store_true",
        help=(
            "Oracle ablation for kvcache-centric routing: only seed sessions "
            "that have more than one turn in the replay trace."
        ),
    )
    benchmark.add_argument(
        "--kvcache-prefill-backup-policy",
        choices=["release-after-transfer", "capacity-backup"],
        default="release-after-transfer",
        help=(
            "For kvcache-centric seed/reseed, release the P-side session after "
            "P->D transfer or keep a capacity-limited P-side backup."
        ),
    )
    benchmark.add_argument(
        "--kvcache-seed-max-inflight-decode",
        type=int,
        default=3,
        help=(
            "For kvcache-centric routing, skip seed/reseed when the router "
            "shadow inflight decode load assigned to the target D exceeds this value. "
            "Use a negative value to disable this filter."
        ),
    )
    benchmark.add_argument(
        "--kvcache-seed-max-decode-transfer-queue-reqs",
        type=int,
        default=None,
        help=(
            "For kvcache-centric worker admission, skip seed/reseed when the "
            "target D reports more transfer-queue requests than this value. "
            "Use a negative value to disable this filter."
        ),
    )
    benchmark.add_argument(
        "--kvcache-direct-max-decode-transfer-queue-reqs",
        type=int,
        default=None,
        help=(
            "For kvcache-centric worker admission, skip direct-to-D append when "
            "the target D reports more transfer-queue requests than this value. "
            "Use a negative value to disable this filter."
        ),
    )
    benchmark.add_argument(
        "--kvcache-prefill-priority-eviction",
        action="store_true",
        help=(
            "For kvcache-centric benchmark-live, launch P workers with priority "
            "radix eviction and mark direct-to-D predicted prefixes lower priority."
        ),
    )
    benchmark.add_argument("--kvcache-prefill-direct-priority", type=int, default=-100)
    benchmark.add_argument("--kvcache-prefill-normal-priority", type=int, default=100)
    benchmark.add_argument(
        "--pool-poll-interval-s",
        type=float,
        default=0.0,
        help=(
            "Poll each P/D worker's /server_info every N seconds and write a "
            "time-series snapshot to <run_dir>/d-pool-timeseries.jsonl. "
            "0 disables polling."
        ),
    )
    benchmark.add_argument(
        "--pool-poll-no-sessions",
        action="store_true",
        help=(
            "Disable per-session detail in the pool timeseries (smaller files)."
        ),
    )
    benchmark.add_argument(
        "--enable-backpressure",
        action="store_true",
        help=(
            "Honor recommended_pause_ms hints from D's admission endpoint."
        ),
    )
    benchmark.add_argument(
        "--backpressure-max-pause-s",
        type=float,
        default=2.0,
        help="Cap on per-request backpressure sleep, regardless of D hint.",
    )
    benchmark.add_argument(
        "--kvcache-migration-reject-threshold",
        type=int,
        default=3,
        help=(
            "Per-(session, D) admission-reject count after which KvAwarePolicy "
            "skips that D for the session (forces migration). 0 disables. "
            "See REFACTOR_PLAN_V1 §6.2 / TEAM_REPORT §2.1."
        ),
    )
    benchmark.add_argument(
        "--sample-profile",
        choices=["default", "small-append"],
        default="default",
        help="Optional session-shape filter applied before live replay.",
    )
    benchmark.add_argument("--min-initial-input-tokens", type=int, default=None)
    benchmark.add_argument("--max-initial-input-tokens", type=int, default=None)
    benchmark.add_argument("--max-append-input-tokens", type=int, default=None)
    benchmark.add_argument("--max-output-tokens", type=int, default=None)
    benchmark.add_argument("--min-overlap-ratio", type=float, default=None)

    args = parser.parse_args()

    if args.command == "print-launch":
        topology = _topology_from_args(args)
        has_pd = bool(topology.prefill_workers and topology.decode_workers)
        has_direct_only = bool(
            topology.direct_workers
            and not topology.prefill_workers
            and not topology.decode_workers
        )
        plan = build_launch_plan(
            topology,
            prefill_policy=args.prefill_policy,
            decode_policy=args.decode_policy,
            include_router=has_pd or has_direct_only,
            naive_dp=has_direct_only,
        )
        print(plan.render())
        return

    if args.command == "replay":
        topology = _topology_from_args(args)
        config = ReplayConfig(
            trace_path=args.trace,
            output_path=args.output,
            policy_name=args.policy,
            mechanism_name=_normalize_mechanism_name(args.mechanism),
            topology=topology,
            router_url=args.router_url,
            model_name=args.model,
            pace=not args.no_pace,
            time_scale=args.time_scale,
            request_limit=args.request_limit,
            concurrency_limit=args.concurrency_limit,
            header_mode=args.header_mode,
            timeout_s=args.timeout_s,
            stream=not args.no_stream,
            stream_idle_timeout_s=args.stream_idle_timeout_s,
            kvcache_direct_max_uncached_tokens=args.kvcache_direct_max_uncached_tokens,
            kvcache_admission_mode=args.kvcache_admission_mode,
            kvcache_seed_max_resident_tokens=args.kvcache_seed_max_resident_tokens,
            kvcache_seed_max_output_tokens=args.kvcache_seed_max_output_tokens,
            kvcache_seed_min_turn_id=args.kvcache_seed_min_turn_id,
            kvcache_seed_only_multiturn_sessions=(
                args.kvcache_seed_only_multiturn_sessions
            ),
            kvcache_prefill_backup_policy=args.kvcache_prefill_backup_policy,
            kvcache_seed_max_inflight_decode=(
                None
                if args.kvcache_seed_max_inflight_decode < 0
                else args.kvcache_seed_max_inflight_decode
            ),
            kvcache_seed_max_decode_transfer_queue_reqs=(
                None
                if args.kvcache_seed_max_decode_transfer_queue_reqs is None
                or args.kvcache_seed_max_decode_transfer_queue_reqs < 0
                else args.kvcache_seed_max_decode_transfer_queue_reqs
            ),
            kvcache_direct_max_decode_transfer_queue_reqs=(
                None
                if args.kvcache_direct_max_decode_transfer_queue_reqs is None
                or args.kvcache_direct_max_decode_transfer_queue_reqs < 0
                else args.kvcache_direct_max_decode_transfer_queue_reqs
            ),
            kvcache_prefill_priority_eviction=(
                args.kvcache_prefill_priority_eviction
            ),
            kvcache_prefill_direct_priority=args.kvcache_prefill_direct_priority,
            kvcache_prefill_normal_priority=args.kvcache_prefill_normal_priority,
            pool_poll_interval_s=args.pool_poll_interval_s,
            pool_poll_include_sessions=not args.pool_poll_no_sessions,
            enable_backpressure=args.enable_backpressure,
            backpressure_max_pause_s=args.backpressure_max_pause_s,
            kvcache_migration_reject_threshold=args.kvcache_migration_reject_threshold,
        )
        results = asyncio.run(replay_trace(config))
        print(
            f"wrote {len(results)} request records to {args.output} and "
            f"{args.output}{'.summary.json'}"
        )
        return

    if args.command == "profile":
        if (args.baseline_metrics is None) != (args.candidate_metrics is None):
            raise ValueError(
                "--baseline-metrics and --candidate-metrics must be provided together"
            )
        report = write_profile(
            ProfileConfig(
                trace_path=args.trace,
                output_path=args.output,
                metrics_path=args.metrics,
                baseline_metrics_path=args.baseline_metrics,
                candidate_metrics_path=args.candidate_metrics,
                direct_max_uncached_tokens=args.direct_max_uncached_tokens,
            )
        )
        print_profile_summary(report)
        if args.output is not None:
            print(f"wrote profile to {args.output}")
        return

    if args.command == "sample-sessions":
        summary = sample_trace_sessions(
            SessionSampleConfig(
                trace_path=args.trace,
                output_path=args.output,
                target_duration_s=args.target_duration_s,
                start_time_s=args.start_time_s,
                session_sample_rate=args.session_sample_rate,
                min_turns=args.min_turns,
                max_requests=args.max_requests,
                profile=args.profile,
                min_initial_input_tokens=args.min_initial_input_tokens,
                max_initial_input_tokens=args.max_initial_input_tokens,
                max_append_input_tokens=args.max_append_input_tokens,
                max_output_tokens=args.max_output_tokens,
                min_overlap_ratio=args.min_overlap_ratio,
            )
        )
        print(
            f"wrote {summary.request_count} requests from {summary.session_count} sessions "
            f"covering {summary.sampled_duration_s:.3f}s to {args.output}"
        )
        return

    if args.command == "normalize-trace-lengths":
        summary = normalize_trace_lengths(
            NormalizeTraceLengthsConfig(
                trace_path=args.trace,
                output_path=args.output,
                initial_input_length=args.initial_input_length,
                append_input_length=args.append_input_length,
                output_length=args.output_length,
                max_requests=args.max_requests,
            )
        )
        print(
            f"wrote {summary.request_count} normalized requests from "
            f"{summary.session_count} sessions to {args.output}"
        )
        return

    if args.command == "make-small-append-trace":
        summary = write_small_append_trace(
            SmallAppendTraceConfig(
                output_path=args.output,
                session_count=args.session_count,
                turns_per_session=args.turns_per_session,
                initial_input_length=args.initial_input_length,
                append_input_length=args.append_input_length,
                output_length=args.output_length,
                inter_turn_gap_s=args.inter_turn_gap_s,
                session_stagger_s=args.session_stagger_s,
            )
        )
        print(
            f"wrote {summary.request_count} requests across {summary.session_count} sessions "
            f"to {args.output}"
        )
        return

    if args.command == "benchmark-live":
        topology = _topology_from_args(args)
        artifacts = run_live_benchmark(
            BenchmarkConfig(
                trace_path=args.trace,
                output_root=args.output_root,
                topology=topology,
                policy_name=args.policy,
                mechanism_name=_normalize_mechanism_name(args.mechanism),
                target_duration_s=args.target_duration_s,
                start_time_s=args.start_time_s,
                session_sample_rate=args.session_sample_rate,
                min_turns=args.min_turns,
                time_scale=args.time_scale,
                concurrency_limit=args.concurrency_limit,
                timeout_s=args.timeout_s,
                request_timeout_s=args.request_timeout_s,
                stream=not args.no_stream,
                stream_idle_timeout_s=args.stream_idle_timeout_s,
                kvcache_direct_max_uncached_tokens=args.kvcache_direct_max_uncached_tokens,
                kvcache_admission_mode=args.kvcache_admission_mode,
                kvcache_seed_max_resident_tokens=args.kvcache_seed_max_resident_tokens,
                kvcache_seed_max_output_tokens=args.kvcache_seed_max_output_tokens,
                kvcache_seed_min_turn_id=args.kvcache_seed_min_turn_id,
                kvcache_seed_only_multiturn_sessions=(
                    args.kvcache_seed_only_multiturn_sessions
                ),
                kvcache_prefill_backup_policy=args.kvcache_prefill_backup_policy,
                kvcache_seed_max_inflight_decode=(
                    None
                    if args.kvcache_seed_max_inflight_decode < 0
                    else args.kvcache_seed_max_inflight_decode
                ),
                kvcache_seed_max_decode_transfer_queue_reqs=(
                    None
                    if args.kvcache_seed_max_decode_transfer_queue_reqs is None
                    or args.kvcache_seed_max_decode_transfer_queue_reqs < 0
                    else args.kvcache_seed_max_decode_transfer_queue_reqs
                ),
                kvcache_direct_max_decode_transfer_queue_reqs=(
                    None
                    if args.kvcache_direct_max_decode_transfer_queue_reqs is None
                    or args.kvcache_direct_max_decode_transfer_queue_reqs < 0
                    else args.kvcache_direct_max_decode_transfer_queue_reqs
                ),
                kvcache_prefill_priority_eviction=(
                    args.kvcache_prefill_priority_eviction
                ),
                kvcache_prefill_direct_priority=(
                    args.kvcache_prefill_direct_priority
                ),
                kvcache_prefill_normal_priority=(
                    args.kvcache_prefill_normal_priority
                ),
                pool_poll_interval_s=args.pool_poll_interval_s,
                pool_poll_include_sessions=not args.pool_poll_no_sessions,
                enable_backpressure=args.enable_backpressure,
                backpressure_max_pause_s=args.backpressure_max_pause_s,
                kvcache_migration_reject_threshold=args.kvcache_migration_reject_threshold,
                sample_profile=args.sample_profile,
                min_initial_input_tokens=args.min_initial_input_tokens,
                max_initial_input_tokens=args.max_initial_input_tokens,
                max_append_input_tokens=args.max_append_input_tokens,
                max_output_tokens=args.max_output_tokens,
                min_overlap_ratio=args.min_overlap_ratio,
                launch_stack=True,
            )
        )
        print(
            f"benchmark artifacts written under {artifacts.run_dir}; "
            f"metrics={artifacts.metrics_path} summary={artifacts.summary_path}"
        )
        return

    raise AssertionError(f"Unhandled command: {args.command}")


def _add_topology_arguments(parser: argparse.ArgumentParser) -> None:
    parser.add_argument(
        "--model-path",
        default="~/models/Qwen/Qwen3-Coder-30B-A3B-Instruct",
    )
    parser.add_argument("--prefill-workers", type=int, default=1)
    parser.add_argument("--decode-workers", type=int, default=1)
    parser.add_argument("--direct-workers", type=int, default=0)
    parser.add_argument("--prefill-tp-size", type=int, default=1)
    parser.add_argument("--decode-tp-size", type=int, default=1)
    parser.add_argument("--direct-tp-size", type=int, default=1)
    parser.add_argument("--gpu-budget", type=int, default=8)
    parser.add_argument(
        "--prefill-gpu-ids",
        default=None,
        help="Comma-separated GPU IDs for prefill workers, e.g. 3,4",
    )
    parser.add_argument(
        "--decode-gpu-ids",
        default=None,
        help="Comma-separated GPU IDs for decode workers, e.g. 5",
    )
    parser.add_argument(
        "--direct-gpu-ids",
        default=None,
        help="Comma-separated GPU IDs for direct workers, e.g. 6",
    )
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--router-port", type=int, default=8000)
    parser.add_argument("--prefill-port-base", type=int, default=30000)
    parser.add_argument("--decode-port-base", type=int, default=31000)
    parser.add_argument("--direct-port-base", type=int, default=32000)
    parser.add_argument("--bootstrap-port-base", type=int, default=8998)
    parser.add_argument("--transfer-backend", default="nixl")
    parser.add_argument(
        "--force-rdma",
        action="store_true",
        help=(
            "Force real RDMA transport for PD KV transfer. "
            "Currently this requires Mooncake plus --ib-device."
        ),
    )
    parser.add_argument("--ib-device", default=None)
    parser.add_argument(
        "--no-trust-remote-code",
        action="store_true",
    )


def _topology_from_args(args: argparse.Namespace):
    transfer_backend = args.transfer_backend
    if args.force_rdma:
        transfer_backend = "mooncake"

    return build_single_node_topology(
        model_path=str(Path(args.model_path).expanduser()),
        prefill_worker_count=args.prefill_workers,
        decode_worker_count=args.decode_workers,
        direct_worker_count=args.direct_workers,
        prefill_tp_size=args.prefill_tp_size,
        decode_tp_size=args.decode_tp_size,
        direct_tp_size=args.direct_tp_size,
        prefill_gpu_ids=_parse_gpu_id_list(args.prefill_gpu_ids),
        decode_gpu_ids=_parse_gpu_id_list(args.decode_gpu_ids),
        direct_gpu_ids=_parse_gpu_id_list(args.direct_gpu_ids),
        total_gpu_budget=args.gpu_budget,
        host=args.host,
        router_port=args.router_port,
        prefill_port_base=args.prefill_port_base,
        decode_port_base=args.decode_port_base,
        direct_port_base=args.direct_port_base,
        bootstrap_port_base=args.bootstrap_port_base,
        transfer_backend=transfer_backend,
        force_rdma=args.force_rdma,
        trust_remote_code=not args.no_trust_remote_code,
        ib_device=args.ib_device,
        direct_extra_server_args=("--enable-streaming-session",),
    )


if __name__ == "__main__":
    main()
