import traceback
import torch
from datasets import load_dataset, concatenate_datasets
from transformers import AutoTokenizer, AutoProcessor, Qwen3_5ForConditionalGeneration
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier
from compressed_tensors.quantization import QuantizationArgs
from compressed_tensors.utils import save_mtp_tensors_to_checkpoint

# NOTE: requires transformers >= v5.9 and llm-compressor >= 0.13.0
# (0.13.0 includes the observer fusion/deletion fix that caused
#  NVFP4 weight corruption when combined with kv_cache_scheme)

MODEL_ID = "Qwen/Qwen3.8-27B"

model = Qwen3_5ForConditionalGeneration.from_pretrained(MODEL_ID, dtype="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
processor = AutoProcessor.from_pretrained(MODEL_ID)  # only for saving

recipe = QuantizationModifier(
    targets="Linear",
    scheme="NVFP4",
    ignore=[
        "lm_head",
        "re:.*visual.*",
        "re:.*linear_attn.*",
    ],
    kv_cache_scheme=QuantizationArgs(
        num_bits=8,
        type="float",
        strategy="tensor",
        dynamic=False,
        symmetric=True,
    ),
)

NUM_CALIBRATION_SAMPLES = 1024
MAX_SEQUENCE_LENGTH = 8192
samples_per_split = NUM_CALIBRATION_SAMPLES // 4  # 256 per domain

# ============================================================
# 1. General conversation
# ============================================================
ds_chat = load_dataset(
    "HuggingFaceH4/ultrachat_200k",
    split=f"train_sft[:{samples_per_split}]",
)

def preprocess_chat(example):
    text = tokenizer.apply_chat_template(example["messages"], tokenize=False)
    return {"text": text}

ds_chat = ds_chat.map(preprocess_chat).select_columns(["text"])

# ============================================================
# 2. Math / reasoning
# ============================================================
ds_math = load_dataset(
    "openai/gsm8k", "main",
    split=f"train[:{samples_per_split}]",
)

def preprocess_math(example):
    messages = [
        {"role": "user", "content": example["question"]},
        {"role": "assistant", "content": example["answer"]},
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False)
    return {"text": text}

ds_math = ds_math.map(preprocess_math).select_columns(["text"])

# ============================================================
# 3. Code
# ============================================================
ds_code = load_dataset(
    "sahil2801/CodeAlpaca-20k",
    split=f"train[:{samples_per_split}]",
)

def preprocess_code(example):
    user_content = example["instruction"]
    if example.get("input"):
        user_content += "\n\n" + example["input"]
    messages = [
        {"role": "user", "content": user_content},
        {"role": "assistant", "content": example["output"]},
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False)
    return {"text": text}

ds_code = ds_code.map(preprocess_code).select_columns(["text"])

# ============================================================
# 4. Multilingual
# ============================================================
ds_multi = load_dataset(
    "CohereForAI/aya_dataset",
    split=f"train[:{samples_per_split}]",
)

def preprocess_multi(example):
    messages = [
        {"role": "user", "content": example["inputs"]},
        {"role": "assistant", "content": example["targets"]},
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False)
    return {"text": text}

ds_multi = ds_multi.map(preprocess_multi).select_columns(["text"])

# ============================================================
# Combine, shuffle, filter
# ============================================================
ds = concatenate_datasets([ds_chat, ds_math, ds_code, ds_multi])
ds = ds.shuffle(seed=42)
ds = ds.filter(lambda x: len(x["text"].strip()) > 0)

def tokenize(sample):
    return tokenizer(
        sample["text"],
        padding=False,
        max_length=MAX_SEQUENCE_LENGTH,
        truncation=True,
        add_special_tokens=False,
    )

ds = ds.map(tokenize, remove_columns=ds.column_names)

# ============================================================
# Patch: attention config from text_config to top-level
# ============================================================
text_cfg = model.config.text_config

for attr in [
    "num_attention_heads",
    "num_key_value_heads",
    "hidden_size",
    "head_dim",
]:
    if not hasattr(model.config, attr) and hasattr(text_cfg, attr):
        setattr(model.config, attr, getattr(text_cfg, attr))

# Apply quantization
oneshot(
    model=model,
    recipe=recipe,
    dataset=ds,
    max_seq_length=MAX_SEQUENCE_LENGTH,
    num_calibration_samples=NUM_CALIBRATION_SAMPLES,
)

# ============================================================
# SAVE IMMEDIATELY — nothing risky between calibration and disk.
# Hours of calibration must not be lost to a downstream exception.
# ============================================================
SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4-FP8KV-MTP"
model.save_pretrained(SAVE_DIR, safe_serialization=True)
processor.save_pretrained(SAVE_DIR)
save_mtp_tensors_to_checkpoint(source_model=MODEL_ID, dest_dir=SAVE_DIR)
print(f"Saved to {SAVE_DIR}")

# ============================================================
# Sanity check generation — AFTER save, wrapped so it can never
# crash the script or be mistaken for a reason to skip saving.
# ============================================================
try:
    from compressed_tensors.offload import dispatch_model

    print("\n=== Sanity check generation (post-save, informational only) ===")
    print("Dispatching model to GPU (may take a few minutes for 27B)...")
    dispatch_model(model)

    test_prompt = "Explain in one paragraph what NVFP4 quantization is."
    inputs = tokenizer(test_prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        output = model.generate(**inputs, max_new_tokens=64, do_sample=False)
    print(tokenizer.decode(output[0], skip_special_tokens=True))
except Exception:
    print("\n[WARN] Post-save sanity generation failed. The saved checkpoint "
          "on disk is unaffected — inspect it manually with a separate script.")
    traceback.print_exc()