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

export PYTHONPATH=$HOME/Desktop/llama.cpp/gguf-py

# ──────────────────────────────────────────────────────────────
# Verify GGUF files
# - Metadata
# - Tensor inventory (optional, skips if missing)
# - SHA256 hash
# - Smoke-test inference
# - Perplexity (optional)
# Usage:
#   ./verify_gguf.sh [--perplexity] [file1.gguf file2.gguf ...]
# ──────────────────────────────────────────────────────────────

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
LLAMA_CPP_DIR="${LLAMA_CPP_DIR:-$HOME/Desktop/llama.cpp}"
OUTPUT_DIR="${SCRIPT_DIR}/gguf_output"
GGUF_PY_DIR="$LLAMA_CPP_DIR/gguf-py"
RUN_PERPLEXITY=false
TARGETS=()

# Add gguf-py to PYTHONPATH
export PYTHONPATH="$GGUF_PY_DIR:$PYTHONPATH"

# Detect number of CPU cores
if command -v nproc >/dev/null 2>&1; then
    NPROC=$(nproc)
else
    NPROC=$(sysctl -n hw.ncpu)
fi

# ─── Parse arguments ─────────────────────────────────────────────
while [[ $# -gt 0 ]]; do
    case "$1" in
        --perplexity|--ppl|--all)
            RUN_PERPLEXITY=true
            shift
            ;;
        --llama-cpp)
            LLAMA_CPP_DIR="$2"
            GGUF_PY_DIR="$LLAMA_CPP_DIR/gguf-py"
            shift 2
            ;;
        *)
            TARGETS+=("$1")
            shift
            ;;
    esac
done

