#!/usr/bin/env python3
"""
Deploy BytePlus Image Generation Studio to HuggingFace Spaces
"""

import os
import shutil
from pathlib import Path

def prepare_hf_deployment():
    """Prepare the app for HuggingFace Spaces deployment"""
    
    print("🚀 Preparing BytePlus Image Generation Studio for HuggingFace Spaces")
    print("=" * 60)
    
    # Create deployment directory
    deploy_dir = Path("hf_deploy")
    if deploy_dir.exists():
        shutil.rmtree(deploy_dir)
    deploy_dir.mkdir()
    
    print(f"📁 Created deployment directory: {deploy_dir}")
    
    # Files to copy to HuggingFace
    files_to_copy = [
        "hf_app.py",  # This will be the main app file
        "requirements.txt",
        ".gitattributes"
    ]
    
    # Copy essential files
    for file in files_to_copy:
        if Path(file).exists():
            shutil.copy2(file, deploy_dir / file)
            print(f"✅ Copied: {file}")
        else:
            print(f"⚠️ Missing: {file}")
    
    # Copy the actual app.py from cache if it exists
    cache_app = Path(".cache/app.py")
    if cache_app.exists():
        shutil.copy2(cache_app, deploy_dir / "app.py")
        print("✅ Copied: BytePlus app.py from cache")
    else:
        print("⚠️ BytePlus app.py not found in cache")
    
    # Create essential directories with README files
    dirs_to_create = ["Generated", "static", "view_session", "Assets"]
    
    for dir_name in dirs_to_create:
        target_dir = deploy_dir / dir_name
        target_dir.mkdir(exist_ok=True)
        
        # Copy README if it exists
        readme_file = Path(dir_name) / "README.md"
        if readme_file.exists():
            shutil.copy2(readme_file, target_dir / "README.md")
            print(f"✅ Created directory with README: {dir_name}")
        else:
            # Create a basic README
            with open(target_dir / "README.md", "w") as f:
                f.write(f"# {dir_name}\n\nThis directory is required for the BytePlus Image Generation Studio.\n")
            print(f"✅ Created directory: {dir_name}")
    
    # Copy Assets if they exist
    assets_dir = Path("Assets")
    if assets_dir.exists():
        shutil.copytree(assets_dir, deploy_dir / "Assets", dirs_exist_ok=True)
        print("✅ Copied Assets directory")
    
    # Create HuggingFace specific README.md
    hf_readme_content = """# BytePlus Image Generation Studio

A secure and feature-rich image generation studio powered by BytePlus API.

## Features

- 🎨 Multiple art styles and generation modes
- 🔒 Secure API key management with encryption
- 📊 Session-based image organization
- 📥 ZIP download functionality
- 👁️ Interactive session viewer
- 🛡️ Rate limiting and abuse prevention
- 📱 Responsive web interface

## Usage

1. Enter your BytePlus API key in the interface
2. Choose your preferred settings (style, size, etc.)
3. Enter a prompt or upload an image
4. Generate images and view them in the gallery
5. Use the session viewer to see all images from a generation session

## Session Viewer

Click the "👁️ View" button next to any session in the history to open a dedicated gallery view showing all images from that session.

## Security Features

- Encrypted API key storage
- Input validation and sanitization  
- Rate limiting (20 requests/minute)
- Secure error handling
- Image validation and size limits

## Directories

- `Generated/`: Stores generated images organized by session
- `static/`: Serves static files and ZIP downloads
- `view_session/`: Supports the session viewing functionality
- `Assets/`: Contains sample images and resources
"""
    
    with open(deploy_dir / "README.md", "w") as f:
        f.write(hf_readme_content)
    print("✅ Created HuggingFace README.md")
    
    # Create .gitignore for HuggingFace
    hf_gitignore = """# Environment files
.env

# Python cache
__pycache__/
*.pyc
*.pyo
*.pyd
.Python
*.so

# Generated content (but keep directories)
Generated/*.jpg
Generated/*.png  
Generated/*.zip
static/*.zip
static/*.jpg
static/*.png

# Logs
*.log

# OS files
.DS_Store
Thumbs.db
"""
    
    with open(deploy_dir / ".gitignore", "w") as f:
        f.write(hf_gitignore)
    print("✅ Created HuggingFace .gitignore")
    
    print("\n" + "=" * 60)
    print("✅ Deployment preparation complete!")
    print(f"📁 Deployment files ready in: {deploy_dir}")
    print("\n📋 Next steps:")
    print("1. Create a new HuggingFace Space")
    print("2. Upload all files from hf_deploy/ directory")
    print("3. Set environment variables:")
    print("   - HF_TOKEN (your HuggingFace token)")
    print("   - REPO_ID (if needed)")
    print("4. The app will automatically create necessary directories")
    print("\n🌐 The Generated, static, and view_session directories will be")
    print("   created automatically when the Space starts!")

if __name__ == "__main__":
    prepare_hf_deployment()