#!/usr/bin/env bash
set -euo pipefail

#=============================================================================
# Convert Jackrong/Qwen3.5-27B-Claude-4.6-Opus-Reasoning-Distilled to GGUF
#
# Produces ALL quantization variants from a single F16 GGUF base.
#
# Prerequisites:
#   - Python 3.10+
#   - pip packages: numpy torch transformers sentencepiece safetensors gguf protobuf
#   - llama.cpp repo cloned (for convert_hf_to_gguf.py and llama-quantize)
#
# Usage:
#   ./convert_to_gguf.sh                        # all quants, default llama.cpp path
#   ./convert_to_gguf.sh /path/to/llama.cpp     # all quants, custom llama.cpp path
#=============================================================================

MODEL_REPO="Jackrong/Qwen3.5-27B-Claude-4.6-Opus-Reasoning-Distilled"
MODEL_NAME="Qwen3.5-27B-Claude-4.6-Opus-Reasoning-Distilled"
LLAMA_CPP_DIR="${1:-$HOME/Desktop/llama.cpp}"
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
OUTPUT_DIR="${SCRIPT_DIR}/gguf_output"

# ─── Every quantization type supported by llama-quantize ─────────────────────
# Sorted roughly by bits-per-weight (smallest → largest)
QUANT_TYPES=(
    # ~1-2 bpw — extreme compression
    IQ1_S        # 1.56 bpw
    IQ1_M        # 1.75 bpw
    TQ1_0        # 1.69 bpw ternary
    TQ2_0        # 2.06 bpw ternary
    IQ2_XXS      # 2.06 bpw
    IQ2_XS       # 2.31 bpw
    IQ2_S        # 2.5  bpw
    IQ2_M        # 2.7  bpw

    # ~3 bpw
    Q2_K         # 2.96G @ 8B, +3.52 ppl
    Q2_K_S       # 2.96G @ 8B, +3.18 ppl
    IQ3_XXS      # 3.06 bpw
    IQ3_XS       # 3.3  bpw
    IQ3_S        # 3.44 bpw
    IQ3_M        # 3.66 bpw
    Q3_K_S       # 3.41G @ 8B, +1.63 ppl
    Q3_K_M       # 3.74G @ 8B, +0.66 ppl
    Q3_K_L       # 4.03G @ 8B, +0.56 ppl

    # ~4 bpw — sweet spot
    IQ4_XS       # 4.25 bpw
    IQ4_NL       # 4.50 bpw
    Q4_0         # 4.34G @ 8B, +0.47 ppl
    Q4_1         # 4.78G @ 8B, +0.45 ppl
    Q4_K_S       # 4.37G @ 8B, +0.27 ppl
    Q4_K_M       # 4.58G @ 8B, +0.18 ppl  ← popular default

    # ~5 bpw
    Q5_0         # 5.21G @ 8B, +0.13 ppl
    Q5_1         # 5.65G @ 8B, +0.11 ppl
    Q5_K_S       # 5.21G @ 8B, +0.10 ppl
    Q5_K_M       # 5.33G @ 8B, +0.06 ppl

    # ~6-8 bpw — high fidelity
    Q6_K         # 6.14G @ 8B, +0.02 ppl
    Q8_0         # 7.96G @ 8B, +0.003 ppl

    # Full precision (format conversion only)
    BF16         # 14G @ 7B
    F16          # 14G @ 7B
    F32          # 26G @ 7B
)

echo "============================================="
echo " HuggingFace → GGUF Conversion Script"
echo "============================================="
echo " Model:       $MODEL_REPO"
echo " llama.cpp:   $LLAMA_CPP_DIR"
echo " Output dir:  $OUTPUT_DIR"
echo " Quant types: ${#QUANT_TYPES[@]} variants"
echo "============================================="

