#!/usr/bin/env python3
"""
Deploy Morphological Transformer to Hugging Face Space
"""

import os
import subprocess
import argparse
from pathlib import Path

def run_command(cmd, check=True):
    """Run a command and handle errors"""
    print(f"Running: {' '.join(cmd)}")
    try:
        result = subprocess.run(cmd, check=check, capture_output=True, text=True)
        if result.stdout:
            print(result.stdout)
        return result.returncode == 0
    except subprocess.CalledProcessError as e:
        print(f"Error: {e}")
        if e.stderr:
            print(f"STDERR: {e.stderr}")
        return False

def main():
    parser = argparse.ArgumentParser(description='Deploy to Hugging Face Space')
    parser.add_argument('--space_name', type=str, required=True, help='Name of the Space')
    parser.add_argument('--username', type=str, required=True, help='Your Hugging Face username')
    parser.add_argument('--hf_token', type=str, help='Hugging Face token')
    parser.add_argument('--data_path', type=str, help='Path to data directory')
    parser.add_argument('--hardware', type=str, default='cpu', choices=['cpu', 'gpu'], help='Hardware type')
    parser.add_argument('--create_space', action='store_true', help='Create new Space')
    args = parser.parse_args()
    
    space_name = args.space_name
    username = args.username
    full_space_name = f"{username}/{space_name}"
    
    print(f"🚀 Deploying to Hugging Face Space: {full_space_name}")
    
    # Check if huggingface-cli is available
    if not run_command(['uv', 'run', 'huggingface-cli', '--help'], check=False):
        print("❌ huggingface-cli not found. Please install it first:")
        print("uv add huggingface_hub")
        return
    
    # Login if token is provided
    if args.hf_token:
        print("🔐 Logging in to Hugging Face...")
        if not run_command(['uv', 'run', 'huggingface-cli', 'login', '--token', args.hf_token]):
            print("❌ Failed to login to Hugging Face")
            return
    else:
        print("⚠️  No token provided. Make sure you're logged in.")
    
    # Create Space if requested
    if args.create_space:
        print(f"📦 Creating Space: {full_space_name}")
        create_cmd = [
            'uv', 'run', 'huggingface-cli', 'repo', 'create',
            full_space_name,
            '--repo-type', 'space',
            '--space_sdk', 'docker'
        ]
        if not run_command(create_cmd):
            print("❌ Failed to create Space")
            return
    
    # Clone the Space repository
    print("📥 Cloning Space repository...")
    if os.path.exists(space_name):
        print(f"⚠️  Directory {space_name} already exists. Removing...")
        import shutil
        shutil.rmtree(space_name)
    
    clone_cmd = ['git', 'clone', f'https://huggingface.co/spaces/{full_space_name}']
    if not run_command(clone_cmd):
        print("❌ Failed to clone Space repository")
        return
    
    # Copy files to Space directory
    print("📋 Copying files to Space...")
    space_dir = Path(space_name)
    
    # Copy essential files for Docker deployment
    files_to_copy = [
        'Dockerfile',
        'app.py',
        'README_SPACE.md',
        'requirements_hf.txt',
        'scripts/',
        'README.md'
    ]
    
    for file_path in files_to_copy:
        src = Path(file_path)
        dst = space_dir / file_path
        
        if src.is_file():
            dst.parent.mkdir(parents=True, exist_ok=True)
            import shutil
            shutil.copy2(src, dst)
            print(f"  ✅ Copied {file_path}")
        elif src.is_dir():
            dst.mkdir(parents=True, exist_ok=True)
            import shutil
            shutil.copytree(src, dst, dirs_exist_ok=True)
            print(f"  ✅ Copied {file_path}/")
        else:
            print(f"  ⚠️  File not found: {file_path}")
    
    # Rename files for Docker deployment
    space_readme = space_dir / 'README_SPACE.md'
    if space_readme.exists():
        space_readme.rename(space_dir / 'README.md')
        print("  ✅ Renamed README_SPACE.md to README.md")
    
    # Copy data if provided
    if args.data_path:
        data_src = Path(args.data_path)
        data_dst = space_dir / 'data'
        
        if data_src.exists():
            print(f"📊 Copying data from {args.data_path}...")
            import shutil
            shutil.copytree(data_src, data_dst, dirs_exist_ok=True)
            print("  ✅ Data copied")
        else:
            print(f"  ⚠️  Data path not found: {args.data_path}")
    
    # Create .gitignore
    gitignore_content = """# Python
__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

# Virtual environments
venv/
env/
ENV/

# IDE
.vscode/
.idea/
*.swp
*.swo

# OS
.DS_Store
Thumbs.db

# Logs
*.log

# Models and data (too large for git)
models/
data/
output/
cache/

# Weights & Biases
wandb/
"""
    
    with open(space_dir / '.gitignore', 'w') as f:
        f.write(gitignore_content)
    print("  ✅ Created .gitignore")
    
    # Commit and push changes
    print("💾 Committing and pushing changes...")
    
    # Change to Space directory
    os.chdir(space_dir)
    
    # Git commands
    git_commands = [
        ['git', 'add', '.'],
        ['git', 'commit', '-m', 'Initial deployment of Morphological Transformer'],
        ['git', 'push']
    ]
    
    for cmd in git_commands:
        if not run_command(cmd):
            print(f"❌ Failed to run: {' '.join(cmd)}")
            return
    
    print(f"🎉 Successfully deployed to {full_space_name}")
    print(f"🌐 View your Space at: https://huggingface.co/spaces/{full_space_name}")
    
    # Instructions
    print("\n📋 Next Steps:")
    print("1. Go to your Space settings")
    print("2. Set environment variables:")
    print("   - HF_TOKEN: Your Hugging Face token")
    print("   - WANDB_TOKEN: Your Weights & Biases token (optional)")
    print("   - WANDB_PROJECT: Project name for experiment tracking")
    print("3. Mount your data directory to /data")
    print("4. Launch your Space!")
    
    # Change back to original directory
    os.chdir('..')

if __name__ == '__main__':
    main()