#!/usr/bin/env python

import os
import sys
import logging
from pathlib import Path
from huggingface_hub import HfApi, login, CommitOperationAdd, CommitOperationDelete

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger(__name__)

def rollback_space(space_id, commit_hash):
    """Rollback a Hugging Face space to a specific commit."""
    try:
        # Initialize API
        api = HfApi()
        logger.info(f"Rolling back space {space_id} to commit {commit_hash}")
        
        # Get the commit info
        commit_info = api.list_repo_commits(repo_id=space_id, repo_type="space")[0]
        logger.info(f"Current commit: {commit_info.commit_id}")
        
        # Get the files at the target commit
        target_files = api.list_repo_files(
            repo_id=space_id,
            repo_type="space",
            revision=commit_hash
        )
        logger.info(f"Found {len(target_files)} files at target commit")
        
        # Download each file from the target commit
        operations = []
        for file_path in target_files:
            try:
                content = api.hf_hub_download(
                    repo_id=space_id,
                    repo_type="space",
                    filename=file_path,
                    revision=commit_hash
                )
                with open(content, 'rb') as f:
                    file_content = f.read()
                operations.append(CommitOperationAdd(path_or_fileobj=file_content, path_in_repo=file_path))
                logger.info(f"Added {file_path} to rollback operations")
            except Exception as e:
                logger.warning(f"Failed to download {file_path}: {str(e)}")
        
        if operations:
            # Create rollback commit
            api.create_commit(
                repo_id=space_id,
                repo_type="space",
                operations=operations,
                commit_message=f"Rollback to {commit_hash}",
                revision="main"
            )
            logger.info(f"Successfully rolled back to commit {commit_hash}")
        else:
            logger.warning("No files to commit")
        
        return True
    except Exception as e:
        logger.error(f"Error rolling back space: {str(e)}")
        return False

def main():
    # Set up environment
    try:
        from dotenv import load_dotenv
        env_path = Path(__file__).parent / ".env"
        if env_path.exists():
            load_dotenv(env_path)
            logger.info(f"Loaded environment variables from {env_path}")
    except ImportError:
        logger.warning("python-dotenv not installed, skipping .env loading")
    
    # Get token
    token = os.environ.get("HF_TOKEN")
    if not token:
        logger.error("HF_TOKEN environment variable not found")
        return False
    
    # Login to Hugging Face
    login(token=token)
    logger.info("Logged in to Hugging Face")
    
    # Rollback space
    space_id = "George-API/DeepSpace"
    commit_hash = "7ba62477b32b389c2b0d5c85138e8b3c531a76cd"
    
    success = rollback_space(space_id, commit_hash)
    if success:
        print(f"\nSpace {space_id} successfully rolled back to commit {commit_hash}")
    else:
        print(f"\nFailed to rollback space {space_id}")
    
    return success

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)