#!/bin/bash
# setup_cloud.sh - Setup script for cloud deployments

# Set error handling
set -e

# Function to print colored output
print_colored() {
  GREEN='\033[0;32m'
  RED='\033[0;31m'
  YELLOW='\033[0;33m'
  BLUE='\033[0;34m'
  NC='\033[0m' # No Color
  
  case $1 in
    "info")
      echo -e "${BLUE}[INFO]${NC} $2"
      ;;
    "success")
      echo -e "${GREEN}[SUCCESS]${NC} $2"
      ;;
    "warning")
      echo -e "${YELLOW}[WARNING]${NC} $2"
      ;;
    "error")
      echo -e "${RED}[ERROR]${NC} $2"
      ;;
    *)
      echo -e "$2"
      ;;
  esac
}

# Check for Python
if ! command -v python3 &> /dev/null; then
  print_colored "error" "Python 3 is required but not found. Please install Python 3 and try again."
  exit 1
fi

# Parse arguments
PLATFORM=""
DEPLOY=false

print_colored "info" "Agentic Defensor Cloud Setup"
print_colored "info" "--------------------------------"

# Parse command line arguments
while [[ $# -gt 0 ]]; do
  case $1 in
    --platform)
      PLATFORM="$2"
      shift 2
      ;;
    --deploy)
      DEPLOY=true
      shift
      ;;
    --help)
      echo "Usage: $0 [options]"
      echo "Options:"
      echo "  --platform PLATFORM   Specify the platform (replicate or huggingface)"
      echo "  --deploy              Deploy immediately after setup"
      echo "  --help                Show this help message"
      exit 0
      ;;
    *)
      print_colored "error" "Unknown argument: $1"
      exit 1
      ;;
  esac
done

# If no platform specified, prompt user
if [ -z "$PLATFORM" ]; then
  echo "Select target platform:"
  echo "1) Replicate"
  echo "2) Hugging Face Spaces"
  read -p "Enter your choice [1/2]: " platform_choice
  
  case $platform_choice in
    1)
      PLATFORM="replicate"
      ;;
    2)
      PLATFORM="huggingface"
      ;;
    *)
      print_colored "error" "Invalid choice. Exiting."
      exit 1
      ;;
  esac
fi

# Validate platform
if [ "$PLATFORM" != "replicate" ] && [ "$PLATFORM" != "huggingface" ]; then
  print_colored "error" "Invalid platform: $PLATFORM. Choose either 'replicate' or 'huggingface'."
  exit 1
fi

print_colored "info" "Setting up for $PLATFORM deployment..."

# Check if we need to install dependencies
read -p "Install required dependencies? [y/N]: " install_deps
if [[ $install_deps =~ ^[Yy]$ ]]; then
  print_colored "info" "Installing dependencies..."
  
  if [ "$PLATFORM" == "replicate" ]; then
    pip install -r requirements.txt
    pip install cog
  else
    pip install -r requirements-huggingface.txt
  fi
  
  print_colored "success" "Dependencies installed."
fi

# Setup platform-specific files
if [ "$PLATFORM" == "replicate" ]; then
  # Replicate setup
  print_colored "info" "Setting up for Replicate..."
  
  # Check for cog.yaml
  if [ ! -f "cog.yaml" ]; then
    print_colored "error" "cog.yaml not found. Please create it first."
    exit 1
  fi
  
  # Check for data download script
  if [ ! -f "download_data.py" ]; then
    print_colored "error" "download_data.py not found. Please create it first."
    exit 1
  fi
  
  # Modify replicate.py if needed
  if [ -f "replicate.py" ] && [ -f "modify_replicate.py" ]; then
    read -p "Enable runtime data downloading? [y/N]: " enable_download
    if [[ $enable_download =~ ^[Yy]$ ]]; then
      print_colored "info" "Modifying replicate.py for runtime data downloading..."
      python modify_replicate.py
    fi
  fi
  
  # Set up data hosts
  read -p "Set up data hosts in download_data.py? [y/N]: " setup_hosts
  if [[ $setup_hosts =~ ^[Yy]$ ]]; then
    print_colored "info" "Please enter your data storage URLs:"
    read -p "Base URL for data storage: " storage_base_url
    
    # Update the download_data.py file with the storage URL
    if [ ! -z "$storage_base_url" ]; then
      sed -i.bak "s|https://your-storage-url.com|$storage_base_url|g" download_data.py
      print_colored "success" "Updated storage URLs in download_data.py"
    fi
  fi
  
  # Deploy to Replicate
  if [ "$DEPLOY" = true ]; then
    print_colored "info" "Deploying to Replicate..."
    cog push
  else
    print_colored "info" "To deploy to Replicate, run: cog push"
  fi

else
  # Hugging Face setup
  print_colored "info" "Setting up for Hugging Face Spaces..."
  
  # Check for Hugging Face files
  if [ ! -f "app_huggingface.py" ]; then
    print_colored "error" "app_huggingface.py not found. Please create it first."
    exit 1
  fi
  
  if [ ! -f "requirements-huggingface.txt" ]; then
    print_colored "error" "requirements-huggingface.txt not found. Please create it first."
    exit 1
  fi
  
  # Check for download script
  if [ ! -f "download_data.py" ]; then
    print_colored "error" "download_data.py not found. Please create it first."
    exit 1
  fi
  
  # Modify app_huggingface.py to include data downloading
  read -p "Enable runtime data downloading? [y/N]: " enable_download
  if [[ $enable_download =~ ^[Yy]$ ]]; then
    print_colored "info" "Please enter your data storage URLs:"
    read -p "Base URL for data storage: " storage_base_url
    
    # Update the download_data.py file with the storage URL
    if [ ! -z "$storage_base_url" ]; then
      sed -i.bak "s|https://your-storage-url.com|$storage_base_url|g" download_data.py
      print_colored "success" "Updated storage URLs in download_data.py"
    fi
  fi
  
  # Deploy to Hugging Face
  if [ "$DEPLOY" = true ]; then
    print_colored "warning" "Automatic deployment to Hugging Face not supported."
    print_colored "info" "Please upload the files manually to your Hugging Face Space."
  fi
  
  # Instructions for manual deployment
  print_colored "info" "To deploy to Hugging Face Spaces:"
  print_colored "info" "1. Create a new Space on huggingface.co/spaces"
  print_colored "info" "2. Select 'Gradio' as the SDK"
  print_colored "info" "3. Upload your files to the Space repository"
fi

print_colored "success" "Setup complete!"
print_colored "info" "Make sure to set your OPENAI_API_KEY as a secret in your deployment platform."