from __future__ import annotations

import json
import os
import shutil
import time
from pathlib import Path

import modal

REPO = "Akahsizrr/Mini-Whale-Flash"
SUBDIR = "merged-v2-full"
VOL_NAME = "fuse2-model-store"
VOL_MOUNT = "/data"
MODEL_CACHE = f"{VOL_MOUNT}/hf-models/mini-whale-flash"
OUTPUT_ROOT = f"{VOL_MOUNT}/quantized"

vol = modal.Volume.from_name(VOL_NAME, create_if_missing=True)
image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install("torch==2.7.0", index_url="https://download.pytorch.org/whl/cu126")
    .pip_install(
        "transformers==5.14.1",
        "accelerate==1.2.1",
        "bitsandbytes==0.49.1",
        "safetensors==0.8.0",
        "huggingface_hub",
    )
    .add_local_file(str(Path(__file__).resolve().parent.parent / "microscope" / "fuse2_model.py"), "/root/fuse2_model.py")
)
app = modal.App("fuse2-quantize")


def _quantization_config(bits: int, quant_type: str = "nf4"):
    from transformers import BitsAndBytesConfig

    if bits == 8:
        return BitsAndBytesConfig(load_in_8bit=True)
    if bits == 4 and quant_type in {"nf4", "fp4"}:
        return BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type=quant_type,
            bnb_4bit_compute_dtype=__import__("torch").bfloat16,
            bnb_4bit_use_double_quant=True,
        )
    raise ValueError("supported formats are 4-bit nf4, 4-bit fp4, and 8-bit int8")


def _apply_runtime_fixes(model):
    import torch
    import torch.nn.functional as F

    scaled_count = 0
    for layer in model.model.layers:
        experts = getattr(layer, "experts", None)
        if experts is None:
            continue
        for expert in experts:
            gate_proj = getattr(expert, "gate_proj", None)
            if gate_proj is None:
                continue
            weight = gate_proj.weight
            if hasattr(weight, "dequantize"):
                weight = weight.dequantize()
            std_val = weight.float().std().item()
            scale = 0.025 / std_val if std_val > 1.0 else 1.0
            if scale != 1.0:
                scaled_count += 1

            expert._fuse2_scale = scale

    patched_count = 0
    for layer in model.model.layers:
        experts = getattr(layer, "experts", None)
        if experts is None:
            continue
        for expert in experts:
            gate_proj = expert.gate_proj
            up_proj = expert.up_proj
            down_proj = expert.down_proj
            scale = getattr(expert, "_fuse2_scale", 1.0)

            def clamped_forward(x, gp=gate_proj, up=up_proj, dp=down_proj, s=scale):
                value = torch.clamp(F.silu(gp(x) * s) * (up(x) * s), -10.0, 10.0)
                return dp(value) * s

            expert.forward = clamped_forward
            patched_count += 1

    if not patched_count:
        raise RuntimeError("Fuse-2 experts were not found")
    return scaled_count, patched_count


