# /// script
# dependencies = [
#     "transformers>=4.36.0",
#     "peft>=0.7.0",
#     "torch>=2.0.0",
#     "accelerate>=0.24.0",
#     "huggingface_hub>=0.20.0",
#     "sentencepiece>=0.1.99",
#     "protobuf>=3.20.0",
#     "numpy",
#     "gguf",
# ]
# ///

"""GGUF Conversion for GRPO model (two-step adapter merge)"""

import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
from huggingface_hub import HfApi
import subprocess

print("🔄 GGUF Conversion Script - GRPO Model")
print("=" * 60)

BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen3-0.6B")
SFT_ADAPTER = os.environ.get("SFT_ADAPTER", "chaddy81/qwen3-0.6b-multicode-sft")
GRPO_ADAPTER = os.environ.get("GRPO_ADAPTER", "chaddy81/qwen3-0.6b-multicode-grpo")
OUTPUT_REPO = os.environ.get("OUTPUT_REPO", "chaddy81/qwen3-0.6b-multicode-grpo-gguf")
QUANT_TYPE = os.environ.get("QUANT_TYPE", "Q8_0")

print(f"\n📦 Configuration:")
print(f"   Base model: {BASE_MODEL}")
print(f"   SFT adapter: {SFT_ADAPTER}")
print(f"   GRPO adapter: {GRPO_ADAPTER}")
print(f"   Output repo: {OUTPUT_REPO}")
print(f"   Quantization: {QUANT_TYPE}")

# Step 1: Load base model
print("\n🔧 Step 1: Loading base model...")
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True,
)
print("   ✅ Base model loaded")

# Step 2: Apply and merge SFT adapter
print("\n🔧 Step 2: Applying SFT adapter...")
sft_model = PeftModel.from_pretrained(base_model, SFT_ADAPTER)
merged_sft = sft_model.merge_and_unload()
print("   ✅ SFT adapter merged")

# Step 3: Apply and merge GRPO adapter
print("\n🔧 Step 3: Applying GRPO adapter...")
grpo_model = PeftModel.from_pretrained(merged_sft, GRPO_ADAPTER)
final_model = grpo_model.merge_and_unload()
print("   ✅ GRPO adapter merged")

# Load tokenizer from GRPO adapter (has latest config)
tokenizer = AutoTokenizer.from_pretrained(GRPO_ADAPTER, trust_remote_code=True)
print("   ✅ Tokenizer loaded")

# Step 4: Save merged model
print("\n💾 Step 4: Saving merged model...")
merged_dir = "/tmp/merged_model"
final_model.save_pretrained(merged_dir, safe_serialization=True)
tokenizer.save_pretrained(merged_dir)
print("   ✅ Merged model saved")

# Step 5: Setup llama.cpp
print("\n📥 Step 5: Setting up llama.cpp...")
subprocess.run(["apt-get", "update", "-qq"], check=True, capture_output=True)
subprocess.run(["apt-get", "install", "-y", "-qq", "build-essential", "cmake"], check=True, capture_output=True)
print("   ✅ Build tools installed")

subprocess.run(["git", "clone", "--depth", "1", "https://github.com/ggerganov/llama.cpp.git", "/tmp/llama.cpp"], check=True, capture_output=True)
print("   ✅ llama.cpp cloned")

subprocess.run(["pip", "install", "-q", "-r", "/tmp/llama.cpp/requirements.txt"], check=True, capture_output=True)
subprocess.run(["pip", "install", "-q", "sentencepiece", "protobuf"], check=True, capture_output=True)
print("   ✅ Dependencies installed")

# Step 6: Convert to GGUF
print("\n🔄 Step 6: Converting to GGUF format...")
gguf_output_dir = "/tmp/gguf_output"
os.makedirs(gguf_output_dir, exist_ok=True)
model_name = GRPO_ADAPTER.split('/')[-1]
gguf_fp16 = f"{gguf_output_dir}/{model_name}-f16.gguf"