#-----------------------------------------------------------------------------
# Step 0: Validate llama.cpp directory
#-----------------------------------------------------------------------------
if [ ! -f "$LLAMA_CPP_DIR/convert_hf_to_gguf.py" ]; then
    echo "❌ convert_hf_to_gguf.py not found in $LLAMA_CPP_DIR"
    echo "   Clone llama.cpp first:  git clone https://github.com/ggml-org/llama.cpp.git"
    exit 1
fi

#-----------------------------------------------------------------------------
# Step 1: Resolve model from HF cache, local dir, or download
#-----------------------------------------------------------------------------
echo ""
echo "▶ Step 1: Resolving model files..."

HF_CACHE_DIR="$HOME/.cache/huggingface/hub"
MODEL_CACHE="$HF_CACHE_DIR/models--$(echo "$MODEL_REPO" | tr '/' '--')"
MODEL_DIR=""

# Helper: check if a directory has config.json + at least one .safetensors
has_model_files() {
    local dir="$1"
    [ -f "$dir/config.json" ] && \
    [ "$(find "$dir" -maxdepth 1 -name '*.safetensors' 2>/dev/null | head -1)" != "" ]
}

# 1) Check HF cache snapshots
if [ -d "$MODEL_CACHE/snapshots" ]; then
    SNAPSHOT=$(ls -t "$MODEL_CACHE/snapshots" 2>/dev/null | head -1)
    if [ -n "$SNAPSHOT" ]; then
        CANDIDATE="$MODEL_CACHE/snapshots/$SNAPSHOT"
        if has_model_files "$CANDIDATE"; then
            MODEL_DIR="$CANDIDATE"
            echo "  Found complete model in HF cache: $MODEL_DIR"
        else
            echo "  HF cache exists but is incomplete (missing safetensors or config.json)"
        fi
    fi
fi

# 2) Check local model_hf directory (from a previous download)
if [ -z "$MODEL_DIR" ] && has_model_files "${SCRIPT_DIR}/model_hf"; then
    MODEL_DIR="${SCRIPT_DIR}/model_hf"
    echo "  Found complete model in local dir: $MODEL_DIR"
fi

# 3) Nothing found — download everything
if [ -z "$MODEL_DIR" ]; then
    MODEL_DIR="${SCRIPT_DIR}/model_hf"
    echo "  ⬇  No complete model found locally. Downloading from HuggingFace..."
    echo "     Repo: $MODEL_REPO"
    echo "     Dest: $MODEL_DIR"
    echo ""
    python3 -c "
from huggingface_hub import snapshot_download
path = snapshot_download(
    '$MODEL_REPO',
    local_dir='$MODEL_DIR',
    resume_download=True,
)
print(f'  ✅ Downloaded to: {path}')
"
    echo ""

    # Verify the download actually worked
    if ! has_model_files "$MODEL_DIR"; then
        echo "❌ Download finished but model files are still missing."
        echo "   Expected config.json and *.safetensors in: $MODEL_DIR"
        echo ""
        echo "   Contents:"
        ls -la "$MODEL_DIR" 2>/dev/null || echo "   (directory does not exist)"
        exit 1
    fi
fi