# ─── If no targets, find all GGUFs in output dir ───────────────
if [ ${#TARGETS[@]} -eq 0 ]; then
    while IFS= read -r f; do
        TARGETS+=("$f")
    done < <(find "$OUTPUT_DIR" -name "*.gguf" -type f | sort)
fi

if [ ${#TARGETS[@]} -eq 0 ]; then
    echo "❌ No GGUF files found in $OUTPUT_DIR"
    exit 1
fi

# ─── Build llama-cli / llama-perplexity if missing ─────────────
CLI_BIN="$LLAMA_CPP_DIR/build/bin/llama-cli"
PPL_BIN="$LLAMA_CPP_DIR/build/bin/llama-perplexity"

build_target() {
    local target="$1"
    local bin="$LLAMA_CPP_DIR/build/bin/$target"
    if [ ! -f "$bin" ]; then
        echo "  Building $target..."
        pushd "$LLAMA_CPP_DIR" > /dev/null
        cmake -B build -DCMAKE_BUILD_TYPE=Release 2>&1 | tail -1
        cmake --build build --target "$target" -j"$NPROC" 2>&1 | tail -3
        popd > /dev/null
    fi
}

build_target "llama-cli"
if $RUN_PERPLEXITY; then
    build_target "llama-perplexity"
fi

# ─── Begin verification ────────────────────────────────────────
PASS=0
FAIL=0

echo "============================================="
echo " GGUF Verification"
echo "============================================="
echo " Files to check: ${#TARGETS[@]}"
echo " Perplexity:     $RUN_PERPLEXITY"
echo "============================================="
echo ""

for GGUF_FILE in "${TARGETS[@]}"; do
    BASENAME=$(basename "$GGUF_FILE")
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "  $BASENAME"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

    # Skip empty or missing files
    if [ ! -s "$GGUF_FILE" ]; then
        echo "  ⚠️  Skipping empty or missing file"
        FAIL=$((FAIL+1))
        continue
    fi

    FILE_SIZE=$(du -h "$GGUF_FILE" | cut -f1)
    echo "  Size: $FILE_SIZE"
    echo ""

    FILE_OK=true

    # ── Check 1: Metadata ─────────────────────────────
    echo "  ┌─ [1/4] Metadata"
    if METADATA=$(python3 "$GGUF_PY_DIR/gguf/scripts/gguf_dump.py" "$GGUF_FILE" --no-tensors 2>&1); then
        ARCH=$(echo "$METADATA" | sed -n 's/.*general\.architecture *= *\([^ ]*\).*/\1/p' | head -1)
        QTYPE=$(echo "$METADATA" | sed -n 's/.*general\.file_type *= *\(.*\)/\1/p' | head -1)
        CTX=$(echo "$METADATA" | sed -n 's/.*context_length *= *\([^ ]*\).*/\1/p' | head -1)
        NLAYER=$(echo "$METADATA" | sed -n 's/.*block_count *= *\([^ ]*\).*/\1/p' | head -1)
        VOCAB=$(echo "$METADATA" | sed -n 's/.*tokens.*length *= *\([^ ]*\).*/\1/p' | head -1)
        echo "  │   Architecture:    ${ARCH:-?}"
        echo "  │   File type:       ${QTYPE:-?}"
        echo "  │   Context length:  ${CTX:-?}"
        echo "  │   Layers:          ${NLAYER:-?}"
        echo "  │   Vocab size:      ${VOCAB:-?}"
        echo "  │   ✅ Metadata OK"
    else
        echo "  │   ❌ Failed to read metadata"
        echo "  │   $METADATA" | head -5
        FILE_OK=false
        # Skip tensor check if metadata fails
        SKIP_TENSORS=true
    fi
    echo "  │"

    # ── Check 2: Tensor inventory ─────────────────────
    echo "  ├─ [2/4] Tensor inventory"
    if [ "${SKIP_TENSORS:-false}" = true ]; then
        echo "  │   ⚠️  Skipping tensors due to missing metadata"
    else
        if TENSOR_INFO=$(python3 "$GGUF_PY_DIR/gguf/scripts/gguf_dump.py" "$GGUF_FILE" --json 2>/dev/null); then
            if [ -z "$TENSOR_INFO" ]; then
                echo "  │   ⚠️  Tensor info empty, skipping"
            else
                TENSOR_COUNT=$(echo "$TENSOR_INFO" | python3 -c "
import sys, json
data = json.load(sys.stdin)
tensors = data.get('tensors', [])
print(len(tensors))
" 2>/dev/null || echo "?")

                echo "  │   Total tensors: ${TENSOR_COUNT:-?}"
                echo "  │   ✅ Tensors OK"
            fi
        else
            echo "  │   ⚠️  Failed to parse tensor JSON, skipping"
        fi
    fi
    echo "  │"

    # ── Check 3: SHA256 hash ──────────────────────────
    echo "  ├─ [3/4] SHA256 hash"
    if command -v shasum >/dev/null 2>&1; then
        HASH=$(shasum -a 256 "$GGUF_FILE" | cut -d' ' -f1)
    else
        HASH=$(sha256sum "$GGUF_FILE" | cut -d' ' -f1)
    fi
    echo "  │   $HASH"
    echo "  │   ✅ Hash computed"
    echo "$HASH  $BASENAME" >> "$OUTPUT_DIR/SHA256SUMS"
    echo "  │"

    # ── Check 4: Smoke-test inference ───────────────
    echo "  └─ [4/4] Smoke-test inference"
    if [ ! -s "$GGUF_FILE" ]; then
        echo "  ⚠️  File empty or missing weights, skipping inference"
        FILE_OK=false
    else
        SMOKE_OUTPUT=$("$CLI_BIN" -m "$GGUF_FILE" -p "The capital of France is" -n 20 -r "" --no-display-prompt --seed 42 --temp 0.0 --threads 4 --log-disable 2>&1) || FILE_OK=false
        TOKEN_COUNT=$(echo "$SMOKE_OUTPUT" | wc -w)
        PREVIEW=$(echo "$SMOKE_OUTPUT" | head -3 | sed 's/^/  │   > /')
        echo "$PREVIEW"
        if [ "$TOKEN_COUNT" -gt 0 ]; then
            echo "  │   ✅ Generated $TOKEN_COUNT words"
        else
            echo "  │   ⚠️  Empty output"
        fi
    fi
    echo ""

    # ── Summary per file ─────────────────────────────
    if $FILE_OK; then
        echo "  ✅ $BASENAME — ALL CHECKS PASSED"
        PASS=$((PASS + 1))
    else
        echo "  ❌ $BASENAME — SOME CHECKS FAILED"
        FAIL=$((FAIL + 1))
    fi
    echo ""
done

# ─── Final summary ─────────────────────────────────────────────
echo "============================================="
echo " Verification Summary"
echo "============================================="
echo " Passed: $PASS"
echo " Failed: $FAIL"
echo " Total:  ${#TARGETS[@]}"
if [ -f "$OUTPUT_DIR/SHA256SUMS" ]; then
    sort -u -o "$OUTPUT_DIR/SHA256SUMS" "$OUTPUT_DIR/SHA256SUMS"
    echo ""
    echo " SHA256 checksums saved to:"
    echo "   $OUTPUT_DIR/SHA256SUMS"
fi
echo "============================================="

exit $FAIL