#!/bin/bash
# Usage: ./run.sh [--epochs N] [--batch_size N] [--num_samples N]
# Example: ./run.sh --epochs 1 --batch_size 64 --num_samples 1000

set -e

GREEN='\033[0;32m'
BLUE='\033[0;34m'
YELLOW='\033[1;33m'
NC='\033[0m'

cd "$(dirname "${BASH_SOURCE[0]}")"

# Default values
EPOCHS=""
BATCH_SIZE=""
NUM_SAMPLES=""

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --epochs)
            EPOCHS="--epochs $2"
            shift 2
            ;;
        --batch_size)
            BATCH_SIZE="--batch_size $2"
            shift 2
            ;;
        --num_samples)
            NUM_SAMPLES="--num_samples $2"
            shift 2
            ;;
        *)
            echo -e "${YELLOW}Unknown option: $1${NC}"
            exit 1
            ;;
    esac
done

# Install uv if not exists
if ! command -v uv &> /dev/null; then
    echo -e "${GREEN}Installing uv...${NC}"
    curl -LsSf https://astral.sh/uv/install.sh | sh
fi

# Sync dependencies
echo -e "${GREEN}✓ Syncing dependencies...${NC}"
uv sync

# Activate venv
source .venv/bin/activate

# Check .env
if [ ! -f ".env" ]; then
    cp .env.example .env
fi

# Quick GPU check
python3 -c "
import torch
if torch.cuda.is_available():
    print(f'✓ GPU: {torch.cuda.get_device_name(0)}')
else:
    print('⚠ CPU mode (slower)')
"

# Build training command
TRAIN_CMD="python3 src/training_pipeline.py config"
if [ -n "$EPOCHS" ]; then
    TRAIN_CMD="$TRAIN_CMD $EPOCHS"
fi
if [ -n "$BATCH_SIZE" ]; then
    TRAIN_CMD="$TRAIN_CMD $BATCH_SIZE"
fi
if [ -n "$NUM_SAMPLES" ]; then
    TRAIN_CMD="$TRAIN_CMD $NUM_SAMPLES"
fi

echo -e "${BLUE}╔════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║  Training GPT-2 on Spotify Dataset     ║${NC}"
echo -e "${BLUE}╚════════════════════════════════════════╝${NC}"
echo ""
echo -e "${YELLOW}Command: $TRAIN_CMD${NC}"
echo ""

eval $TRAIN_CMD

echo ""
echo -e "${GREEN}✓ Training Complete!${NC}"
echo -e "  Model: ${BLUE}./model${NC}"
echo ""

echo -e "${BLUE}╔════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║  Testing Model                         ║${NC}"
echo -e "${BLUE}╚════════════════════════════════════════╝${NC}"
echo ""

python3 test_model.py ./model