@app.function(
    image=image,
    gpu="A100",
    cpu=8,
    memory=65536,
    timeout=7200,
    volumes={VOL_MOUNT: vol},
)
def quantize(bits: int = 4, quant_type: str = "nf4", force: bool = False):
    import torch
    from huggingface_hub import snapshot_download
    from transformers import AutoModelForCausalLM, AutoTokenizer

    if bits not in (4, 8):
        raise ValueError("bits must be 4 or 8")
    if bits == 8:
        quant_type = "int8"
        output_name = "fuse2-8bit-bnb"
    elif quant_type == "nf4":
        output_name = "fuse2-4bit-bnb"
    else:
        output_name = f"fuse2-4bit-{quant_type}-bnb"
    output_dir = Path(OUTPUT_ROOT) / output_name
    marker = output_dir / "quantization_metadata.json"
    if marker.exists() and not force:
        return {"status": "exists", "path": str(output_dir), "bits": bits}

    model_path = Path(
        snapshot_download(
            REPO,
            allow_patterns=[f"{SUBDIR}/*"],
            local_dir=MODEL_CACHE,
        )
    ) / SUBDIR
    shutil.copy2("/root/fuse2_model.py", model_path / "fuse2_model.py")
    output_dir.mkdir(parents=True, exist_ok=True)

    started = time.time()
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        trust_remote_code=True,
        quantization_config=_quantization_config(bits, quant_type),
        device_map="cuda:0",
        torch_dtype=torch.bfloat16,
    )
    model.eval()
    scaled_experts, patched_experts = _apply_runtime_fixes(model)

    text = tokenizer.apply_chat_template(
        [{"role": "user", "content": "Say hello in one sentence."}],
        tokenize=False,
        add_generation_prompt=True,
    )
    inputs = tokenizer(text, return_tensors="pt").to("cuda:0")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=16,
            do_sample=False,
            use_cache=False,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    sample = tokenizer.decode(
        outputs[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True
    )
    if not sample and outputs.shape[-1] <= inputs["input_ids"].shape[-1]:
        raise RuntimeError("quantized model generated no tokens")

    model.save_pretrained(
        output_dir,
        safe_serialization=True,
        max_shard_size="4GB",
    )
    tokenizer.save_pretrained(output_dir)
    source_code = model_path / "fuse2_model.py"
    if source_code.exists():
        shutil.copy2(source_code, output_dir / "fuse2_model.py")

    metadata = {
        "base_model": REPO,
        "base_subdir": SUBDIR,
        "quantization": "bitsandbytes",
        "bits": bits,
        "quant_type": quant_type,
        "compute_dtype": "bfloat16",
        "gpu_validation": "NVIDIA A100-SXM4-40GB",
        "runtime_fixes": {
            "scaled_shared_experts": scaled_experts,
            "clamped_experts": patched_experts,
        },
        "sample_output": sample,
        "parameter_count": sum(p.numel() for p in model.parameters()),
        "elapsed_seconds": round(time.time() - started, 1),
    }
    (output_dir / "quantization_metadata.json").write_text(
        json.dumps(metadata, indent=2) + "\n", encoding="utf-8"
    )
    vol.commit()
    return {"status": "created", "path": str(output_dir), **metadata}


@app.function(
    image=image,
    gpu="A100",
    cpu=8,
    memory=65536,
    timeout=7200,
    volumes={VOL_MOUNT: vol},
)
def validate(bits: int = 4, quant_type: str = "nf4"):
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    if bits not in (4, 8):
        raise ValueError("bits must be 4 or 8")
    if bits == 8:
        quant_type = "int8"
        output_name = "fuse2-8bit-bnb"
    elif quant_type == "nf4":
        output_name = "fuse2-4bit-bnb"
    else:
        output_name = f"fuse2-4bit-{quant_type}-bnb"
    output_dir = Path(OUTPUT_ROOT) / output_name
    if not (output_dir / "quantization_metadata.json").exists():
        raise FileNotFoundError(output_dir)

    tokenizer = AutoTokenizer.from_pretrained(output_dir, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        output_dir,
        trust_remote_code=True,
        quantization_config=_quantization_config(bits, quant_type),
        device_map="cuda:0",
        torch_dtype=torch.bfloat16,
    )
    model.eval()
    scaled_experts, patched_experts = _apply_runtime_fixes(model)
    text = tokenizer.apply_chat_template(
        [{"role": "user", "content": "Say hello in one sentence."}],
        tokenize=False,
        add_generation_prompt=True,
    )
    inputs = tokenizer(text, return_tensors="pt").to("cuda:0")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=16,
            do_sample=False,
            use_cache=False,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    sample = tokenizer.decode(
        outputs[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True
    )
    if not sample:
        raise RuntimeError("saved quantized artifact generated no text")
    result = {
        "status": "validated",
        "path": str(output_dir),
        "bits": bits,
        "sample_output": sample,
        "scaled_shared_experts": scaled_experts,
        "clamped_experts": patched_experts,
    }
    (output_dir / "validation.json").write_text(
        json.dumps(result, indent=2) + "\n", encoding="utf-8"
    )
    vol.commit()
    return result


@app.local_entrypoint()
def main(
    action: str = "quantize",
    bits: int = 4,
    quant_type: str = "nf4",
    force: bool = False,
):
    if action == "quantize":
        result = quantize.remote(bits=bits, quant_type=quant_type, force=force)
    elif action == "validate":
        result = validate.remote(bits=bits, quant_type=quant_type)
    else:
        raise ValueError("action must be quantize or validate")
    print(json.dumps(result, indent=2))