#!/bin/bash

# ===============================================================
# Hugging Face Deployment Script
# ===============================================================
# This script deploys your model to Hugging Face.
#
# SECURITY NOTICE:
# - NEVER hardcode tokens or secrets in this script
# - Use environment variables for sensitive information
# - Store secrets in a secure location, not in your code
#
# Required environment variables from .env file in project root:
# - HUGGING_FACE_TOKEN: Your Hugging Face API token
# - HUGGING_FACE_SPACE: Your Hugging Face space name
#
# How to use:
# 1. Make sure .env file exists in project root with proper variables
#
# 2. Run the script:
#    ./scripts/deploy_to_hf.sh
# ===============================================================

# Color definitions for better readability
RED="\033[0;31m"
GREEN="\033[0;32m"
YELLOW="\033[0;33m"
BLUE="\033[0;34m"
PURPLE="\033[0;35m"
CYAN="\033[0;36m"
NC="\033[0m" # No Color

# Function for logging messages with color
log_info() {
  echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
  echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
  echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
  echo -e "${RED}[ERROR]${NC} $1"
}

log_section() {
  echo -e "\n${PURPLE}=== $1 ===${NC}"
}

# Function to check if a command exists
command_exists() {
  command -v "$1" >/dev/null 2>&1
}

# Function to validate token format (basic check)
validate_token() {
  local token="$1"
  # Check if token starts with "hf_" prefix (typical for Hugging Face tokens)
  if [[ ! "$token" =~ ^hf_[a-zA-Z0-9]+$ ]]; then
    return 1
  fi
  return 0
}

# Exit on any error
set -e

# Script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"

# Change to project root directory
cd "$PROJECT_ROOT" || {
  log_error "Failed to change to project root directory: $PROJECT_ROOT"
  exit 1
}

log_section "Starting Hugging Face Deployment"

# Check for required dependencies
log_info "Checking dependencies..."

MISSING_DEPS=0
for cmd in git huggingface-cli python pip; do
  if ! command_exists "$cmd"; then
    log_error "Required command not found: $cmd"
    MISSING_DEPS=1
  fi
done

if [ $MISSING_DEPS -eq 1 ]; then
  log_error "Please install missing dependencies and try again."
  exit 1
fi

# Load environment variables from .env file
log_info "Loading environment variables from .env file..."
if [ -f "$PROJECT_ROOT/.env" ]; then
  source "$PROJECT_ROOT/.env"
else
  log_error ".env file not found in project root directory!"
  exit 1
fi

# Check for required environment variables
log_info "Checking environment variables..."

if [ -z "$HUGGING_FACE_TOKEN" ]; then
  log_error "HUGGING_FACE_TOKEN environment variable is not set in .env file!"
  exit 1
fi

if [ -z "$HUGGING_FACE_SPACE" ]; then
  log_error "HUGGING_FACE_SPACE environment variable is not set in .env file!"
  exit 1
fi

# Validate token format
if ! validate_token "$HUGGING_FACE_TOKEN"; then
  log_warning "The HUGGING_FACE_TOKEN format doesn't appear valid (should start with 'hf_')."
  read -p "Continue anyway? (y/n): " -n 1 -r
  echo
  if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    log_info "Deployment canceled."
    exit 1
  fi
fi

log_info "Deploying to Hugging Face space: $HUGGING_FACE_SPACE"

# Login to Hugging Face
log_info "Logging in to Hugging Face..."
if ! HUGGINGFACE_HUB_TOKEN="$HUGGING_FACE_TOKEN" huggingface-cli login --token "$HUGGING_FACE_TOKEN" 2>/dev/null; then
  log_error "Failed to log in to Hugging Face. Please check your token."
  exit 1
fi
log_success "Logged in to Hugging Face successfully."

# Check if we're in a git repository
if [ ! -d .git ]; then
  log_warning "Not in a git repository. Initializing..."
  git init
  git config --local user.name "HIM Deployment"
  git config --local user.email "no-reply@teleologyhi.com"
fi

# Check if repository is configured properly
if ! git remote -v | grep -q "$HUGGING_FACE_SPACE"; then
  log_info "Setting up git remote for Hugging Face..."
  git remote add huggingface "https://huggingface.co/spaces/$HUGGING_FACE_SPACE"
  # Alternatively, use the token in the URL (less secure but sometimes needed)
  # git remote add huggingface "https://USER:$HUGGING_FACE_TOKEN@huggingface.co/spaces/$HUGGING_FACE_SPACE"
fi

# Make sure dependencies are installed
log_info "Checking if requirements.txt exists..."
if [ -f requirements.txt ]; then
  log_info "Installing dependencies from requirements.txt..."
  if ! pip install -r requirements.txt; then
    log_warning "Some dependencies could not be installed. This might cause deployment issues."
  fi
else
  log_warning "requirements.txt not found. Make sure all dependencies are installed."
fi

# Prepare deployment
log_section "Preparing Deployment"

# Check for untracked files
UNTRACKED_FILES=$(git ls-files --others --exclude-standard)
if [ -n "$UNTRACKED_FILES" ]; then
  log_warning "Untracked files found:"
  echo "$UNTRACKED_FILES"
  read -p "Add these files to git? (y/n): " -n 1 -r
  echo
  if [[ $REPLY =~ ^[Yy]$ ]]; then
    git add .
    log_success "Files added to git."
  else
    log_warning "Continuing without adding untracked files."
  fi
fi

# Commit changes if needed
if ! git diff-index --quiet HEAD --; then
  log_info "Changes detected, creating a commit..."
  read -p "Enter commit message (default: 'Update deployment'): " COMMIT_MSG
  COMMIT_MSG=${COMMIT_MSG:-"Update deployment"}
  git commit -am "$COMMIT_MSG"
  log_success "Changes committed."
fi

# Deploy to Hugging Face
log_section "Deploying to Hugging Face"
log_info "Pushing to Hugging Face Space: $HUGGING_FACE_SPACE"

if ! git push -u huggingface main; then
  log_error "Failed to push to Hugging Face."
  log_info "Try running: git push -u huggingface main --force"
  exit 1
fi

log_success "Deployment successful!"
log_info "Your model is now available at: https://huggingface.co/spaces/$HUGGING_FACE_SPACE"

# Optional: Generate a client usage example
log_section "Client Usage Example"
cat << EOF
Python client example:

from gradio_client import Client

client = Client("$HUGGING_FACE_SPACE")

result = client.predict(
    message="Hello!!",
    system_message="You are a friendly Chatbot.",
    max_tokens=512,
    temperature=0.7,
    top_p=0.95,
    api_name="/chat"
)

print(result)
EOF

exit 0