try:
    result = subprocess.run(
        ["python", "/tmp/llama.cpp/convert_hf_to_gguf.py", merged_dir, "--outfile", gguf_fp16, "--outtype", "f16"],
        check=True, capture_output=True, text=True
    )
    print("   ✅ FP16 GGUF created")
except subprocess.CalledProcessError as e:
    print(f"❌ Conversion failed! STDERR: {e.stderr[-2000:]}")
    raise

# Step 7: Quantize
print(f"\n⚙️  Step 7: Creating {QUANT_TYPE} quantization...")
os.makedirs("/tmp/llama.cpp/build", exist_ok=True)
subprocess.run(["cmake", "-B", "/tmp/llama.cpp/build", "-S", "/tmp/llama.cpp", "-DGGML_CUDA=OFF"], check=True, capture_output=True, text=True)
subprocess.run(["cmake", "--build", "/tmp/llama.cpp/build", "--target", "llama-quantize", "-j", "4"], check=True, capture_output=True, text=True)
print("   ✅ Quantize tool built")

quantize_bin = "/tmp/llama.cpp/build/bin/llama-quantize"
gguf_quant = f"{gguf_output_dir}/{model_name}-{QUANT_TYPE.lower()}.gguf"
subprocess.run([quantize_bin, gguf_fp16, gguf_quant, QUANT_TYPE], check=True, capture_output=True)
size_mb = os.path.getsize(gguf_quant) / (1024 * 1024)
print(f"   ✅ {QUANT_TYPE}: {size_mb:.1f} MB")

# Step 8: Upload to Hub
print("\n☁️  Step 8: Uploading to Hub...")
api = HfApi()
api.create_repo(repo_id=OUTPUT_REPO, repo_type="model", exist_ok=True)
print("   ✅ Repository ready")

api.upload_file(path_or_fileobj=gguf_quant, path_in_repo=f"{model_name}-{QUANT_TYPE.lower()}.gguf", repo_id=OUTPUT_REPO)
print(f"   ✅ {QUANT_TYPE} uploaded")

# Create README
readme = f"""---
base_model: {BASE_MODEL}
tags:
- gguf
- llama.cpp
- quantized
- trl
- grpo
---

# {OUTPUT_REPO.split('/')[-1]}

GGUF conversion of [{GRPO_ADAPTER}](https://huggingface.co/{GRPO_ADAPTER}).

**Training Pipeline:**
1. Base: [{BASE_MODEL}](https://huggingface.co/{BASE_MODEL})
2. SFT: [{SFT_ADAPTER}](https://huggingface.co/{SFT_ADAPTER})
3. GRPO: [{GRPO_ADAPTER}](https://huggingface.co/{GRPO_ADAPTER})

## Available Files

| File | Quant | Size |
|------|-------|------|
| {model_name}-{QUANT_TYPE.lower()}.gguf | {QUANT_TYPE} | {size_mb:.1f} MB |

## Usage

### With Ollama
```bash
huggingface-cli download {OUTPUT_REPO} {model_name}-{QUANT_TYPE.lower()}.gguf
echo "FROM ./{model_name}-{QUANT_TYPE.lower()}.gguf" > Modelfile
ollama create {model_name} -f Modelfile
ollama run {model_name}
```

### With llama.cpp
```bash
./llama-cli -m {model_name}-{QUANT_TYPE.lower()}.gguf -p "Your prompt"
```
"""
api.upload_file(path_or_fileobj=readme.encode(), path_in_repo="README.md", repo_id=OUTPUT_REPO)
print("   ✅ README uploaded")

print("\n" + "=" * 60)
print("✅ GGUF Conversion Complete!")
print(f"📦 https://huggingface.co/{OUTPUT_REPO}")
print(f"📥 huggingface-cli download {OUTPUT_REPO} {model_name}-{QUANT_TYPE.lower()}.gguf")
print("=" * 60)