"""Produce vLLM-loadable quantized checkpoints.

compressor methods (int8/awq/gptq) use llm-compressor `oneshot`.
bnb methods (bnb8/bnb4) use transformers + BitsAndBytesConfig + save_pretrained.
baseline/online (bf16/fp8) need no build.
"""
from __future__ import annotations

from pathlib import Path
from typing import List

from scripts.quant.recipes import METHOD_SPECS


def needs_build(method: str) -> bool:
    return METHOD_SPECS[method]["kind"] in ("compressor", "bnb")


def _is_built(out_dir: Path) -> bool:
    return (out_dir / "config.json").exists()


CALIB_MAX_SEQ_LEN = 8192
# Never quantize the LM head or the vision tower: vision modules are shape-fragile
# (see the FP8 Marlin failure) and contribute little to decode cost.
QUANT_IGNORE = ["re:.*lm_head", "re:.*visual.*", "re:.*vision.*"]


def _calibration_dataset(calib_images: List[Path], processor):
    """HF Dataset of pre-tokenized multimodal samples (batch dim kept), matching
    llm-compressor's multimodal-vision examples; re-tensorized by _data_collator.
    Activation calibration only needs representative forward passes, not the
    exact training prompt."""
    from datasets import Dataset
    from PIL import Image

    samples = []
    for path in calib_images:
        image = Image.open(path).convert("RGB")
        messages = [{
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": "OCR this document."},
            ],
        }]
        prompt = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
        inputs = processor(
            text=[prompt], images=[image],
            padding=False, truncation=True, max_length=CALIB_MAX_SEQ_LEN,
        )
        samples.append({k: (v.tolist() if hasattr(v, "tolist") else v) for k, v in inputs.items()})
    return Dataset.from_list(samples)


def _data_collator(batch):
    import torch

    assert len(batch) == 1
    return {key: torch.tensor(value) for key, value in batch[0].items()}


def _build_compressor(method: str, base_model: str, out_dir: Path, calib_images: List[Path]) -> None:
    from llmcompressor import oneshot
    from llmcompressor.modifiers.quantization import GPTQModifier
    from transformers import AutoProcessor

    try:
        from transformers import AutoModelForImageTextToText as _AutoModel
    except ImportError:  # older transformers
        from transformers import AutoModelForCausalLM as _AutoModel

    spec = METHOD_SPECS[method]
    model = _AutoModel.from_pretrained(base_model, dtype="auto", device_map="auto")
    processor = AutoProcessor.from_pretrained(base_model)
    dataset = _calibration_dataset(calib_images, processor)

    if spec["modifier"] == "awq":
        # llm-compressor main: AWQ is a transform paired with a QuantizationModifier.
        from llmcompressor.modifiers.quantization import QuantizationModifier
        try:
            from llmcompressor.modifiers.transform.awq import AWQModifier
        except ImportError:  # older layouts keep AWQModifier under modifiers.awq
            from llmcompressor.modifiers.awq import AWQModifier
        recipe = [
            AWQModifier(duo_scaling=False),
            QuantizationModifier(scheme=spec["scheme"], ignore=QUANT_IGNORE),
        ]
    else:
        recipe = []
        if spec.get("smoothquant"):
            from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
            recipe.append(SmoothQuantModifier(smoothing_strength=0.8))
        recipe.append(GPTQModifier(targets="Linear", scheme=spec["scheme"], ignore=QUANT_IGNORE))

    oneshot(
        model=model,
        dataset=dataset,
        recipe=recipe,
        max_seq_length=CALIB_MAX_SEQ_LEN,
        num_calibration_samples=len(dataset),
        data_collator=_data_collator,
    )
    out_dir.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(out_dir, save_compressed=True)
    processor.save_pretrained(out_dir)


def _build_bnb(method: str, base_model: str, out_dir: Path) -> None:
    import torch
    from transformers import AutoProcessor, BitsAndBytesConfig

    try:
        from transformers import AutoModelForImageTextToText as _AutoModel
    except ImportError:
        from transformers import AutoModelForCausalLM as _AutoModel

    bits = METHOD_SPECS[method]["bits"]
    if bits == 8:
        config = BitsAndBytesConfig(load_in_8bit=True)
    else:
        config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True,
        )
    model = _AutoModel.from_pretrained(base_model, quantization_config=config, device_map="auto")
    model.save_pretrained(out_dir)
    AutoProcessor.from_pretrained(base_model).save_pretrained(out_dir)


def build_model(method: str, base_model: str, out_dir: Path, calib_images: List[Path]) -> Path:
    if not needs_build(method):
        return Path(base_model)
    out_dir = Path(out_dir)
    if _is_built(out_dir):
        return out_dir
    kind = METHOD_SPECS[method]["kind"]
    if kind == "compressor":
        _build_compressor(method, base_model, out_dir, calib_images)
    else:
        _build_bnb(method, base_model, out_dir)
    return out_dir