#!/bin/bash

# HuggingFace deployment script for PuppyCompanion FastAPI
# Usage: ./deploy_hf.sh [--space-name SPACE_NAME]

set -e

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

log() {
    echo -e "${GREEN}[DEPLOY]${NC} $1"
}

warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

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

info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

# Default configuration
SPACE_NAME="puppycompanion-v3"
APP_DIR="."
HF_USERNAME="JTh34"

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --space-name)
            SPACE_NAME="$2"
            shift 2
            ;;
        *)
            error "Unknown option: $1"
            exit 1
            ;;
    esac
done

# Ask for space name if not provided (with default)
if [ -z "$SPACE_NAME" ]; then
    read -p "HuggingFace space name [puppycompanion-v3]: " SPACE_NAME
    SPACE_NAME=${SPACE_NAME:-puppycompanion-v3}
fi

if [ -z "$SPACE_NAME" ]; then
    error "Space name required"
    exit 1
fi

# Check that the app directory exists
if [ ! -d "$APP_DIR" ]; then
    error "Application directory not found: $APP_DIR"
    exit 1
fi

cd "$APP_DIR"

log "Preparing deployment for FastAPI space: $HF_USERNAME/$SPACE_NAME"

# Pre-deployment checks
log "Pre-deployment checks..."

# Check required files for FastAPI application
REQUIRED_FILES=(
    "main.py"
    "Dockerfile"
    "README.MD"
    "all_books_preprocessed_chunks.json"
    "rag_system.py"
    "agent_workflow.py"
    "embedding_models.py"
    "books_config.json"
    "static/index.html"
)

for file in "${REQUIRED_FILES[@]}"; do
    if [ ! -f "$file" ]; then
        error "Missing required file: $file"
        exit 1
    fi
done

log "✅ All required files found"

# Check that API keys are not in files
ENV_FILES=(".env" ".env.local" ".env.prod" ".env.development" ".env.staging")
for env_file in "${ENV_FILES[@]}"; do
    if [ -f "$env_file" ]; then
        if grep -q "sk-\|OPENAI_API_KEY\|TAVILY_API_KEY" "$env_file" 2>/dev/null; then
            warn "API keys detected in $env_file - file will be excluded from deployment"
            warn "Make sure to configure your secrets on HuggingFace Spaces"
        fi
    fi
done

