#!/bin/bash
# Setup script for automated notebook execution and plot extraction
# This script creates an hourly cron job (dependencies must be installed separately)

set -e  # Exit on error

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Get the directory where this script is located
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
cd "$SCRIPT_DIR"

echo -e "${GREEN}Setting up automated notebook job...${NC}"
echo "Note: This script only configures cron. Install dependencies first with:"
echo "  pip install -r requirements.txt"
echo

# Find conda base directory
if [ -d "$HOME/anaconda3" ]; then
    CONDA_BASE="$HOME/anaconda3"
elif [ -d "$HOME/miniconda3" ]; then
    CONDA_BASE="$HOME/miniconda3"
else
    CONDA_BASE=$(conda info --base 2>/dev/null || echo "$HOME/anaconda3")
fi

ENV_NAME="suno_env_dev"
echo "Using conda at: $CONDA_BASE"
echo "Using conda environment: $ENV_NAME"

# Check if config.yaml is set up
if grep -q "YOUR_WEBHOOK_URL_HERE" config.yaml; then
    echo -e "${YELLOW}Warning: config.yaml still has placeholder values${NC}"
    echo "Please edit config.yaml and add your Slack webhook URL before running"
else
    echo -e "${GREEN}✓ Config file looks good${NC}"
fi

# Check if dependencies are installed
echo "Checking Python dependencies..."
source "$CONDA_BASE/etc/profile.d/conda.sh"
conda activate "$ENV_NAME"

if ! python -c "import yaml, requests, nbconvert; from PIL import Image" 2>/dev/null; then
    echo -e "${RED}Error: Required dependencies not installed${NC}"
    echo "Please run: pip install -r requirements.txt"
    exit 1
else
    echo -e "${GREEN}✓ Dependencies are installed${NC}"
fi

# Make the Python script executable
chmod +x run_and_notify.py

# Check if run_hourly.sh already exists (don't overwrite if manually edited)
WRAPPER_SCRIPT="$SCRIPT_DIR/run_hourly.sh"
if [ ! -f "$WRAPPER_SCRIPT" ]; then
    cat > "$WRAPPER_SCRIPT" << 'EOF'
#!/bin/bash
# Wrapper script for cron execution
# Usage: ./run_hourly.sh [--skip-notebook]
#   --skip-notebook: Skip notebook execution, only upload existing plot

# Get the directory where this script is located
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

# Initialize conda
CONDA_BASE=$(conda info --base 2>/dev/null)
if [ -z "$CONDA_BASE" ]; then
    # Try common conda locations if conda not in PATH
    for CONDA_LOC in "$HOME/anaconda3" "$HOME/miniconda3" "/opt/conda"; do
        if [ -d "$CONDA_LOC" ]; then
            CONDA_BASE="$CONDA_LOC"
            break
        fi
    done
fi

# Source conda
if [ -f "$CONDA_BASE/etc/profile.d/conda.sh" ]; then
    source "$CONDA_BASE/etc/profile.d/conda.sh"
    conda activate suno_env_dev
else
    echo "Error: Cannot find conda initialization script" >&2
    exit 1
fi

# Run the Python script with arguments passed through
python run_and_notify.py "$@"

# Exit with the same code as the Python script
exit $?
EOF

    chmod +x "$WRAPPER_SCRIPT"
    echo -e "${GREEN}Created wrapper script: $WRAPPER_SCRIPT${NC}"
else
    echo -e "${YELLOW}Wrapper script already exists, skipping${NC}"
fi

# Set up cron job (runs at the start of every hour)
CRON_SCHEDULE="0 * * * *"
CRON_COMMAND="$WRAPPER_SCRIPT >> $SCRIPT_DIR/logs/cron.log 2>&1"
CRON_JOB="$CRON_SCHEDULE $CRON_COMMAND"

# Check if cron job already exists
EXISTING_CRON=$(crontab -l 2>/dev/null | grep -F "$WRAPPER_SCRIPT" || true)

if [ -n "$EXISTING_CRON" ]; then
    echo -e "${YELLOW}Cron job already exists:${NC}"
    echo "$EXISTING_CRON"
    read -p "Replace it? (y/n) " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        # Remove old cron job
        crontab -l 2>/dev/null | grep -v -F "$WRAPPER_SCRIPT" | crontab -
        echo "Old cron job removed"
    else
        echo "Keeping existing cron job"
        exit 0
    fi
fi

# Add new cron job
(crontab -l 2>/dev/null; echo "$CRON_JOB") | crontab -

echo -e "${GREEN}✓ Cron job installed successfully!${NC}"
echo ""
echo "Schedule: Every hour at minute 0"
echo "Command: $WRAPPER_SCRIPT"
echo "Logs: $SCRIPT_DIR/logs/"
echo ""
echo "Current crontab:"
crontab -l | grep -F "$WRAPPER_SCRIPT"
echo ""
echo -e "${GREEN}Setup complete!${NC}"
echo ""
echo "To test the job manually, run:"
echo "  $WRAPPER_SCRIPT"
echo ""
echo "To remove the cron job, run:"
echo "  crontab -e"
echo "  # Then delete the line containing: $WRAPPER_SCRIPT"