#!/bin/bash

# Text colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

# Function to print colored messages
print_success() {
    echo -e "${GREEN}[✔] $1${NC}"
}

print_error() {
    echo -e "${RED}[✘] $1${NC}"
}

print_warning() {
    echo -e "${YELLOW}[!] $1${NC}"
}

print_info() {
    echo -e "[*] $1"
}

# Header
echo "======================================"
echo "    BLOOM AI Setup for Hugging Face   "
echo "======================================"
echo

# Create necessary directories
print_info "Setting up project directories..."
directories=("static" "templates" "cache" "cache/torch")
for dir in "${directories[@]}"; do
    if [ ! -d "$dir" ]; then
        mkdir -p "$dir"
        chmod 777 "$dir"
        print_success "Created directory: $dir"
    else
        print_info "Directory exists: $dir"
        chmod 777 "$dir"
    fi
done

# Check required files
print_info "Checking required files..."
required_files=(
    "app.py"
    "requirements.txt"
    "Dockerfile"
    "docker-compose.yml"
    "templates/index.html"
    "README.md"
)

files_ok=true
for file in "${required_files[@]}"; do
    if [ ! -f "$file" ]; then
        print_error "Missing required file: $file"
        files_ok=false
    else
        print_success "Found $file"
    fi
done

# Create README.md if it doesn't exist
if [ ! -f "README.md" ]; then
    cat > README.md <<EOF
---
title: BLOOM AI Assistant
emoji: 🤖
colorFrom: blue
colorTo: purple
sdk: docker
pinned: false
---

# BLOOM AI Assistant

This is a FastAPI-based AI assistant using BLOOM model.
EOF
    print_success "Created README.md with Hugging Face metadata"
fi

# Check static directory content
if [ -z "$(ls -A static 2>/dev/null)" ]; then
    print_warning "Static directory is empty. Make sure to add your static files."
fi

# Set correct permissions
print_info "Setting permissions..."
chmod -R 777 cache templates static
print_success "Permissions updated"

# Create .dockerignore
cat > .dockerignore <<EOF
.git
.gitignore
README.md
*.pyc
__pycache__
.pytest_cache
.env
.venv
EOF
print_success "Created .dockerignore"

# Create .gitignore
cat > .gitignore <<EOF
__pycache__/
*.py[cod]
*$py.class
*.so
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
*.egg-info/
.installed.cfg
*.egg
.env
.venv
.idea/
.vscode/
cache/
logs/
EOF
print_success "Created .gitignore"

# Final status
echo
echo "======================================"
if [ "$files_ok" = true ]; then
    print_success "Setup completed successfully!"
    echo
    print_info "Ready for Hugging Face deployment"
else
    print_error "Setup completed with warnings. Please fix the issues above."
fi
echo "======================================"