#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""
CLI for vllm-mlx.

Commands:
    vllm-mlx serve <model> --port 8000    Start OpenAI-compatible server
    vllm-mlx bench <model>                Run benchmark

Usage:
    vllm-mlx serve mlx-community/Llama-3.2-3B-Instruct-4bit --port 8000
    vllm-mlx bench mlx-community/Llama-3.2-1B-Instruct-4bit --num-prompts 10
"""

import argparse
import json
import sys

from .cli_arg_types import make_json_object_arg_parser, make_positive_int_arg_parser


def serve_command(args):
    """Start the OpenAI-compatible server."""
    import logging
    import os
    import sys

    import uvicorn

    # Import unified server
    from . import server
    from .model_registry import RegistryServeDefaults
    from .server import RateLimiter, app, load_model, load_model_registry

    logger = logging.getLogger(__name__)
    model_arg = getattr(args, "model", None)
    models_config = getattr(args, "models_config", None)

    if models_config and model_arg:
        print("Error: use either positional MODEL or --models-config, not both")
        sys.exit(1)
    if not models_config and not model_arg:
        print("Error: MODEL is required unless --models-config is provided")
        sys.exit(1)
    if models_config and args.served_model_name:
        print("Error: --served-model-name cannot be used with --models-config")
        sys.exit(1)

    # Validate tool calling arguments
    if args.enable_auto_tool_choice and not args.tool_call_parser:
        print("Error: --enable-auto-tool-choice requires --tool-call-parser")
        print("Example: --enable-auto-tool-choice --tool-call-parser mistral")
        sys.exit(1)

    # Validate gpu-memory-utilization range
    if not (0.0 < args.gpu_memory_utilization <= 1.0):
        print(
            "Error: --gpu-memory-utilization must be between 0.0 (exclusive) and 1.0 (inclusive)"
        )
        sys.exit(1)
    if args.max_tokens < 1:
        print("Error: --max-tokens must be at least 1")
        sys.exit(1)
    max_request_tokens = getattr(args, "max_request_tokens", args.max_tokens)
    max_kv_size = getattr(args, "max_kv_size", None)
    trust_remote_code = getattr(args, "trust_remote_code", False)
    if max_request_tokens < 1:
        print("Error: --max-request-tokens must be at least 1")
        sys.exit(1)
    if args.max_tokens > max_request_tokens:
        print("Error: --max-tokens cannot exceed --max-request-tokens")
        sys.exit(1)
    mllm_draft_model = getattr(args, "mllm_draft_model", None)
    mllm_draft_kind = getattr(args, "mllm_draft_kind", None)
    mllm_draft_block_size = getattr(args, "mllm_draft_block_size", None)
    if mllm_draft_model and models_config:
        print("Error: --mllm-draft-model cannot be used with --models-config")
        sys.exit(1)
    if mllm_draft_model and not getattr(args, "mllm", False):
        print("Error: --mllm-draft-model requires --mllm")
        sys.exit(1)
    if mllm_draft_block_size is not None and mllm_draft_block_size <= 0:
        print("Error: --mllm-draft-block-size must be a positive integer")
        sys.exit(1)
    if mllm_draft_model and args.continuous_batching:
        print(
            "Error: --mllm-draft-model is supported only without --continuous-batching"
        )
        sys.exit(1)
    if mllm_draft_model and (args.auto_unload_idle_seconds > 0 or args.lazy_load_model):
        print("Error: --mllm-draft-model is not supported with lifecycle residency yet")
        sys.exit(1)

    # Configure server security settings
    server._api_key = args.api_key
    server._default_timeout = args.timeout
    server._metrics_enabled = args.enable_metrics
    server._metrics.configure(enabled=args.enable_metrics)
    server._max_request_tokens = max_request_tokens
    if args.rate_limit > 0:
        server._rate_limiter = RateLimiter(
            requests_per_minute=args.rate_limit, enabled=True
        )

    # Configure tool calling
    if args.enable_auto_tool_choice and args.tool_call_parser:
        server._enable_auto_tool_choice = True
        server._tool_call_parser = args.tool_call_parser
    else:
        server._enable_auto_tool_choice = False
        server._tool_call_parser = None

    # Configure generation defaults
    if args.default_temperature is not None:
        server._default_temperature = args.default_temperature
    if args.default_top_p is not None:
        server._default_top_p = args.default_top_p
    server._default_chat_template_kwargs = getattr(
        args, "default_chat_template_kwargs", None
    )
    if args.default_top_k is not None:
        server._default_top_k = args.default_top_k
    if args.default_min_p is not None:
        server._default_min_p = args.default_min_p
    if args.default_presence_penalty is not None:
        server._default_presence_penalty = args.default_presence_penalty
    if args.default_repetition_penalty is not None:
        server._default_repetition_penalty = args.default_repetition_penalty
    max_audio_upload_mb = getattr(args, "max_audio_upload_mb", 25)
    max_tts_input_chars = getattr(args, "max_tts_input_chars", 4096)
    server._max_audio_upload_bytes = max_audio_upload_mb * 1024 * 1024
    server._max_tts_input_chars = max_tts_input_chars

    # Configure thinking token budget
    default_thinking_token_budget = getattr(args, "default_thinking_token_budget", None)
    if default_thinking_token_budget is not None:
        server._default_thinking_token_budget = default_thinking_token_budget

    # Configure reasoning parser
    if args.reasoning_parser:
        try:
            from .reasoning import get_parser

            parser_cls = get_parser(args.reasoning_parser)
            server._reasoning_parser = parser_cls()
            server._reasoning_parser_name = args.reasoning_parser
            logger.info(f"Reasoning parser enabled: {args.reasoning_parser}")
        except KeyError as e:
            print(f"Error: {e}")
            sys.exit(1)
        except ImportError as e:
            print(f"Error: Failed to import reasoning module: {e}")
            sys.exit(1)
        except Exception as e:
            print(
                f"Error: Failed to initialize reasoning parser "
                f"'{args.reasoning_parser}': {e}"
            )
            sys.exit(1)
    else:
        server._reasoning_parser = None
        server._reasoning_parser_name = None

    # Security summary at startup
    print("=" * 60)
    print("SECURITY CONFIGURATION")
    print("=" * 60)
    if args.api_key:
        print("  Authentication: ENABLED (API key required)")
    else:
        print("  Authentication: DISABLED - Use --api-key to enable")
    if args.rate_limit > 0:
        print(f"  Rate limiting: ENABLED ({args.rate_limit} req/min)")
    else:
        print("  Rate limiting: DISABLED - Use --rate-limit to enable")
    print(f"  Request timeout: {args.timeout}s")
    if args.enable_metrics:
        print("  Metrics: ENABLED (/metrics, unauthenticated)")
    else:
        print("  Metrics: DISABLED - Use --enable-metrics to expose /metrics")
    if trust_remote_code:
        print("  Remote code loading: ENABLED (--trust-remote-code)")
    else:
        print("  Remote code loading: DISABLED (default)")
    if args.auto_unload_idle_seconds > 0:
        print(f"  Idle auto-unload: ENABLED ({args.auto_unload_idle_seconds:.0f}s)")
    else:
        print("  Idle auto-unload: DISABLED")
    if args.enable_auto_tool_choice:
        print(f"  Tool calling: ENABLED (parser: {args.tool_call_parser})")
    else:
        print("  Tool calling: Use --enable-auto-tool-choice to enable")
    if args.reasoning_parser:
        print(f"  Reasoning: ENABLED (parser: {args.reasoning_parser})")
    else:
        print("  Reasoning: Use --reasoning-parser to enable")
    if default_thinking_token_budget is not None:
        print(f"  Thinking budget: {default_thinking_token_budget} tokens")
    print(
        f"  Audio upload limit: {max_audio_upload_mb} MiB, "
        f"TTS input limit: {max_tts_input_chars} chars"
    )
    print("=" * 60)

    # Pre-download model with retry/timeout
    from .api.utils import is_mllm_model
    from .utils.download import DownloadConfig, ensure_model_downloaded

    download_config = DownloadConfig(
        download_timeout=args.download_timeout,
        max_retries=args.download_retries,
        offline=getattr(args, "offline", False),
    )
    if model_arg:
        ensure_model_downloaded(
            model_arg,
            config=download_config,
            is_mllm=is_mllm_model(model_arg),
        )
        if args.lazy_load_model:
            print(f"Registering model for lazy load: {model_arg}")
            print("Model will load on the first request.")
        else:
            print(f"Loading model: {model_arg}")
    else:
        print(f"Loading models config: {models_config}")
    print(f"Default max tokens: {args.max_tokens}")
    print(f"Max request tokens: {max_request_tokens}")
    if max_kv_size is not None:
        print(f"Max KV size: {max_kv_size} (RotatingKVCache)")

    # Store MCP config path for FastAPI startup
    if args.mcp_config:
        print(f"MCP config: {args.mcp_config}")
        os.environ["VLLM_MLX_MCP_CONFIG"] = args.mcp_config

    # Pre-load embedding model if specified
    embedding_model = getattr(args, "embedding_model", None)
    if embedding_model:
        print(f"Pre-loading embedding model: {embedding_model}")
        server.load_embedding_model(embedding_model, lock=True)
        print(f"Embedding model loaded: {embedding_model}")

    # Pre-load reranker model if specified
    rerank_model = getattr(args, "rerank_model", None)
    if rerank_model:
        print(f"Pre-loading reranker model: {rerank_model}")
        server.load_reranker_model(rerank_model, lock=True)
        print(f"Reranker model loaded: {rerank_model}")

    # Build scheduler config for batched mode
    scheduler_config = None
    if args.continuous_batching:
        from .scheduler import SchedulerConfig

        # Handle prefix cache flags
        enable_prefix_cache = args.enable_prefix_cache and not args.disable_prefix_cache

        scheduler_config = SchedulerConfig(
            max_num_seqs=args.max_num_seqs,
            prefill_batch_size=args.prefill_batch_size,
            completion_batch_size=args.completion_batch_size,
            enable_prefix_cache=enable_prefix_cache,
            prefix_cache_size=args.prefix_cache_size,
            # Memory-aware cache options
            use_memory_aware_cache=not args.no_memory_aware_cache,
            cache_memory_mb=args.cache_memory_mb,
            cache_memory_percent=args.cache_memory_percent,
            # Paged cache options
            use_paged_cache=args.use_paged_cache,
            paged_cache_block_size=args.paged_cache_block_size,
            max_cache_blocks=args.max_cache_blocks,
            # Chunked prefill
            chunked_prefill_tokens=args.chunked_prefill_tokens,
            # MTP
            enable_mtp=args.enable_mtp,
            mtp_num_draft_tokens=args.mtp_num_draft_tokens,
            mtp_optimistic=args.mtp_optimistic,
            # KV cache quantization
            kv_cache_quantization=args.kv_cache_quantization,
            kv_cache_quantization_bits=args.kv_cache_quantization_bits,
            kv_cache_quantization_group_size=args.kv_cache_quantization_group_size,
            kv_cache_min_quantize_tokens=args.kv_cache_min_quantize_tokens,
            mllm_prefill_step_size=(
                args.mllm_prefill_step_size if args.mllm_prefill_step_size > 0 else None
            ),
            # SSD cache tiering
            ssd_cache_dir=getattr(args, "ssd_cache_dir", None),
            ssd_cache_max_gb=getattr(args, "ssd_cache_max_gb", 10.0),
            # KV cache size limit
            max_kv_size=max_kv_size or 0,
        )

        print("Mode: Continuous batching (for multiple concurrent users)")
        if args.chunked_prefill_tokens > 0:
            print(f"Chunked prefill: {args.chunked_prefill_tokens} tokens per step")
        if args.enable_mtp:
            print(f"MTP: enabled, requested_draft_tokens={args.mtp_num_draft_tokens}")
            if args.mllm:
                print(
                    "MTP: MLLM path currently uses effective_draft_tokens=1 "
                    "per verify step; inspect /v1/status for attempts and acceptance"
                )
        print(f"Stream interval: {args.stream_interval} tokens")
        if args.use_paged_cache:
            print(
                f"Paged cache: block_size={args.paged_cache_block_size}, max_blocks={args.max_cache_blocks}"
            )
        elif enable_prefix_cache and not args.no_memory_aware_cache:
            cache_info = (
                f"{args.cache_memory_mb}MB"
                if args.cache_memory_mb
                else f"{args.cache_memory_percent*100:.0f}% of RAM"
            )
            print(f"Memory-aware cache: {cache_info}")
            if args.kv_cache_quantization:
                print(
                    f"KV cache quantization: {args.kv_cache_quantization_bits}-bit, "
                    f"group_size={args.kv_cache_quantization_group_size}"
                )
        elif enable_prefix_cache:
            print(f"Prefix cache: max_entries={args.prefix_cache_size}")
    else:
        print("Mode: Simple (maximum throughput)")
        if args.enable_mtp:
            print("MTP: enabled (native speculative decoding)")
        if args.enable_mtp and getattr(args, "mllm", False):
            print("MTP + MLLM: per-request routing (text-only → MTP, media → MLLM)")
        if args.specprefill and args.specprefill_draft_model:
            print(
                f"SpecPrefill: enabled (draft={args.specprefill_draft_model}, "
                f"threshold={args.specprefill_threshold}, "
                f"keep={args.specprefill_keep_pct*100:.0f}%)"
            )
        if mllm_draft_model:
            print(
                "MLLM draft model: enabled "
                f"(draft={mllm_draft_model}, kind={mllm_draft_kind}, "
                f"block_size={mllm_draft_block_size})"
            )

    if models_config:
        defaults = RegistryServeDefaults(
            continuous_batching=args.continuous_batching,
            force_mllm=getattr(args, "mllm", False),
            enable_mtp=args.enable_mtp,
            prefill_step_size=args.prefill_step_size,
            specprefill_enabled=args.specprefill,
            specprefill_threshold=args.specprefill_threshold,
            specprefill_keep_pct=args.specprefill_keep_pct,
            specprefill_draft_model=args.specprefill_draft_model,
            stream_interval=args.stream_interval if args.continuous_batching else 1,
            gpu_memory_utilization=args.gpu_memory_utilization,
            scheduler_config=scheduler_config,
            max_tokens=args.max_tokens,
            download_config=download_config,
        )
        load_model_registry(models_config, defaults=defaults)
    else:
        # Load model with unified server
        load_model(
            model_arg,
            use_batching=args.continuous_batching,
            scheduler_config=scheduler_config,
            stream_interval=args.stream_interval if args.continuous_batching else 1,
            max_tokens=args.max_tokens,
            max_request_tokens=max_request_tokens,
            force_mllm=getattr(args, "mllm", False),
            gpu_memory_utilization=args.gpu_memory_utilization,
            served_model_name=args.served_model_name,
            trust_remote_code=trust_remote_code,
            mtp=args.enable_mtp,
            prefill_step_size=args.prefill_step_size,
            specprefill_enabled=args.specprefill,
            specprefill_threshold=args.specprefill_threshold,
            specprefill_keep_pct=args.specprefill_keep_pct,
            specprefill_draft_model=args.specprefill_draft_model,
            mllm_draft_model=mllm_draft_model,
            mllm_draft_kind=mllm_draft_kind,
            mllm_draft_block_size=mllm_draft_block_size,
            warm_prompts_path=getattr(args, "warm_prompts", None),
            auto_unload_idle_seconds=args.auto_unload_idle_seconds,
            lazy_load_model=args.lazy_load_model,
        )

    # Start server
    print(f"Starting server at http://{args.host}:{args.port}")
    uvicorn.run(app, host=args.host, port=args.port, log_level="info")


def download_command(args):
    """Download a model to local cache without starting a server."""
    from .utils.download import DownloadConfig, ensure_model_downloaded

    config = DownloadConfig(
        download_timeout=args.timeout,
        max_retries=args.retries,
    )
    print(f"Downloading model: {args.model}")
    path = ensure_model_downloaded(
        args.model,
        config=config,
        is_mllm=args.mllm,
    )
    print(f"Model ready at: {path}")


def model_command(args):
    """Run model lifecycle helper commands."""
    from .model_workflow import (
        AcquisitionOptions,
        ConversionOptions,
        QualificationOptions,
        RegistrationOptions,
        acquire_model,
        convert_model,
        inspect_model,
        qualify_model,
        register_model,
    )

    if args.model_command == "inspect":
        payload = inspect_model(
            args.model,
            revision=args.revision,
            local_files_only=args.local_files_only,
        )
    elif args.model_command == "acquire":
        payload = acquire_model(
            args.model,
            options=AcquisitionOptions(
                revision=args.revision,
                target_dir=args.target_dir,
                staging_dir=args.staging_dir,
                is_mllm=args.mllm,
                fast_transfer=not args.no_fast_transfer,
                local_files_only=args.local_files_only,
            ),
        )
    elif args.model_command == "convert":
        payload = convert_model(
            ConversionOptions(
                source_path=args.source,
                output_path=args.output,
                quantize=args.quantize,
                q_bits=args.q_bits,
                q_group_size=args.q_group_size,
                q_mode=args.q_mode,
                quant_predicate=args.quant_predicate,
                dtype=args.dtype,
                trust_remote_code=args.trust_remote_code,
                dry_run=args.dry_run,
            )
        )
        if payload.get("status") == "failed":
            print(json.dumps(payload, indent=2))
            sys.exit(payload.get("returncode") or 1)
    elif args.model_command == "register":
        payload = register_model(
            RegistrationOptions(
                artifact_path=args.artifact,
                model_id=args.model_id,
                served_model_name=args.served_model_name,
                preset_alias=args.preset_alias,
                output_path=args.output,
                mllm=args.mllm,
                tool_call_parser=args.tool_call_parser,
                reasoning_parser=args.reasoning_parser,
                default_temperature=args.default_temperature,
                default_top_p=args.default_top_p,
                default_top_k=args.default_top_k,
                default_min_p=args.default_min_p,
                default_presence_penalty=args.default_presence_penalty,
                default_repetition_penalty=args.default_repetition_penalty,
                chat_template_kwargs=args.default_chat_template_kwargs,
                feature_flags=args.feature_flag,
            )
        )
    elif args.model_command == "qualify":
        payload = qualify_model(
            QualificationOptions(
                model_id=args.model_id,
                server_url=args.url,
                workload_path=args.workload,
                output_path=args.output,
                result_path=args.result_output,
                repetitions=args.repetitions,
                dry_run=args.dry_run,
                extra_args=args.extra_arg,
            )
        )
        if payload.get("status") == "failed":
            print(json.dumps(payload, indent=2))
            sys.exit(payload.get("returncode") or 1)
    else:
        raise ValueError(f"Unsupported model command: {args.model_command}")

    print(json.dumps(payload, indent=2))


def bench_command(args):
    """Run benchmark."""
    import asyncio
    import time

    from mlx_lm import load

    from .engine_core import AsyncEngineCore, EngineConfig
    from .request import SamplingParams
    from .scheduler import SchedulerConfig

    # Handle prefix cache flags
    enable_prefix_cache = args.enable_prefix_cache and not args.disable_prefix_cache

    async def run_benchmark():
        print(f"Loading model: {args.model}")
        model, tokenizer = load(args.model)

        scheduler_config = SchedulerConfig(
            max_num_seqs=args.max_num_seqs,
            prefill_batch_size=args.prefill_batch_size,
            completion_batch_size=args.completion_batch_size,
            enable_prefix_cache=enable_prefix_cache,
            prefix_cache_size=args.prefix_cache_size,
            # Memory-aware cache options
            use_memory_aware_cache=not args.no_memory_aware_cache,
            cache_memory_mb=args.cache_memory_mb,
            cache_memory_percent=args.cache_memory_percent,
            # Paged cache options
            use_paged_cache=args.use_paged_cache,
            paged_cache_block_size=args.paged_cache_block_size,
            max_cache_blocks=args.max_cache_blocks,
            # KV cache quantization
            kv_cache_quantization=args.kv_cache_quantization,
            kv_cache_quantization_bits=args.kv_cache_quantization_bits,
            kv_cache_quantization_group_size=args.kv_cache_quantization_group_size,
            kv_cache_min_quantize_tokens=args.kv_cache_min_quantize_tokens,
        )

        engine_config = EngineConfig(
            model_name=args.model,
            scheduler_config=scheduler_config,
        )

        if args.use_paged_cache:
            print(
                f"Paged cache: block_size={args.paged_cache_block_size}, max_blocks={args.max_cache_blocks}"
            )

        # Generate prompts
        prompts = [
            f"Write a short poem about {topic}."
            for topic in [
                "nature",
                "love",
                "technology",
                "space",
                "music",
                "art",
                "science",
                "history",
                "food",
                "travel",
            ][: args.num_prompts]
        ]

        params = SamplingParams(
            max_tokens=args.max_tokens,
            temperature=0.7,
        )

        print(
            f"\nRunning benchmark with {len(prompts)} prompts, max_tokens={args.max_tokens}"
        )
        print("-" * 50)

        total_prompt_tokens = 0
        total_completion_tokens = 0

        async with AsyncEngineCore(model, tokenizer, engine_config) as engine:
            await asyncio.sleep(0.1)  # Warm up

            start_time = time.perf_counter()

            # Add all requests
            request_ids = []
            for prompt in prompts:
                rid = await engine.add_request(prompt, params)
                request_ids.append(rid)

            # Collect all outputs
            async def get_output(rid):
                async for out in engine.stream_outputs(rid, timeout=120):
                    if out.finished:
                        return out
                return None

            results = await asyncio.gather(*[get_output(r) for r in request_ids])

            total_time = time.perf_counter() - start_time

        # Calculate stats
        for r in results:
            if r:
                total_prompt_tokens += r.prompt_tokens
                total_completion_tokens += r.completion_tokens

        total_tokens = total_prompt_tokens + total_completion_tokens

        print("\nResults:")
        print(f"  Total time: {total_time:.2f}s")
        print(f"  Prompts: {len(prompts)}")
        print(f"  Prompts/second: {len(prompts)/total_time:.2f}")
        print(f"  Total prompt tokens: {total_prompt_tokens}")
        print(f"  Total completion tokens: {total_completion_tokens}")
        print(f"  Total tokens: {total_tokens}")
        print(f"  Tokens/second: {total_completion_tokens/total_time:.2f}")
        print(f"  Throughput: {total_tokens/total_time:.2f} tok/s")

    asyncio.run(run_benchmark())


def bench_detok_command(args):
    """Benchmark streaming detokenizer optimization."""
    import statistics
    import time

    from mlx_lm import load
    from mlx_lm.generate import generate

    print("=" * 70)
    print(" Streaming Detokenizer Benchmark")
    print("=" * 70)
    print()

    print(f"Loading model: {args.model}")
    model, tokenizer = load(args.model)

    # Generate tokens for benchmark
    prompt = "Write a detailed explanation of how machine learning works and its applications in modern technology."
    print(f"Generating tokens with prompt: {prompt[:50]}...")

    output = generate(
        model=model,
        tokenizer=tokenizer,
        prompt=prompt,
        max_tokens=2000,
        verbose=False,
    )

    prompt_tokens = tokenizer.encode(prompt)
    all_tokens = tokenizer.encode(output)
    generated_tokens = all_tokens[len(prompt_tokens) :]
    print(f"Generated {len(generated_tokens)} tokens for benchmark")
    print()

    iterations = args.iterations

    # Benchmark naive decode (old method)
    print("Benchmarking Naive Decode (OLD method)...")
    naive_times = []
    for _ in range(iterations):
        start = time.perf_counter()
        for t in generated_tokens:
            _ = tokenizer.decode([t])
        elapsed = time.perf_counter() - start
        naive_times.append(elapsed)

    naive_mean = statistics.mean(naive_times) * 1000

    # Benchmark streaming decode (new method)
    print("Benchmarking Streaming Detokenizer (NEW method)...")
    streaming_times = []
    detok_class = tokenizer._detokenizer_class
    for _ in range(iterations):
        detok = detok_class(tokenizer)
        detok.reset()
        start = time.perf_counter()
        for t in generated_tokens:
            detok.add_token(t)
            _ = detok.last_segment
        detok.finalize()
        elapsed = time.perf_counter() - start
        streaming_times.append(elapsed)

    streaming_mean = statistics.mean(streaming_times) * 1000

    # Results
    speedup = naive_mean / streaming_mean
    time_saved = naive_mean - streaming_mean

    print()
    print("=" * 70)
    print(f" RESULTS: {len(generated_tokens)} tokens, {iterations} iterations")
    print("=" * 70)
    print(f"{'Method':<25} {'Time':>12} {'Speedup':>10}")
    print("-" * 70)
    print(f"{'Naive decode():':<25} {naive_mean:>10.2f}ms {'1.00x':>10}")
    print(f"{'Streaming detokenizer:':<25} {streaming_mean:>10.2f}ms {speedup:>9.2f}x")
    print("-" * 70)
    print(f"{'Time saved per request:':<25} {time_saved:>10.2f}ms")
    print(
        f"{'Per-token savings:':<25} {(time_saved/len(generated_tokens)*1000):>10.1f}µs"
    )
    print()

    # Verify correctness (strip for BPE edge cases with leading/trailing spaces)
    print("Verifying correctness...")
    detok = detok_class(tokenizer)
    detok.reset()
    for t in generated_tokens:
        detok.add_token(t)
    detok.finalize()

    batch_result = tokenizer.decode(generated_tokens)
    # BPE tokenizers may have minor edge case differences with spaces
    # Compare stripped versions for functional correctness
    streaming_stripped = detok.text.strip()
    batch_stripped = batch_result.strip()
    if streaming_stripped == batch_stripped:
        print("  ✓ Streaming output matches batch decode")
    elif streaming_stripped in batch_stripped or batch_stripped in streaming_stripped:
        print("  ✓ Streaming output matches (minor BPE edge case)")
    else:
        # Check if most of the content matches (BPE edge cases at boundaries)
        common_len = min(len(streaming_stripped), len(batch_stripped)) - 10
        if (
            common_len > 0
            and streaming_stripped[:common_len] == batch_stripped[:common_len]
        ):
            print("  ✓ Streaming output matches (BPE boundary difference)")
        else:
            print("  ✗ MISMATCH! Results differ")
            print(f"    Streaming: {repr(detok.text[:100])}...")
            print(f"    Batch: {repr(batch_result[:100])}...")


def bench_kv_cache_command(args):
    """Benchmark KV cache quantization memory savings and quality."""
    import time

    import mlx.core as mx
    from mlx_lm.models.cache import KVCache

    from .memory_cache import (
        _dequantize_cache,
        _quantize_cache,
        estimate_kv_cache_memory,
    )

    print("=" * 70)
    print(" KV Cache Quantization Benchmark")
    print("=" * 70)
    print()

    n_layers = args.layers
    seq_len = args.seq_len
    n_heads = args.heads
    head_dim = args.head_dim

    print(
        f"Config: {n_layers} layers, seq_len={seq_len}, "
        f"n_heads={n_heads}, head_dim={head_dim}"
    )
    print()

    # Create synthetic KV cache with random data
    print("Creating synthetic KV cache...")
    cache = []
    for _ in range(n_layers):
        kv = KVCache()
        kv.keys = mx.random.normal((1, n_heads, seq_len, head_dim))
        kv.values = mx.random.normal((1, n_heads, seq_len, head_dim))
        kv.offset = seq_len
        cache.append(kv)
    mx.eval(*[kv.keys for kv in cache], *[kv.values for kv in cache])

    fp16_mem = estimate_kv_cache_memory(cache)
    print(f"FP16 cache memory: {fp16_mem / 1024 / 1024:.2f} MB")
    print()

    # Test each bit width
    results = []
    for bits in [8, 4]:
        group_size = args.group_size

        # Quantize
        start = time.perf_counter()
        quantized = _quantize_cache(cache, bits=bits, group_size=group_size)
        mx.eval(
            *[
                layer.keys[0]
                for layer in quantized
                if hasattr(layer, "keys") and layer.keys is not None
            ]
        )
        quant_time = (time.perf_counter() - start) * 1000

        quant_mem = estimate_kv_cache_memory(quantized)

        # Dequantize
        start = time.perf_counter()
        restored = _dequantize_cache(quantized)
        mx.eval(
            *[
                layer.keys
                for layer in restored
                if hasattr(layer, "keys") and layer.keys is not None
            ]
        )
        dequant_time = (time.perf_counter() - start) * 1000

        # Measure quality
        total_error = 0.0
        max_error = 0.0
        count = 0
        for orig, rest in zip(cache, restored):
            if orig.keys is not None and rest.keys is not None:
                mx.eval(orig.keys, rest.keys, orig.values, rest.values)
                key_err = mx.abs(orig.keys - rest.keys).mean().item()
                val_err = mx.abs(orig.values - rest.values).mean().item()
                key_max = mx.abs(orig.keys - rest.keys).max().item()
                val_max = mx.abs(orig.values - rest.values).max().item()
                total_error += (key_err + val_err) / 2
                max_error = max(max_error, key_max, val_max)
                count += 1

        mean_error = total_error / count if count > 0 else 0.0
        ratio = fp16_mem / quant_mem if quant_mem > 0 else 0.0

        results.append(
            {
                "bits": bits,
                "mem_mb": quant_mem / 1024 / 1024,
                "ratio": ratio,
                "mean_err": mean_error,
                "max_err": max_error,
                "quant_ms": quant_time,
                "dequant_ms": dequant_time,
            }
        )

    # Print results
    fp16_mb = fp16_mem / 1024 / 1024
    print(
        f"{'Mode':<12} {'Memory':>10} {'Savings':>10} "
        f"{'Mean Err':>10} {'Max Err':>10} {'Quant':>10} {'Dequant':>10}"
    )
    print("-" * 72)
    print(
        f"{'FP16':<12} {fp16_mb:>8.2f}MB {'1.00x':>10} "
        f"{'0.000':>10} {'0.000':>10} {'-':>10} {'-':>10}"
    )

    for r in results:
        print(
            f"{r['bits']}-bit{'':<7} {r['mem_mb']:>8.2f}MB "
            f"{r['ratio']:>9.2f}x "
            f"{r['mean_err']:>10.5f} {r['max_err']:>10.5f} "
            f"{r['quant_ms']:>8.1f}ms {r['dequant_ms']:>8.1f}ms"
        )

    print()

    # Recommendation
    best = results[0]  # 8-bit
    print(
        f"Recommendation: 8-bit quantization gives {best['ratio']:.1f}x memory savings "
        f"with mean error {best['mean_err']:.5f}"
    )
    print(
        f"Use 4-bit for maximum compression if quality loss of "
        f"{results[1]['mean_err']:.4f} is acceptable."
    )
    print()
    print("Usage:")
    print("  vllm-mlx serve <model> --continuous-batching --kv-cache-quantization")
    print(
        "  vllm-mlx serve <model> --continuous-batching --kv-cache-quantization "
        "--kv-cache-quantization-bits 4"
    )


def bench_serve_command(args):
    """Run serving benchmark."""
    import asyncio

    from .bench_serve import run_bench_serve, run_bench_serve_workload

    if args.workload:
        sweep_only_warnings = []
        if args.prompts != "short,medium,long":
            sweep_only_warnings.append(f"--prompts={args.prompts}")
        if args.concurrency != "1,4":
            sweep_only_warnings.append(f"--concurrency={args.concurrency}")
        if args.warmup != 1:
            sweep_only_warnings.append(f"--warmup={args.warmup}")
        if sweep_only_warnings:
            import sys as _sys

            print(
                f"Warning: --workload mode ignores sweep-only args: "
                f"{', '.join(sweep_only_warnings)}",
                file=_sys.stderr,
            )
        request_timeout_s = (
            None if args.request_timeout_s <= 0 else args.request_timeout_s
        )
        output_format = "json" if args.format == "auto" else args.format
        asyncio.run(
            run_bench_serve_workload(
                url=args.url,
                workload_path=args.workload,
                model=args.model,
                output_path=args.output,
                output_format=output_format,
                scrape=args.scrape_metrics == "true",
                include_content=args.include_content,
                request_timeout_s=request_timeout_s,
                repetitions=args.repetitions,
                cache_policy=args.cache_policy,
            )
        )
        return

    prompt_sets = args.prompts.split(",")
    concurrencies = [int(c) for c in args.concurrency.split(",")]

    # Parse thinking values
    thinking_values = [None]
    if args.enable_thinking:
        thinking_values = []
        for v in args.enable_thinking.split(","):
            v = v.strip().lower()
            if v == "true":
                thinking_values.append(True)
            elif v == "false":
                thinking_values.append(False)

    # Parse extra body (comma-separated JSON dicts)
    extra_bodies = [""]
    if args.extra_body:
        # Handle both '{"a":1}','{"b":2}' and {"a":1},{"b":2}
        import re

        extra_bodies = [
            s.strip().strip("'\"")
            for s in re.split(r"(?<=})\s*,\s*(?={)", args.extra_body)
        ]

    # Parse override fields
    overrides = {}
    for kv in args.override_field or []:
        if "=" in kv:
            k, v = kv.split("=", 1)
            overrides[k] = v

    output_format = "table" if args.format == "auto" else args.format
    asyncio.run(
        run_bench_serve(
            url=args.url,
            model=args.model,
            prompt_sets=prompt_sets,
            prompt_file=args.prompt_file,
            concurrencies=concurrencies,
            max_tokens=args.max_tokens,
            repetitions=args.repetitions,
            warmup=args.warmup,
            thinking_values=thinking_values,
            extra_bodies=extra_bodies,
            output_path=args.output,
            fmt=output_format,
            do_validate=args.validate == "true",
            scrape=args.scrape_metrics == "true",
            tag=args.tag,
            override_fields=overrides,
            system_prompt_file=args.system_prompt_file,
            # Auto-enable skip-preflight when a system-prompt-file is set:
            # the whole point of that flag is measuring warm-cache behavior,
            # and the preflight count_prompt_tokens request pollutes the cache.
            skip_preflight_token_count=(
                args.skip_preflight_token_count or bool(args.system_prompt_file)
            ),
        )
    )


def create_parser() -> argparse.ArgumentParser:
    """Build the top-level CLI parser."""
    parser = argparse.ArgumentParser(
        description="vllm-mlx: Apple Silicon MLX backend for vLLM",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  vllm-mlx serve mlx-community/Llama-3.2-3B-Instruct-4bit --port 8000
  vllm-mlx bench mlx-community/Llama-3.2-1B-Instruct-4bit --num-prompts 10
        """,
    )
    subparsers = parser.add_subparsers(dest="command", help="Commands")

    # Serve command
    serve_parser = subparsers.add_parser("serve", help="Start OpenAI-compatible server")
    serve_parser.add_argument("model", nargs="?", type=str, help="Model to serve")
    serve_parser.add_argument(
        "--models-config",
        type=str,
        default=None,
        help="YAML file describing a registry of models for lazy multi-model serving",
    )
    serve_parser.add_argument(
        "--served-model-name",
        type=str,
        default=None,
        help="The model name used in the API. If not specified, the model argument is used.",
    )
    serve_parser.add_argument(
        "--host",
        type=str,
        default="127.0.0.1",
        help="Host to bind (default: localhost; use 0.0.0.0 to expose externally)",
    )
    serve_parser.add_argument("--port", type=int, default=8000, help="Port to bind")
    serve_parser.add_argument(
        "--max-num-seqs", type=int, default=256, help="Max concurrent sequences"
    )
    serve_parser.add_argument(
        "--prefill-batch-size", type=int, default=8, help="Prefill batch size"
    )
    serve_parser.add_argument(
        "--completion-batch-size", type=int, default=32, help="Completion batch size"
    )
    serve_parser.add_argument(
        "--mllm-prefill-step-size",
        type=int,
        default=0,
        help="Override MLLM prefill-step guard (0=use MLLM default: 1024)",
    )
    serve_parser.add_argument(
        "--enable-prefix-cache",
        action="store_true",
        default=True,
        help="Enable prefix caching for repeated prompts (default: enabled)",
    )
    serve_parser.add_argument(
        "--disable-prefix-cache",
        action="store_true",
        help="Disable prefix caching",
    )
    serve_parser.add_argument(
        "--prefix-cache-size",
        type=int,
        default=100,
        help="Max entries in prefix cache (default: 100, legacy mode only)",
    )
    # Memory-aware cache options (recommended for large models)
    serve_parser.add_argument(
        "--cache-memory-mb",
        type=int,
        default=None,
        help="Cache memory limit in MB (default: auto-detect ~20%% of RAM)",
    )
    serve_parser.add_argument(
        "--cache-memory-percent",
        type=float,
        default=0.20,
        help="Fraction of available RAM for cache if auto-detecting (default: 0.20)",
    )
    serve_parser.add_argument(
        "--no-memory-aware-cache",
        action="store_true",
        help="Disable memory-aware cache, use legacy entry-count based cache",
    )
    # KV cache quantization options
    serve_parser.add_argument(
        "--kv-cache-quantization",
        action="store_true",
        help="Quantize stored KV caches to reduce memory (8-bit by default)",
    )
    serve_parser.add_argument(
        "--kv-cache-quantization-bits",
        type=int,
        default=8,
        choices=[4, 8],
        help="Bit width for KV cache quantization (default: 8)",
    )
    serve_parser.add_argument(
        "--kv-cache-quantization-group-size",
        type=int,
        default=64,
        help="Group size for KV cache quantization (default: 64)",
    )
    serve_parser.add_argument(
        "--kv-cache-min-quantize-tokens",
        type=int,
        default=256,
        help="Minimum tokens for quantization to apply (default: 256)",
    )
    # SSD cache tiering options
    serve_parser.add_argument(
        "--ssd-cache-dir",
        type=str,
        default=None,
        help="Directory for SSD KV cache tier (default: disabled)",
    )
    serve_parser.add_argument(
        "--ssd-cache-max-gb",
        type=float,
        default=10.0,
        help="Maximum SSD cache size in GB (default: 10.0)",
    )
    # Prompt warm-up options
    serve_parser.add_argument(
        "--warm-prompts",
        type=str,
        default=None,
        help=(
            "Path to a JSON file with prompts to pre-run at startup. Populates "
            "the prefix cache so the first real request hits warm (cold TTFT "
            "drops 1.3-2.3x on agent workloads). File format is a list of "
            "message arrays, same shape as /v1/chat/completions messages. "
            "Prompts are warmed concurrently — keep the file small (1-3 entries "
            "for typical agent deployments) to avoid memory pressure at boot."
        ),
    )
    serve_parser.add_argument(
        "--stream-interval",
        type=int,
        default=1,
        help="Tokens to batch before streaming (1=smooth, higher=throughput)",
    )
    serve_parser.add_argument(
        "--max-kv-size",
        type=int,
        default=None,
        help="Maximum KV cache size per sequence. When set, uses RotatingKVCache "
        "which bounds memory at the cost of losing early context. Reasoning "
        "models (e.g. Qwen3, DeepSeek-R1) should use >= 32768 to avoid "
        "evicting the think block mid-generation.",
    )
    serve_parser.add_argument(
        "--max-tokens",
        type=int,
        default=32768,
        help="Default max tokens for generation (default: 32768)",
    )
    serve_parser.add_argument(
        "--max-request-tokens",
        type=int,
        default=32768,
        help="Maximum max_tokens accepted from API clients (default: 32768)",
    )
    serve_parser.add_argument(
        "--continuous-batching",
        action="store_true",
        help="Enable continuous batching for multiple concurrent users (slower for single user)",
    )
    serve_parser.add_argument(
        "--gpu-memory-utilization",
        type=float,
        default=0.90,
        help="Fraction of device memory for Metal allocation limit and emergency "
        "cache clear threshold (0.0-1.0, default: 0.90). Increase to 0.95 for "
        "large models (200GB+) that need more memory headroom.",
    )
    # Paged cache options (experimental)
    serve_parser.add_argument(
        "--use-paged-cache",
        action="store_true",
        help="Use paged KV cache for memory efficiency (experimental)",
    )
    serve_parser.add_argument(
        "--paged-cache-block-size",
        type=int,
        default=64,
        help="Tokens per cache block (default: 64)",
    )
    serve_parser.add_argument(
        "--max-cache-blocks",
        type=int,
        default=1000,
        help="Maximum number of cache blocks (default: 1000)",
    )
    # Chunked prefill
    serve_parser.add_argument(
        "--chunked-prefill-tokens",
        type=int,
        default=0,
        help="Max prefill tokens per scheduler step (0=disabled). "
        "Prevents starvation of active requests during long prefills.",
    )
    # MTP (Multi-Token Prediction)
    serve_parser.add_argument(
        "--enable-mtp",
        action="store_true",
        default=False,
        help="Enable MTP (Multi-Token Prediction) for models with built-in MTP heads. "
        "Uses cache snapshot/restore for speculative generation.",
    )
    serve_parser.add_argument(
        "--mtp-num-draft-tokens",
        type=int,
        default=1,
        help="Number of draft tokens per MTP step (default: 1)",
    )
    serve_parser.add_argument(
        "--mtp-optimistic",
        action="store_true",
        default=False,
        help="Skip MTP acceptance check for maximum speed. "
        "~5-10%% wrong tokens. Best for chat, not for code.",
    )
    # Prefill step size
    serve_parser.add_argument(
        "--prefill-step-size",
        type=int,
        default=2048,
        help="Chunk size for prompt prefill processing. Larger values use more memory "
        "but can improve prefill throughput. (default: 2048)",
    )
    # SpecPrefill (attention-based sparse prefill using draft model)
    serve_parser.add_argument(
        "--specprefill",
        action="store_true",
        default=False,
        help="Enable SpecPrefill: use a small draft model to score token importance, "
        "then sparse-prefill only the important tokens on the target model. "
        "Reduces TTFT on long prompts. Requires --specprefill-draft-model.",
    )
    serve_parser.add_argument(
        "--specprefill-threshold",
        type=int,
        default=8192,
        help="Minimum suffix tokens to trigger SpecPrefill (default: 8192). "
        "Shorter prompts use full prefill (scoring overhead > savings).",
    )
    serve_parser.add_argument(
        "--specprefill-keep-pct",
        type=float,
        default=0.3,
        help="Fraction of tokens to keep during sparse prefill (default: 0.3). "
        "Lower = faster prefill but more quality loss.",
    )
    serve_parser.add_argument(
        "--specprefill-draft-model",
        type=str,
        default=None,
        help="Path to small draft model for SpecPrefill importance scoring. "
        "Must share the same tokenizer as the target model.",
    )
    # MLLM speculative draft/assistant model
    serve_parser.add_argument(
        "--mllm-draft-model",
        type=str,
        default=None,
        help="Path to an mlx-vlm MLLM draft/assistant model. "
        "For Gemma 4 assistant drafters, use with --mllm-draft-kind mtp.",
    )
    serve_parser.add_argument(
        "--mllm-draft-kind",
        type=str,
        default=None,
        choices=["mtp"],
        help="mlx-vlm draft kind for --mllm-draft-model.",
    )
    serve_parser.add_argument(
        "--mllm-draft-block-size",
        type=make_positive_int_arg_parser("--mllm-draft-block-size"),
        default=None,
        help="Draft block size passed to mlx-vlm for --mllm-draft-model.",
    )
    # MCP options
    serve_parser.add_argument(
        "--mcp-config",
        type=str,
        default=None,
        help="Path to MCP configuration file (JSON/YAML) for tool integration",
    )
    # Security options
    serve_parser.add_argument(
        "--api-key",
        type=str,
        default=None,
        help="API key for authentication (if not set, no auth required)",
    )
    serve_parser.add_argument(
        "--rate-limit",
        type=int,
        default=0,
        help="Rate limit requests per minute per client (0 = disabled)",
    )
    serve_parser.add_argument(
        "--timeout",
        type=float,
        default=300.0,
        help="Default request timeout in seconds (default: 300)",
    )
    serve_parser.add_argument(
        "--enable-metrics",
        action="store_true",
        help="Expose Prometheus metrics on /metrics (disabled by default)",
    )
    serve_parser.add_argument(
        "--auto-unload-idle-seconds",
        type=float,
        default=0.0,
        help="Unload the main model after this many idle seconds (0 = disabled)",
    )
    serve_parser.add_argument(
        "--lazy-load-model",
        action="store_true",
        help="Register the main model at startup but defer loading until first request",
    )
    serve_parser.add_argument(
        "--max-audio-upload-mb",
        type=int,
        default=25,
        help="Maximum size of uploaded audio files in MiB (default: 25)",
    )
    serve_parser.add_argument(
        "--max-tts-input-chars",
        type=int,
        default=4096,
        help="Maximum number of characters accepted by /v1/audio/speech (default: 4096)",
    )
    # Tool calling options
    serve_parser.add_argument(
        "--enable-auto-tool-choice",
        action="store_true",
        help="Enable auto tool choice for supported models. Use --tool-call-parser to specify which parser to use.",
    )
    serve_parser.add_argument(
        "--tool-call-parser",
        type=str,
        default=None,
        choices=[
            "auto",
            "mistral",
            "qwen",
            "qwen3_coder",
            "llama",
            "hermes",
            "harmony",
            "gpt-oss",
            "deepseek",
            "kimi",
            "granite",
            "nemotron",
            "xlam",
            "functionary",
            "gemma4",
            "glm47",
            "minimax",
        ],
        help=(
            "Select the tool call parser for the model. Options: "
            "auto (auto-detect), mistral, qwen, qwen3_coder, llama, hermes, "
            "harmony, gpt-oss, deepseek, gemma4, kimi, granite, nemotron, "
            "xlam, functionary, glm47, minimax. "
            "Required for --enable-auto-tool-choice."
        ),
    )
    # Reasoning parser options - choices loaded dynamically from registry
    from .reasoning import list_parsers

    reasoning_choices = list_parsers()
    serve_parser.add_argument(
        "--reasoning-parser",
        type=str,
        default=None,
        choices=reasoning_choices,
        help=(
            "Enable reasoning content extraction with specified parser. "
            "Extracts <think>...</think> tags into reasoning_content field. "
            f"Options: {', '.join(reasoning_choices)}."
        ),
    )
    # Multimodal option
    serve_parser.add_argument(
        "--mllm",
        action="store_true",
        help="Force load model as multimodal (vision) even if name doesn't match auto-detection patterns",
    )
    serve_parser.add_argument(
        "--trust-remote-code",
        action="store_true",
        help="Allow HuggingFace remote code execution during model/tokenizer loading",
    )
    # Generation defaults
    serve_parser.add_argument(
        "--default-temperature",
        type=float,
        default=None,
        help="Override default temperature for all requests (default: use model default)",
    )
    serve_parser.add_argument(
        "--default-top-p",
        type=float,
        default=None,
        help="Override default top_p for all requests (default: use model default)",
    )
    serve_parser.add_argument(
        "--default-thinking-token-budget",
        type=int,
        default=None,
        help=(
            "Default thinking token budget for reasoning models. Caps reasoning "
            "tokens by forcing the end-think sequence when the budget is exhausted. "
            "Per-request thinking_token_budget overrides this. (default: None = unlimited)"
        ),
    )
    serve_parser.add_argument(
        "--default-chat-template-kwargs",
        type=make_json_object_arg_parser("--default-chat-template-kwargs"),
        default=None,
        help=(
            "Default chat template kwargs to apply to all requests when request "
            "chat_template_kwargs is omitted or empty; empty request kwargs use "
            'existing server defaults (JSON object, e.g. {"enable_thinking": true})'
        ),
    )
    serve_parser.add_argument(
        "--default-top-k",
        type=int,
        default=None,
        help="Override default top_k for all requests (default: use model default)",
    )
    serve_parser.add_argument(
        "--default-min-p",
        type=float,
        default=None,
        help="Override default min_p for all requests (default: use model default)",
    )
    serve_parser.add_argument(
        "--default-presence-penalty",
        type=float,
        default=None,
        help=(
            "Override default presence_penalty for all requests "
            "(default: use model default)"
        ),
    )
    serve_parser.add_argument(
        "--default-repetition-penalty",
        type=float,
        default=None,
        help=(
            "Override default repetition_penalty for all requests "
            "(default: use model default)"
        ),
    )
    # Embedding model option
    serve_parser.add_argument(
        "--embedding-model",
        type=str,
        default=None,
        help="Pre-load an embedding model at startup (e.g. mlx-community/embeddinggemma-300m-6bit)",
    )
    # Reranker model option
    serve_parser.add_argument(
        "--rerank-model",
        type=str,
        default=None,
        help="Pre-load a reranker model at startup (e.g. mlx-community/jina-reranker-v2-base-multilingual)",
    )
    # Download options
    serve_parser.add_argument(
        "--download-timeout",
        type=int,
        default=300,
        help="Per-file download timeout in seconds (default: 300)",
    )
    serve_parser.add_argument(
        "--download-retries",
        type=int,
        default=3,
        help="Number of download retry attempts (default: 3)",
    )
    serve_parser.add_argument(
        "--offline",
        action="store_true",
        help="Offline mode — only use locally cached models",
    )
    # Bench command
    bench_parser = subparsers.add_parser("bench", help="Run benchmark")
    bench_parser.add_argument("model", type=str, help="Model to benchmark")
    bench_parser.add_argument(
        "--num-prompts", type=int, default=10, help="Number of prompts"
    )
    bench_parser.add_argument(
        "--max-tokens", type=int, default=100, help="Max tokens per prompt"
    )
    bench_parser.add_argument(
        "--max-num-seqs", type=int, default=32, help="Max concurrent sequences"
    )
    bench_parser.add_argument(
        "--prefill-batch-size", type=int, default=8, help="Prefill batch size"
    )
    bench_parser.add_argument(
        "--completion-batch-size", type=int, default=16, help="Completion batch size"
    )
    bench_parser.add_argument(
        "--enable-prefix-cache",
        action="store_true",
        default=True,
        help="Enable prefix caching (default: enabled)",
    )
    bench_parser.add_argument(
        "--disable-prefix-cache",
        action="store_true",
        help="Disable prefix caching",
    )
    bench_parser.add_argument(
        "--prefix-cache-size",
        type=int,
        default=100,
        help="Max entries in prefix cache (default: 100, legacy mode only)",
    )
    # Memory-aware cache options (recommended for large models)
    bench_parser.add_argument(
        "--cache-memory-mb",
        type=int,
        default=None,
        help="Cache memory limit in MB (default: auto-detect ~20%% of RAM)",
    )
    bench_parser.add_argument(
        "--cache-memory-percent",
        type=float,
        default=0.20,
        help="Fraction of available RAM for cache if auto-detecting (default: 0.20)",
    )
    bench_parser.add_argument(
        "--no-memory-aware-cache",
        action="store_true",
        help="Disable memory-aware cache, use legacy entry-count based cache",
    )
    # KV cache quantization options
    bench_parser.add_argument(
        "--kv-cache-quantization",
        action="store_true",
        help="Quantize stored KV caches to reduce memory (8-bit by default)",
    )
    bench_parser.add_argument(
        "--kv-cache-quantization-bits",
        type=int,
        default=8,
        choices=[4, 8],
        help="Bit width for KV cache quantization (default: 8)",
    )
    bench_parser.add_argument(
        "--kv-cache-quantization-group-size",
        type=int,
        default=64,
        help="Group size for KV cache quantization (default: 64)",
    )
    bench_parser.add_argument(
        "--kv-cache-min-quantize-tokens",
        type=int,
        default=256,
        help="Minimum tokens for quantization to apply (default: 256)",
    )
    # Paged cache options (experimental)
    bench_parser.add_argument(
        "--use-paged-cache",
        action="store_true",
        help="Use paged KV cache for memory efficiency (experimental)",
    )
    bench_parser.add_argument(
        "--paged-cache-block-size",
        type=int,
        default=64,
        help="Tokens per cache block (default: 64)",
    )
    bench_parser.add_argument(
        "--max-cache-blocks",
        type=int,
        default=1000,
        help="Maximum number of cache blocks (default: 1000)",
    )

    # Detokenizer benchmark
    detok_parser = subparsers.add_parser(
        "bench-detok", help="Benchmark streaming detokenizer optimization"
    )
    detok_parser.add_argument(
        "model",
        type=str,
        nargs="?",
        default="mlx-community/Qwen3-0.6B-8bit",
        help="Model to use for tokenizer (default: mlx-community/Qwen3-0.6B-8bit)",
    )
    detok_parser.add_argument(
        "--iterations", type=int, default=5, help="Benchmark iterations (default: 5)"
    )

    # KV cache quantization benchmark
    kv_cache_parser = subparsers.add_parser(
        "bench-kv-cache", help="Benchmark KV cache quantization memory savings"
    )
    kv_cache_parser.add_argument(
        "--layers", type=int, default=32, help="Number of layers (default: 32)"
    )
    kv_cache_parser.add_argument(
        "--seq-len", type=int, default=512, help="Sequence length (default: 512)"
    )
    kv_cache_parser.add_argument(
        "--heads", type=int, default=32, help="Number of attention heads (default: 32)"
    )
    kv_cache_parser.add_argument(
        "--head-dim", type=int, default=128, help="Head dimension (default: 128)"
    )
    kv_cache_parser.add_argument(
        "--group-size",
        type=int,
        default=64,
        help="Quantization group size (default: 64)",
    )

    # Download command
    download_parser = subparsers.add_parser(
        "download", help="Download a model to local cache without starting a server"
    )
    download_parser.add_argument("model", type=str, help="Model to download")
    download_parser.add_argument(
        "--timeout",
        type=int,
        default=300,
        help="Per-file download timeout in seconds (default: 300)",
    )
    download_parser.add_argument(
        "--retries",
        type=int,
        default=3,
        help="Number of retry attempts (default: 3)",
    )
    download_parser.add_argument(
        "--mllm",
        action="store_true",
        help="Download as multimodal model (broader file patterns)",
    )

    # Model lifecycle helpers
    model_parser = subparsers.add_parser(
        "model",
        help="Inspect, acquire, or convert model artifacts",
    )
    model_subparsers = model_parser.add_subparsers(
        dest="model_command", help="Model workflow command", required=True
    )

    model_inspect_parser = model_subparsers.add_parser(
        "inspect",
        help="Inspect a local path or Hugging Face model without loading weights",
    )
    model_inspect_parser.add_argument(
        "model",
        type=str,
        help="Local model path or Hugging Face model id",
    )
    model_inspect_parser.add_argument(
        "--revision",
        type=str,
        default=None,
        help="Hugging Face revision to inspect",
    )
    model_inspect_parser.add_argument(
        "--local-files-only",
        action="store_true",
        help="Use only local Hugging Face cache files",
    )

    model_acquire_parser = model_subparsers.add_parser(
        "acquire",
        help="Download a Hugging Face model and write an artifact manifest",
    )
    model_acquire_parser.add_argument("model", type=str, help="Hugging Face model id")
    model_acquire_parser.add_argument(
        "--revision",
        type=str,
        default=None,
        help="Hugging Face revision to download",
    )
    model_acquire_parser.add_argument(
        "--target-dir",
        type=str,
        default=None,
        help="Final local directory. Defaults to Hugging Face cache.",
    )
    model_acquire_parser.add_argument(
        "--staging-dir",
        type=str,
        default=None,
        help="Directory for temporary staged downloads before finalizing target-dir",
    )
    model_acquire_parser.add_argument(
        "--mllm",
        action="store_true",
        help="Acquire multimodal model files using broader allow patterns",
    )
    model_acquire_parser.add_argument(
        "--no-fast-transfer",
        action="store_true",
        help="Do not set HF_HUB_ENABLE_HF_TRANSFER=1 during download",
    )
    model_acquire_parser.add_argument(
        "--local-files-only",
        action="store_true",
        help="Use only local Hugging Face cache files",
    )

    model_convert_parser = model_subparsers.add_parser(
        "convert",
        help="Run mlx-lm conversion and write a conversion manifest",
    )
    model_convert_parser.add_argument(
        "source",
        type=str,
        help="Hugging Face model id or local source path",
    )
    model_convert_parser.add_argument(
        "--output",
        required=True,
        type=str,
        help="Output directory for the converted MLX model",
    )
    model_convert_parser.add_argument(
        "--quantize",
        action="store_true",
        help="Generate a quantized MLX model",
    )
    model_convert_parser.add_argument(
        "--q-bits",
        type=int,
        default=None,
        help="Quantization bit width (e.g. 3, 4, 8)",
    )
    model_convert_parser.add_argument(
        "--q-group-size",
        type=int,
        default=None,
        help="Quantization group size (default: mlx-lm default)",
    )
    model_convert_parser.add_argument(
        "--q-mode",
        choices=["affine", "mxfp4", "nvfp4", "mxfp8"],
        default=None,
    )
    model_convert_parser.add_argument(
        "--quant-predicate",
        choices=["mixed_2_6", "mixed_3_4", "mixed_3_6", "mixed_4_6"],
        default=None,
        help="mlx-lm mixed-bit quantization recipe",
    )
    model_convert_parser.add_argument(
        "--dtype",
        choices=["float16", "bfloat16", "float32"],
        default=None,
        help="Non-quantized parameter dtype",
    )
    model_convert_parser.add_argument(
        "--trust-remote-code",
        action="store_true",
        help="Allow Hugging Face remote code during mlx-lm conversion",
    )
    model_convert_parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print the conversion command and manifest without executing",
    )

    model_register_parser = model_subparsers.add_parser(
        "register",
        help="Write a portable registration manifest for a finalized artifact",
    )
    model_register_parser.add_argument(
        "artifact",
        type=str,
        help="Finalized local model artifact directory",
    )
    model_register_parser.add_argument(
        "--model-id",
        type=str,
        default=None,
        help="Override model ID (default: directory name of artifact)",
    )
    model_register_parser.add_argument(
        "--served-model-name",
        type=str,
        default=None,
        help="Model name exposed by the API (default: model-id)",
    )
    model_register_parser.add_argument(
        "--preset-alias",
        type=str,
        default=None,
        help="Optional alias for preset lookup in registry",
    )
    model_register_parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Manifest path. Defaults to artifact/vllm_mlx_registration_manifest.json",
    )
    mllm_group = model_register_parser.add_mutually_exclusive_group()
    mllm_group.add_argument(
        "--mllm",
        action="store_true",
        default=None,
        help="Mark the artifact as an MLLM serving candidate",
    )
    mllm_group.add_argument(
        "--no-mllm",
        action="store_false",
        dest="mllm",
        help="Explicitly mark the artifact as text-only",
    )
    model_register_parser.add_argument(
        "--tool-call-parser",
        type=str,
        default=None,
        help="Tool call parser name for the model (e.g. qwen3_coder, mistral)",
    )
    model_register_parser.add_argument(
        "--reasoning-parser",
        type=str,
        default=None,
        help="Reasoning parser name for thinking models (e.g. qwen3)",
    )
    model_register_parser.add_argument(
        "--default-temperature",
        type=float,
        default=None,
        help="Default temperature for all requests",
    )
    model_register_parser.add_argument(
        "--default-top-p",
        type=float,
        default=None,
        help="Default top_p for all requests",
    )
    model_register_parser.add_argument(
        "--default-top-k",
        type=int,
        default=None,
        help="Default top_k for all requests",
    )
    model_register_parser.add_argument(
        "--default-min-p",
        type=float,
        default=None,
        help="Default min_p for all requests",
    )
    model_register_parser.add_argument(
        "--default-presence-penalty",
        type=float,
        default=None,
        help="Default presence_penalty for all requests",
    )
    model_register_parser.add_argument(
        "--default-repetition-penalty",
        type=float,
        default=None,
        help="Default repetition_penalty for all requests",
    )
    model_register_parser.add_argument(
        "--default-chat-template-kwargs",
        type=make_json_object_arg_parser("--default-chat-template-kwargs"),
        default=None,
        help='Default chat template kwargs as JSON, e.g. {"enable_thinking": true}',
    )
    model_register_parser.add_argument(
        "--feature-flag",
        action="append",
        default=[],
        help="Feature flag to record in the registration manifest. Repeatable.",
    )

    model_qualify_parser = model_subparsers.add_parser(
        "qualify",
        help="Create or run a bench-serve qualification handoff",
    )
    model_qualify_parser.add_argument(
        "model_id",
        type=str,
        help="Model ID to qualify against the running server",
    )
    model_qualify_parser.add_argument(
        "--url",
        type=str,
        default="http://127.0.0.1:8080",
        help="Running server URL for bench-serve",
    )
    model_qualify_parser.add_argument(
        "--workload",
        type=str,
        default=None,
        help="bench-serve workload contract path",
    )
    model_qualify_parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Qualification request manifest path",
    )
    model_qualify_parser.add_argument(
        "--result-output",
        type=str,
        default=None,
        help="Result output path passed to bench-serve",
    )
    model_qualify_parser.add_argument(
        "--repetitions",
        type=int,
        default=None,
        help="Number of repetitions per benchmark sweep configuration",
    )
    model_qualify_parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Write or print the qualification command without running it",
    )
    model_qualify_parser.add_argument(
        "--extra-arg",
        action="append",
        default=[],
        help="Extra argument passed through to bench-serve. Repeatable.",
    )

    # Serving benchmark
    bench_serve_parser = subparsers.add_parser(
        "bench-serve", help="Benchmark a running vllm-mlx server via HTTP API"
    )
    bench_serve_parser.add_argument(
        "--url",
        type=str,
        default="http://127.0.0.1:8080",
        help="Base URL of the running server (default: http://127.0.0.1:8080)",
    )
    bench_serve_parser.add_argument(
        "--model",
        type=str,
        default=None,
        help="Model ID to benchmark (default: auto-detected from server)",
    )
    bench_serve_parser.add_argument(
        "--workload",
        type=str,
        default=None,
        help=(
            "Path to a declarative workload JSON file. When set, bench-serve "
            "runs contract-style cases with per-case quality checks and "
            "comparison-only policy timeouts instead of the prompt sweep."
        ),
    )
    bench_serve_parser.add_argument(
        "--prompts",
        type=str,
        default="short,medium,long",
        help="Comma-separated prompt set names or paths (default: short,medium,long)",
    )
    bench_serve_parser.add_argument(
        "--prompt-file",
        type=str,
        default=None,
        help="Path to an additional prompt file (JSON list of message dicts)",
    )
    bench_serve_parser.add_argument(
        "--system-prompt-file",
        type=str,
        default=None,
        help=(
            "Path to a text file whose contents are prepended as a system "
            "message to every prompt. Use this together with --warm-prompts "
            "to benchmark the warm-cache path (the warmup populates the "
            "prefix cache with this same system, so every request in the "
            "bench hits the cache)."
        ),
    )
    bench_serve_parser.add_argument(
        "--skip-preflight-token-count",
        action="store_true",
        help=(
            "Skip the pre-flight max_tokens=1 request that counts prompt "
            "tokens per prompt set. That request populates the prefix cache "
            "with the full prompt, which defeats cold-vs-warm comparisons. "
            "Auto-enabled when --system-prompt-file is set; pass this flag "
            "explicitly to force-enable regardless."
        ),
    )
    bench_serve_parser.add_argument(
        "--concurrency",
        type=str,
        default="1,4",
        help="Comma-separated concurrency levels to sweep (default: 1,4)",
    )
    bench_serve_parser.add_argument(
        "--max-tokens",
        type=int,
        default=256,
        help="Maximum tokens to generate per request (default: 256)",
    )
    bench_serve_parser.add_argument(
        "--repetitions",
        type=int,
        default=3,
        help="Number of repetitions per sweep configuration or workload case (default: 3)",
    )
    bench_serve_parser.add_argument(
        "--warmup",
        type=int,
        default=1,
        help="Warmup rounds before the first measured repetition (default: 1)",
    )
    bench_serve_parser.add_argument(
        "--enable-thinking",
        type=str,
        default=None,
        help='Enable thinking mode: "true", "false", or "true,false" to sweep both',
    )
    bench_serve_parser.add_argument(
        "--extra-body",
        type=str,
        default=None,
        help="Comma-separated JSON dicts to pass as extra body parameters",
    )
    bench_serve_parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="File path to write results to (default: stdout)",
    )
    bench_serve_parser.add_argument(
        "--format",
        type=str,
        default="auto",
        choices=["auto", "table", "json", "csv", "sql", "sqlite"],
        help=(
            "Output format (auto = table for prompt sweeps, json for workloads; "
            "sqlite requires --output)"
        ),
    )
    bench_serve_parser.add_argument(
        "--validate",
        type=str,
        default="true",
        choices=["true", "false"],
        help="Validate responses (default: true)",
    )
    bench_serve_parser.add_argument(
        "--scrape-metrics",
        type=str,
        default="true",
        choices=["true", "false"],
        help="Scrape /metrics before and after each run (default: true)",
    )
    bench_serve_parser.add_argument(
        "--include-content",
        action="store_true",
        help="Include full generated content in workload JSON output",
    )
    bench_serve_parser.add_argument(
        "--request-timeout-s",
        type=float,
        default=300.0,
        help=(
            "HTTP transport timeout for workload mode in seconds (default: 300). "
            "Use 0 to disable; product policy timeouts belong in the workload."
        ),
    )
    bench_serve_parser.add_argument(
        "--cache-policy",
        type=str,
        default=None,
        choices=["preserve", "before-run", "before-case"],
        help=(
            "Workload cache handling (default: workload defaults or preserve). "
            "Use before-case for cold, uncontaminated per-case qualification. "
            "Workload JSON may also spell these with underscores."
        ),
    )
    bench_serve_parser.add_argument(
        "--tag",
        type=str,
        default=None,
        help="Optional tag string stored in every result row",
    )
    bench_serve_parser.add_argument(
        "--override-field",
        nargs="*",
        default=[],
        help="Override result fields as key=value pairs (e.g. chip=M4Pro)",
    )

    return parser


# Alias for test compatibility
build_parser = create_parser


def main():
    parser = create_parser()
    args = parser.parse_args()

    if args.command == "serve":
        serve_command(args)
    elif args.command == "bench":
        bench_command(args)
    elif args.command == "bench-detok":
        bench_detok_command(args)
    elif args.command == "bench-kv-cache":
        bench_kv_cache_command(args)
    elif args.command == "download":
        download_command(args)
    elif args.command == "model":
        model_command(args)
    elif args.command == "bench-serve":
        bench_serve_command(args)
    else:
        parser.print_help()
        sys.exit(1)


if __name__ == "__main__":
    main()