# Check for PDF files that should not be deployed
if ls data/*.pdf >/dev/null 2>&1; then
    warn "PDF files detected in data/ directory"
    warn "These will be excluded from deployment for security"
fi

# Validate chunks file size
if [ -f "all_books_preprocessed_chunks.json" ]; then
    FILE_SIZE=$(stat -f%z "all_books_preprocessed_chunks.json" 2>/dev/null || stat -c%s "all_books_preprocessed_chunks.json" 2>/dev/null)
    FILE_SIZE_MB=$((FILE_SIZE / 1024 / 1024))
    log "Chunks file size: ${FILE_SIZE_MB}MB"
    if [ $FILE_SIZE_MB -gt 100 ]; then
        warn "Large chunks file detected (${FILE_SIZE_MB}MB)"
        warn "This may cause slower deployment and startup times"
    fi
fi

# Create requirements.txt from pyproject.toml if needed and requirements.txt doesn't exist
if [ ! -f "requirements.txt" ] && [ -f "pyproject.toml" ]; then
    log "Generating requirements.txt from pyproject.toml..."
    if command -v pip-compile &> /dev/null; then
        pip-compile pyproject.toml
    else
        warn "pip-compile not found. Basic requirements.txt generation..."
        # Basic extraction of dependencies from pyproject.toml
        if command -v python &> /dev/null; then
            python -c "
import tomllib
with open('pyproject.toml', 'rb') as f:
    data = tomllib.load(f)
deps = data.get('project', {}).get('dependencies', [])
with open('requirements.txt', 'w') as f:
    for dep in deps:
        f.write(dep + '\n')
" 2>/dev/null || {
                warn "Could not generate requirements.txt automatically"
                warn "Please create the requirements.txt file manually"
                exit 1
            }
        fi
    fi
elif [ -f "requirements.txt" ]; then
    log "✅ requirements.txt found"
fi

# Validate Dockerfile for FastAPI
if grep -q "main.py" Dockerfile && grep -q "7860" Dockerfile; then
    log "✅ Dockerfile configured for FastAPI"
else
    warn "Dockerfile may not be properly configured for FastAPI deployment"
fi

# Check SSH connection to HuggingFace
log "Checking HuggingFace SSH connection..."
if ! ssh -T git@hf.co -o ConnectTimeout=10 -o BatchMode=yes 2>/dev/null; then
    warn "SSH connection to HuggingFace failed"
    info "Check your SSH configuration or use: ssh-keygen -t ed25519 -C 'your_email@example.com'"
    info "Then add the public key to your HuggingFace profile"
    
    # Fallback to HTTPS if available
    if command -v huggingface-cli &> /dev/null && huggingface-cli whoami &> /dev/null; then
        warn "Using HTTPS as fallback..."
        USE_SSH=false
    else
        error "No authentication method available"
        exit 1
    fi
else
    log "✅ HuggingFace SSH connection OK"
    USE_SSH=true
fi

# Create the space if it does not exist (requires huggingface-cli)
if command -v huggingface-cli &> /dev/null; then
    log "Creating/updating HuggingFace space..."
    huggingface-cli repo create "$SPACE_NAME" --type space --space_sdk docker 2>/dev/null || true
else
    warn "huggingface-cli not available, make sure the space exists"
fi

# Clone or update the repo
TEMP_DIR="/tmp/hf_deploy_$$"
log "Cloning repository..."

if [ "$USE_SSH" = true ]; then
    # Use SSH
    git clone "git@hf.co:spaces/$HF_USERNAME/$SPACE_NAME" "$TEMP_DIR" || {
        error "SSH clone failed. Check username and space name."
        exit 1
    }
else
    # Use HTTPS
    git clone "https://huggingface.co/spaces/$HF_USERNAME/$SPACE_NAME" "$TEMP_DIR" || {
        error "HTTPS clone failed. Check username and space name."
        exit 1
    }
fi

# Copy files with FastAPI-specific exclusions
log "Copying FastAPI application files..."
rsync -av \
    --exclude='.git' \
    --exclude='venv_*' \
    --exclude='.env*' \
    --exclude='__pycache__' \
    --exclude='.chainlit' \
    --exclude='*.pdf' \
    --exclude='*.PDF' \
    --exclude='data/*.pdf' \
    --exclude='data/*.PDF' \
    --exclude='.DS_Store' \
    --exclude='*.log' \
    --exclude='qdrant_storage' \
    --exclude='deploy_hf.sh' \
    --exclude='.gitattributes' \
    --exclude='uv.lock' \
    --exclude='document_loader_preproc.py' \
    ./ "$TEMP_DIR/"

# Go to temp directory
cd "$TEMP_DIR"

# Configure Git if needed
git config user.email "action@github.com" || true
git config user.name "Deploy Script" || true

# Verify critical files are present
log "Verifying deployment files..."
CRITICAL_FILES=("main.py" "Dockerfile" "all_books_preprocessed_chunks.json")
for file in "${CRITICAL_FILES[@]}"; do
    if [ ! -f "$file" ]; then
        error "Critical file missing after copy: $file"
        exit 1
    fi
done

log "✅ All critical files verified"

# Add and commit
log "Committing changes..."
git add .
git commit -m "Deploy PuppyCompanion FastAPI $(date '+%Y-%m-%d %H:%M:%S')" || {
    warn "No changes detected"
}

# Push to HuggingFace
log "Pushing to HuggingFace..."
git push

# Clean up
cd - > /dev/null
rm -rf "$TEMP_DIR"

log "🚀 FastAPI deployment completed successfully!"
info "Your app will be available at: https://huggingface.co/spaces/$HF_USERNAME/$SPACE_NAME"
info "Deployment may take a few minutes to build and start..."
info "Check the logs on HuggingFace Spaces for any issues"