SAFETENSOR_COUNT=$(find "$MODEL_DIR" -maxdepth 1 -name "*.safetensors" 2>/dev/null | wc -l)
TOTAL_SIZE=$(du -sh "$MODEL_DIR"/*.safetensors 2>/dev/null | tail -1 | cut -f1)
echo "  Found config.json and $SAFETENSOR_COUNT safetensor shard(s) (~${TOTAL_SIZE:-?} total)"

#-----------------------------------------------------------------------------
# Step 2: Convert to F16 GGUF (base for all quantizations)
#-----------------------------------------------------------------------------
echo ""
echo "▶ Step 2: Converting safetensors → GGUF (F16 base)..."

mkdir -p "$OUTPUT_DIR"
F16_GGUF="$OUTPUT_DIR/${MODEL_NAME}-F16.gguf"

if [ -f "$F16_GGUF" ]; then
    echo "  F16 GGUF already exists, skipping conversion."
else
    python3 "$LLAMA_CPP_DIR/convert_hf_to_gguf.py" \
        "$MODEL_DIR" \
        --outfile "$F16_GGUF" \
        --outtype f16

    echo "  ✅ Created: $F16_GGUF"
fi

F16_SIZE=$(du -h "$F16_GGUF" | cut -f1)
echo "  F16 size: $F16_SIZE"

#-----------------------------------------------------------------------------
# Step 3: Build llama-quantize if needed
#-----------------------------------------------------------------------------
echo ""
echo "▶ Step 3: Ensuring llama-quantize is built..."

QUANTIZE_BIN="$LLAMA_CPP_DIR/build/bin/llama-quantize"
if [ ! -f "$QUANTIZE_BIN" ]; then
    echo "  Building llama-quantize..."
    pushd "$LLAMA_CPP_DIR" > /dev/null
    cmake -B build -DCMAKE_BUILD_TYPE=Release 2>&1 | tail -3
    cmake --build build --target llama-quantize -j"$(nproc)" 2>&1 | tail -5
    popd > /dev/null

    if [ ! -f "$QUANTIZE_BIN" ]; then
        echo "❌ Failed to build llama-quantize"
        exit 1
    fi
fi
echo "  ✅ llama-quantize ready: $QUANTIZE_BIN"

#-----------------------------------------------------------------------------
# Step 4: Quantize ALL variants
#-----------------------------------------------------------------------------
echo ""
echo "▶ Step 4: Quantizing ${#QUANT_TYPES[@]} variants..."
echo ""

SUCCEEDED=0
FAILED=0
SKIPPED=0
FAILED_LIST=()

for QTYPE in "${QUANT_TYPES[@]}"; do
    QUANT_GGUF="$OUTPUT_DIR/${MODEL_NAME}-${QTYPE}.gguf"

    # Skip F16 since that's our base file
    if [ "$QTYPE" = "F16" ]; then
        echo "  [SKIP] $QTYPE — already created as base"
        SKIPPED=$((SKIPPED + 1))
        continue
    fi

    # Skip if already exists
    if [ -f "$QUANT_GGUF" ]; then
        SIZE=$(du -h "$QUANT_GGUF" | cut -f1)
        echo "  [SKIP] $QTYPE — already exists ($SIZE)"
        SKIPPED=$((SKIPPED + 1))
        continue
    fi

    echo -n "  [QUANT] $QTYPE ... "
    if "$QUANTIZE_BIN" "$F16_GGUF" "$QUANT_GGUF" "$QTYPE" > "$OUTPUT_DIR/.quantize_${QTYPE}.log" 2>&1; then
        SIZE=$(du -h "$QUANT_GGUF" | cut -f1)
        echo "✅ ($SIZE)"
        SUCCEEDED=$((SUCCEEDED + 1))
    else
        echo "❌ FAILED (see $OUTPUT_DIR/.quantize_${QTYPE}.log)"
        FAILED=$((FAILED + 1))
        FAILED_LIST+=("$QTYPE")
    fi
done

#-----------------------------------------------------------------------------
# Summary
#-----------------------------------------------------------------------------
echo ""
echo "============================================="
echo " ✅ Conversion complete!"
echo "============================================="
echo ""
echo " Results: $SUCCEEDED succeeded, $SKIPPED skipped, $FAILED failed"
if [ ${#FAILED_LIST[@]} -gt 0 ]; then
    echo " Failed:  ${FAILED_LIST[*]}"
fi
echo ""
echo " Output directory: $OUTPUT_DIR"
echo ""
echo " Files:"
echo " ─────"
ls -lhS "$OUTPUT_DIR"/*.gguf 2>/dev/null | awk '{printf "   %-12s %s\n", $5, $NF}'
echo ""
echo " Test with llama.cpp:"
echo "   $LLAMA_CPP_DIR/build/bin/llama-cli -m $OUTPUT_DIR/${MODEL_NAME}-Q4_K_M.gguf -p 'Hello!' -n 128"
echo "============================================="