#!/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

# Ask if fresh build or existing
echo ""
echo "Would you like to:"
echo "  1) Clone and build a fresh copy of llama.cpp"
echo "  2) Use an existing llama.cpp build"
echo ""
read -p "Enter choice (1-2): " BUILD_CHOICE

LLAMA_CPP_PATH=""

if [ "$BUILD_CHOICE" = "1" ]; then
    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

    echo ""
    echo "Cloning llama.cpp..."
    git clone https://github.com/ggml-org/llama.cpp.git

    # 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 -DBUILD_SHARED_LIBS=ON -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..."
            read -p "Enter GPU target (e.g., gfx1030, gfx942, gfx1100): " ROCM_TARGET
            [ -z "$ROCM_TARGET" ] && ROCM_TARGET="gfx1030"
            HIPCXX="$(hipconfig -l)/clang" HIP_PATH="$(hipconfig -R)" \
                cmake -S . -B build -DBUILD_SHARED_LIBS=ON -DGGML_HIP=ON -DGPU_TARGETS="$ROCM_TARGET" -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 -DBUILD_SHARED_LIBS=ON -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 -DBUILD_SHARED_LIBS=ON || 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 -DBUILD_SHARED_LIBS=ON || BUILD_FAILED=1
            [ $BUILD_FAILED -eq 0 ] && cmake --build build --config Release -j$(nproc) || BUILD_FAILED=1
            ;;
    esac

    LLAMA_CPP_PATH=$(pwd)
    cd ..
else
    read -p "Enter path to existing llama.cpp build: " LLAMA_CPP_PATH

    # Validate path
    if [ ! -d "$LLAMA_CPP_PATH" ]; then
        echo "Error: Directory $LLAMA_CPP_PATH does not exist."
        exit 1
    fi

    # Resolve to absolute path
    LLAMA_CPP_PATH=$(realpath "$LLAMA_CPP_PATH")

    # Check for shared library
    if [ ! -f "$LLAMA_CPP_PATH/build/bin/libllama.so" ]; then
        echo "Error: Could not find libllama.so at $LLAMA_CPP_PATH/build/bin/libllama.so"
        echo "Make sure llama.cpp was built with -DBUILD_SHARED_LIBS=ON"
        exit 1
    fi

    echo ""
    echo "Using existing build: $LLAMA_CPP_PATH"
    ln -sfn "$LLAMA_CPP_PATH" llama.cpp
fi

# Find the shared library
if [ -f "$LLAMA_CPP_PATH/build/bin/libllama.so" ]; then
    LLAMA_CPP_LIB="$LLAMA_CPP_PATH/build/bin/libllama.so"
else
    echo "Error: Could not locate libllama.so"
    exit 1
fi

echo ""
echo "Using shared library: $LLAMA_CPP_LIB"

# Install llama-cpp-python with existing build
echo ""
echo "Installing llama-cpp-python..."

source venv/bin/activate
LLAMA_CPP_LIB="$LLAMA_CPP_LIB" LLAMA_CPP_LIB_PATH="$LLAMA_CPP_LIB" CMAKE_ARGS="-DLLAMA_BUILD=OFF" 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 [ -n "$BUILD_FAILED" ] && [ $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