#!/bin/bash

# Safe deployment script for Hugging Face Spaces
# This script prepares the repository for deployment while avoiding the 1GB limit

set -e

echo "🚀 Preparing MIMO for Hugging Face Spaces deployment..."

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

# Function to print colored output
print_status() {
    echo -e "${GREEN}✅ $1${NC}"
}

print_warning() {
    echo -e "${YELLOW}⚠️  $1${NC}"
}

print_error() {
    echo -e "${RED}❌ $1${NC}"
}

# Check if we're in the right directory
if [ ! -f "app.py" ] || [ ! -f "requirements.txt" ]; then
    print_error "Please run this script from the mimo-demo root directory"
    exit 1
fi

# Remove large files from git tracking
print_status "Removing large files from git tracking..."

# Remove pretrained weights from git
if [ -d "pretrained_weights" ]; then
    git rm -r --cached pretrained_weights/ 2>/dev/null || true
    print_status "Removed pretrained_weights from git tracking"
fi

# Remove video_decomp from git
if [ -d "video_decomp" ]; then
    git rm -r --cached video_decomp/ 2>/dev/null || true
    print_status "Removed video_decomp from git tracking"
fi

# Remove large asset files
if [ -f "assets/matting_human.pb" ]; then
    git rm --cached assets/matting_human.pb 2>/dev/null || true
    print_status "Removed large segmenter model from git tracking"
fi

# Remove any remaining large files
find . -size +50M -type f -not -path "./.git/*" | while read -r file; do
    if git ls-files --error-unmatch "$file" >/dev/null 2>&1; then
        git rm --cached "$file" 2>/dev/null || true
        print_warning "Removed large file from git tracking: $file"
    fi
done

# Ensure .gitignore is up to date
print_status "Updated .gitignore file"

# Create README for HF Spaces if it doesn't exist
if [ ! -f "README_HF.md" ]; then
    cat > README_HF.md << 'EOF'
---
title: MIMO - Controllable Character Video Synthesis
emoji: 🎭
colorFrom: purple
colorTo: pink
sdk: gradio
sdk_version: 5.33.0
app_file: app_hf.py
pinned: false
license: mit
---

# MIMO: Controllable Character Video Synthesis

MIMO enables controllable character video synthesis with spatial decomposed modeling. Upload a reference image and pose video to generate realistic character animations.

## Features

- 🎭 Controllable character animation
- 🖼️ Reference image-based generation
- 🕺 Pose-guided video synthesis
- ⚡ Optimized for HuggingFace Spaces

## Usage

1. Upload a reference character image
2. Upload a pose video or select from examples
3. Click "Generate Video" to create your animation

The model will automatically download weights from HuggingFace Hub on first use.

## Model Details

Based on the CVPR 2025 paper "MIMO: Controllable Character Video Synthesis with Spatial Decomposed Modeling"

- Model weights: ~8GB (downloaded at runtime)
- Supports both CPU and GPU inference
- Optimized for HuggingFace Spaces deployment
EOF
    print_status "Created README_HF.md for HuggingFace Spaces"
fi

# Check repository size
print_status "Checking repository size..."
REPO_SIZE=$(du -sh . --exclude=.git | cut -f1)
echo "Current repository size (excluding .git): $REPO_SIZE"

# Count files that will be uploaded
TRACKED_FILES=$(git ls-files | wc -l)
echo "Number of tracked files: $TRACKED_FILES"

# Check if any large files are still tracked
LARGE_FILES=$(git ls-files | xargs -I {} sh -c 'if [ -f "{}" ]; then du -h "{}" | awk "\$1 ~ /[0-9]+M/ || \$1 ~ /[0-9]+G/"; fi' | wc -l)

if [ "$LARGE_FILES" -gt 0 ]; then
    print_warning "Found $LARGE_FILES large files still tracked by git:"
    git ls-files | xargs -I {} sh -c 'if [ -f "{}" ]; then du -h "{}" | awk "$1 ~ /[0-9]+M/ || $1 ~ /[0-9]+G/ {print $2}"; fi'
    echo ""
    print_warning "These files may cause deployment issues. Consider adding them to .gitignore"
fi

# Commit changes
print_status "Staging changes for commit..."
git add .gitignore
git add requirements.txt
git add app_hf.py
git add README_HF.md 2>/dev/null || true

# Check if there are changes to commit
if git diff --staged --quiet; then
    print_status "No changes to commit"
else
    print_status "Committing changes..."
    git commit -m "Optimize for HuggingFace Spaces deployment

- Add .gitignore for large files (pretrained_weights/, video_decomp/)
- Update requirements.txt for HF Spaces
- Optimize app_hf.py for automatic model downloading
- Remove large files from git tracking to stay under 1GB limit"
fi

echo ""
print_status "Repository prepared for HuggingFace Spaces deployment!"
echo ""
echo "Next steps:"
echo "1. Push to your HuggingFace Space:"
echo "   git push origin main"
echo ""
echo "2. Or create a new Space:"
echo "   - Visit https://huggingface.co/new-space"
echo "   - Choose Gradio SDK"
echo "   - Set app_file to 'app_hf.py'"
echo "   - Push this repository to the Space"
echo ""
print_status "The app will automatically download model weights (~8GB) on first startup"
print_warning "Initial startup may take 10-15 minutes for weight downloading"
echo ""