#!/bin/bash

echo "============================================"
echo "  Unsloth Fine-Tune Setup"
echo "============================================"

# Create main virtual environment
echo ""
echo "Creating main virtual environment..."
python -m venv venv
source venv/bin/activate
pip install -r requirements.txt

# Prompt for backend
echo ""
echo "Select llama.cpp backend:"
echo "  1) CUDA (NVIDIA GPU)"
echo "  2) ROCm (AMD GPU)"
echo "  3) Vulkan (Cross-vendor GPU)"
echo "  4) CPU only"
echo ""
read -p "Enter choice (1-4): " BACKEND

# Clone llama.cpp
echo ""
echo "Cloning llama.cpp..."
if [ ! -d "llama.cpp" ]; then
    git clone https://github.com/ggml-org/llama.cpp.git
else
    echo "llama.cpp already exists, skipping clone."
fi

# Build llama.cpp with correct flags
echo ""
echo "Building llama.cpp..."
cd llama.cpp

BUILD_FAILED=0

case $BACKEND in
    1)
        echo "Building with CUDA support..."
        cmake -B build -DGGML_CUDA=ON || BUILD_FAILED=1
        [ $BUILD_FAILED -eq 0 ] && cmake --build build --config Release -j$(nproc) || BUILD_FAILED=1
        ;;
    2)
        echo "Building with ROCm support..."
        HIPCXX="$(hipconfig -l)/clang" HIP_PATH="$(hipconfig -R)" \
            cmake -S . -B build -DGGML_HIP=ON -DGPU_TARGETS=gfx1030 -DCMAKE_BUILD_TYPE=Release || BUILD_FAILED=1
        [ $BUILD_FAILED -eq 0 ] && cmake --build build --config Release -j$(nproc) || BUILD_FAILED=1
        ;;
    3)
        echo "Building with Vulkan support..."
        cmake -B build -DGGML_VULKAN=1 || BUILD_FAILED=1
        [ $BUILD_FAILED -eq 0 ] && cmake --build build --config Release -j$(nproc) || BUILD_FAILED=1
        ;;
    4)
        echo "Building CPU-only..."
        cmake -B build || BUILD_FAILED=1
        [ $BUILD_FAILED -eq 0 ] && cmake --build build --config Release -j$(nproc) || BUILD_FAILED=1
        ;;
    *)
        echo "Invalid choice. Building CPU-only."
        cmake -B build || BUILD_FAILED=1
        [ $BUILD_FAILED -eq 0 ] && cmake --build build --config Release -j$(nproc) || BUILD_FAILED=1
        ;;
esac

cd ..

# Install llama-cpp-python in main venv
echo ""
echo "Installing llama-cpp-python..."

source venv/bin/activate

case $BACKEND in
    1) CMAKE_ARGS="-DGGML_CUDA=on" ;;
    2) CMAKE_ARGS="-DGGML_HIP=on" ;;
    3) CMAKE_ARGS="-DGGML_VULKAN=on" ;;
    *) CMAKE_ARGS="" ;;
esac

eval "CMAKE_ARGS=\"$CMAKE_ARGS\" pip install llama-cpp-python"

# Create convertgguf_venv for llama.cpp Python tools
echo ""
echo "Creating convertgguf_venv..."
python -m venv llama.cpp/convertgguf_venv
source llama.cpp/convertgguf_venv/bin/activate
pip install -r llama.cpp/requirements.txt

echo ""
echo "Setup complete! Configure the scripts and run:"
echo "  bash run-pipeline.sh"

if [ $BUILD_FAILED -ne 0 ]; then
    echo ""
    echo "Build failed. See the build guide for help:"
    echo "  https://github.com/ggml-org/llama.cpp/blob/master/docs/build.md"
fi