"""Per-method quantization specs and vLLM serve-argument construction.

kind:
  baseline   - serve the unquantized model as-is (the accuracy reference)
  online     - vLLM quantizes at load (fp8 dynamic); serve the base model
  compressor - llm-compressor produced a checkpoint; quant config travels with it
  bnb        - transformers+bitsandbytes produced a checkpoint; serve with bnb flags
"""
from __future__ import annotations

from typing import Dict, List

METHOD_SPECS: Dict[str, Dict] = {
    "bf16": {"kind": "baseline", "t4_deployable": True},
    "fp8": {"kind": "online", "vllm_quant": "fp8", "t4_deployable": False},
    # smoothquant=False: SmoothQuant's default mappings fail to resolve for the
    # qwen3_5 layer layout (each mapping matches all 24 input_layernorms), so
    # int8 is plain GPTQ W8A8.
    "int8": {"kind": "compressor", "modifier": "gptq", "scheme": "W8A8", "smoothquant": False, "t4_deployable": True},
    "awq": {"kind": "compressor", "modifier": "awq", "scheme": "W4A16", "smoothquant": False, "t4_deployable": True},
    "gptq": {"kind": "compressor", "modifier": "gptq", "scheme": "W4A16", "smoothquant": False, "t4_deployable": True},
    "bnb8": {"kind": "bnb", "bits": 8, "t4_deployable": True},
    "bnb4": {"kind": "bnb", "bits": 4, "t4_deployable": True},
}


def method_names() -> List[str]:
    return list(METHOD_SPECS.keys())


def vllm_serve_args(method: str, model_path: str, base_model: str, port: int) -> List[str]:
    spec = METHOD_SPECS[method]
    kind = spec["kind"]
    serve_model = base_model if kind in ("baseline", "online") else model_path
    args = [
        "--host", "127.0.0.1",
        "--port", str(port),
        "--model", serve_model,
        "--served-model-name", "datalab-to/surya-ocr-2",
        "--max-model-len", "18000",
        "--max-num-seqs", "16",
        "--gpu-memory-utilization", "0.85",
        "--enable-prefix-caching",
        "--mm-processor-kwargs", '{"min_pixels":3136,"max_pixels":6291456}',
    ]
    if kind == "online":
        args += ["--quantization", spec["vllm_quant"]]
    elif kind == "bnb":
        args += ["--quantization", "bitsandbytes", "--load-format", "bitsandbytes"]
    return args