#!/bin/bash
# install_deps.sh - Install Python dependencies for the LLM Gateway fine-tuner.
# Tested on macOS with Apple Silicon (MPS) and Python 3.9+.

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT_DIR="$(dirname "$SCRIPT_DIR")"

echo "=== LLM Gateway Fine-Tuner: Dependency Installer ==="
echo "Root: $ROOT_DIR"
echo ""

# Verify Python version
PYTHON_VERSION=$(python3 --version 2>&1 | awk '{print $2}')
echo "Python: $PYTHON_VERSION"
MAJOR=$(echo "$PYTHON_VERSION" | cut -d. -f1)
MINOR=$(echo "$PYTHON_VERSION" | cut -d. -f2)
if [ "$MAJOR" -lt 3 ] || { [ "$MAJOR" -eq 3 ] && [ "$MINOR" -lt 9 ]; }; then
    echo "ERROR: Python 3.9+ required (found $PYTHON_VERSION)"
    exit 1
fi

# Check for MPS availability
python3 -c "import torch; print('MPS available:', torch.backends.mps.is_available())" 2>/dev/null || true

# Core ML dependencies
echo ""
echo "--- Installing core ML packages ---"
pip3 install \
    "peft>=0.7.0" \
    "trl>=0.7.4" \
    "datasets>=2.16.0" \
    "psycopg2-binary>=2.9.9" \
    "pyyaml>=6.0.1" \
    "requests>=2.31.0" \
    "huggingface-hub>=0.20.0"

# torch / transformers / accelerate should already be present per the
# environment spec, but install if missing
echo ""
echo "--- Verifying torch / transformers / accelerate ---"
pip3 install \
    "torch>=2.1.0" \
    "transformers>=4.36.0" \
    "accelerate>=0.25.0" \
    --upgrade 2>/dev/null || echo "WARNING: Could not upgrade torch stack — ensure versions are compatible"

# llama.cpp for GGUF conversion
echo ""
echo "--- Checking llama.cpp (for GGUF conversion) ---"
if command -v llama-quantize &>/dev/null; then
    echo "OK: llama-quantize found at $(which llama-quantize)"
else
    echo "llama-quantize not found — attempting brew install..."
    if command -v brew &>/dev/null; then
        brew install llama.cpp || echo "WARNING: brew install llama.cpp failed — GGUF conversion will be unavailable"
    else
        echo "WARNING: brew not found. Install llama.cpp manually: https://github.com/ggerganov/llama.cpp"
        echo "  Or via pip:  pip3 install llama-cpp-python"
        echo "  The fine-tuner will still train but cannot convert to GGUF."
    fi
fi

# Check for convert_hf_to_gguf.py (may ship with llama-cpp-python)
if python3 -c "import llama_cpp; import pathlib; p = pathlib.Path(llama_cpp.__file__).parent / 'convert_hf_to_gguf.py'; exit(0 if p.exists() else 1)" 2>/dev/null; then
    echo "OK: convert_hf_to_gguf.py found via llama-cpp-python"
else
    echo "NOTE: convert_hf_to_gguf.py not found in llama-cpp-python package."
    echo "      If you need GGUF conversion, install llama.cpp via brew or build from source."
fi

echo ""
echo "--- Verifying installation ---"
python3 - <<'PYEOF'
import sys
required = ["torch", "transformers", "peft", "trl", "datasets", "accelerate", "psycopg2", "yaml", "requests", "huggingface_hub"]
missing = []
for mod in required:
    try:
        __import__(mod)
        print(f"  OK  {mod}")
    except ImportError:
        print(f"  MISSING  {mod}")
        missing.append(mod)

if missing:
    print(f"\nERROR: Missing modules: {missing}")
    sys.exit(1)

import torch
print(f"\nTorch version: {torch.__version__}")
print(f"MPS available: {torch.backends.mps.is_available()}")
print(f"MPS built: {torch.backends.mps.is_built()}")
PYEOF

echo ""
echo "=== Fine-tuner dependencies installed successfully ==="
echo ""
echo "Next steps:"
echo "  1. Ensure PostgreSQL is running with the llm_gateway database"
echo "  2. Copy config/fine_tuner.yaml and adjust URLs if needed"
echo "  3. Start the service: python3 -m src.main"