app.py

import os
import sys
import shutil
from pathlib import Path
from huggingface_hub import snapshot_download
import importlib.util

from dotenv import load_dotenv
load_dotenv()

def ensure_root_directories(main_workspace: Path) -> None:
    """Create required root-level folders and subfolders before any symlinks.

    Folders: Generated/, static/{css,js,images}/, view_session/
    Also drops a .gitkeep in each to ensure they can be tracked in a repo.
    """
    print("🏗️  Creating root-level directories (before symlinking)...")
    required = [
        main_workspace / "Generated",
        main_workspace / "static",
        main_workspace / "static" / "css",
        main_workspace / "static" / "js",
        main_workspace / "static" / "images",
        main_workspace / "view_session",
    ]

    for d in required:
        try:
            # If a symlink exists at the path, replace it with a real directory
            if d.is_symlink():
                try:
                    d.unlink()
                    print(f"🔁 Replaced symlink with real directory: {d}")
                except Exception as e:
                    print(f"⚠️  Could not remove symlink {d}: {e}")
            # Ensure directory exists
            d.mkdir(parents=True, exist_ok=True)
            try:
                d.chmod(0o755)
            except Exception:
                pass
            # Add README.md so directories are tracked per current .gitignore rules
            try:
                readme = d / "README.md"
                if not readme.exists():
                    rel = d.relative_to(main_workspace)
                    readme.write_text(f"# {rel}\n\nPlaceholder to ensure this directory is tracked.\n")
            except Exception:
                pass
            print(f"✅ Ensured directory: {d}")
        except Exception as e:
            print(f"⚠️  Warning: could not ensure {d}: {e}")

def create_symlinks_to_real_folders(cache_dir):
    """Create symbolic links in cache directory to real folders in main workspace."""
    # Use the directory containing this script as root workspace
    main_workspace = Path(__file__).parent.resolve()
    
    # Define the folders that need symbolic links
    folders_to_link = ["Generated", "static", "view_session"]
    
    print("🔗 Creating symbolic links to real folders...")
    
    for folder_name in folders_to_link:
        # Real folder path in main workspace
        real_folder = main_workspace / folder_name
        # Symbolic link path in cache
        link_path = cache_dir / folder_name
        
        try:
            # Ensure real folder exists (defensive)
            real_folder.mkdir(parents=True, exist_ok=True)

            # Remove existing file/folder/link if it exists
            if link_path.exists() or link_path.is_symlink():
                if link_path.is_symlink():
                    link_path.unlink()
                elif link_path.is_dir():
                    shutil.rmtree(link_path)
                else:
                    link_path.unlink()
            
            # Create symbolic link
            link_path.symlink_to(real_folder, target_is_directory=True)
            print(f"✅ Created symlink: {link_path} -> {real_folder}")
            
        except Exception as e:
            print(f"⚠️  Warning: Could not create symlink for {folder_name}: {e}")
            # Fallback: create empty directory
            link_path.mkdir(exist_ok=True)
            print(f"📁 Created fallback directory: {link_path}")
    
    print("🎯 All symbolic links are ready!")
    return True

def modify_byteplus_for_directories(cache_dir):
    """Modify the BytePlus app to ensure proper directory handling."""
    app_file = cache_dir / "app.py"
    
    if not app_file.exists():
        print("⚠️  BytePlus app.py not found in cache")
        return False
    
    try:
        # Read the current app content
        with open(app_file, 'r', encoding='utf-8') as f:
            app_content = f.read()
        
        # Add directory creation code at the beginning of the app
        directory_setup_code = '''
import os
from pathlib import Path

def ensure_directories_exist():
    """Ensure required directories exist for BytePlus operation."""
    required_dirs = ["Generated", "static", "view_session"]
    
    for dir_name in required_dirs:
        dir_path = Path(dir_name)
        if not dir_path.exists():
            dir_path.mkdir(parents=True, exist_ok=True)
            print(f"Created directory: {dir_path}")
        
        # Set proper permissions
        try:
            dir_path.chmod(0o755)
        except:
            pass  # Ignore permission errors
        
        # Create subdirectories
        if dir_name == "static":
            for subdir in ["css", "js", "images"]:
                (dir_path / subdir).mkdir(exist_ok=True)
    
    print("✅ All required directories are ready")

# Initialize directories when the app starts
ensure_directories_exist()

'''
        
        # Insert directory setup after the imports but before the main code
        import_end = app_content.find('\n# Configure secure logging')
        if import_end == -1:
            import_end = app_content.find('\n# Security Configuration')
        if import_end == -1:
            import_end = app_content.find('\nlogger = logging.getLogger')
        
        if import_end != -1:
            modified_content = (app_content[:import_end] + 
                              '\n' + directory_setup_code + 
                              app_content[import_end:])
            
            # Write the modified content back
            with open(app_file, 'w', encoding='utf-8') as f:
                f.write(modified_content)
            
            print("✅ Modified BytePlus app to ensure directory creation")
            return True
        else:
            print("⚠️  Could not find insertion point in BytePlus app")
            return False
            
    except Exception as e:
        print(f"❌ Error modifying BytePlus app: {e}")
        return False

def create_native_symlinks_in_cache(cache_dir):
    """Create native directories directly in the cache for BytePlus to use.
    
    Since HuggingFace repo doesn't have Generated, static, view_session folders,
    we create them directly in cache with proper permissions so BytePlus can work.
    """
    directories_to_create = ["Generated", "static", "view_session"]
    
    print("Creating required directories directly in cache for BytePlus...")
    
    for dir_name in directories_to_create:
        # Directory directly in cache
        cache_dir_path = cache_dir / dir_name
        
        try:
            # Create the directory in cache if it doesn't exist
            if not cache_dir_path.exists():
                cache_dir_path.mkdir(parents=True, exist_ok=True)
                print(f"✅ Created cache directory: {cache_dir_path}")
            
            # Set proper permissions (755 - readable/writable by owner, readable by others)
            cache_dir_path.chmod(0o755)
            print(f"✅ Set permissions 755 on: {cache_dir_path}")
            
            # Verify the directory is accessible and writable
            if cache_dir_path.is_dir() and os.access(cache_dir_path, os.W_OK):
                print(f"✅ Verified: {dir_name} is accessible and writable in cache")
            else:
                print(f"⚠️  Warning: {dir_name} may not be fully accessible")
                
        except Exception as e:
            print(f"❌ Error creating cache directory for {dir_name}: {e}")
            # Still try to create basic directory
            try:
                cache_dir_path.mkdir(parents=True, exist_ok=True)
                print(f"📁 Created basic directory: {cache_dir_path}")
            except Exception as fe:
                print(f"❌ Complete failure for {dir_name}: {fe}")
    
    # Create additional subdirectories that might be needed
    try:
        # Create common subdirectories for static files
        (cache_dir / "static" / "css").mkdir(parents=True, exist_ok=True)
        (cache_dir / "static" / "js").mkdir(parents=True, exist_ok=True)
        (cache_dir / "static" / "images").mkdir(parents=True, exist_ok=True)
        
        # Ensure view_session has proper structure
        (cache_dir / "view_session").chmod(0o755)
        
        print("✅ Created additional subdirectory structure")
    except Exception as e:
        print(f"⚠️  Warning: Could not create additional subdirectories: {e}")
    
    # Create a simple index file for static directory
    try:
        index_file = cache_dir / "static" / "index.html"
        if not index_file.exists():
            index_file.write_text("""<!DOCTYPE html>
<html>
<head><title>BytePlus Static Files</title></head>
<body><h1>BytePlus Image Generation Studio</h1><p>Static files directory</p></body>
</html>""")
        print("✅ Created static index file")
    except Exception as e:
        print(f"⚠️  Warning: Could not create static index: {e}")
    
    print("🎯 Cache directory structure ready for BytePlus operation")
    return True

def setup_cache_directory():
    """Create a hidden cache directory '.cache'.

    If an existing 'cache' directory exists, migrate its contents into '.cache'.
    Set restrictive permissions (owner rwx) and on macOS set the Finder hidden flag
    unless the environment variable `HIDE_CACHE` is explicitly set to '0'.
    """
    hidden_cache = Path(".cache")
    public_cache = Path("cache")

    # If the hidden cache doesn't exist but a public one does, move it.
    try:
        if not hidden_cache.exists():
            if public_cache.exists():
                # Prefer move to preserve contents and metadata
                public_cache.rename(hidden_cache)
            else:
                hidden_cache.mkdir(exist_ok=True)
        else:
            # ensure it exists
            hidden_cache.mkdir(exist_ok=True)

        # Restrict permissions to owner only (rwx------)
        try:
            hidden_cache.chmod(0o700)
        except Exception:
            # chmod may fail on some filesystems or platforms; ignore
            pass

        # On macOS, optionally set Finder hidden flag for extra concealment
        if sys.platform == "darwin":
            hide_flag = os.environ.get("HIDE_CACHE", "1")
            if hide_flag != "0":
                try:
                    # 'chflags hidden <path>' will make the folder hidden in Finder
                    os.system(f"/usr/bin/chflags hidden {hidden_cache}")
                except Exception:
                    pass

        return hidden_cache
    except Exception:
        # Fallback: try to create a simple cache folder named 'cache'
        public_cache.mkdir(exist_ok=True)
        return public_cache

def download_space(cache_dir):
    """Download the BytePlus space from HuggingFace."""
    token = os.environ.get("HF_TOKEN")
    repo_id = os.environ.get("REPO_ID")
    
    if not token or not repo_id:
        print("❌ HF_TOKEN or REPO_ID not found in environment variables")
        return False
    
    try:
        print(f"📥 Downloading BytePlus space: {repo_id}")
        snapshot_download(
            repo_id=repo_id,
            repo_type="space",
            local_dir=cache_dir,
            token=token
        )
        print("✅ Successfully downloaded BytePlus space")
        return True
    except Exception as e:
        print(f"❌ Error downloading space: {e}")
        return False

def load_app(cache_dir):
    sys.path.insert(0, str(cache_dir))
    app_path = cache_dir / "app.py"
    spec = importlib.util.spec_from_file_location("app", app_path)
    app = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(app)

    # First try the common 'demo' attribute used by many Gradio Spaces
    if hasattr(app, "demo"):
        return app.demo

    # Otherwise, search for any attribute with a callable 'launch' method
    for name in dir(app):
        try:
            attr = getattr(app, name)
        except Exception:
            continue
        # Skip classes (types) since their 'launch' will be an unbound function
        if isinstance(attr, type):
            continue
        if hasattr(attr, "launch") and callable(getattr(attr, "launch")):
            return attr

    # Next, accept top-level callables that return an object with a 'launch' method.
    # This covers Spaces that expose a factory function instead of a 'demo' object.
    factory_names = [
        "create_secure_interface",
        "create_app",
        "create_demo",
        "main",
        "generator",
    ]

    for name in factory_names:
        if hasattr(app, name):
            try:
                candidate = getattr(app, name)
                if callable(candidate):
                    created = candidate()
                    if hasattr(created, "launch") and callable(getattr(created, "launch")):
                        return created
                    # If the factory itself is a Gradio Interface (callable with launch), return it
                    if hasattr(candidate, "launch") and callable(getattr(candidate, "launch")):
                        return candidate
            except Exception:
                # ignore failures from calling the factory and continue searching
                pass

    # Also accept any top-level callable that when called returns an object with 'launch'
    for name in dir(app):
        if name.startswith("__"):
            continue
        try:
            attr = getattr(app, name)
        except Exception:
            continue
        if callable(attr):
            try:
                created = attr()
                if hasattr(created, "launch") and callable(getattr(created, "launch")):
                    return created
            except Exception:
                # if calling fails, skip
                continue

    # If nothing found, raise a helpful error describing the problem
    available = [n for n in dir(app) if not n.startswith("__")]
    raise AttributeError(
        "Could not find a demo application in the downloaded space."
        f" Looked for attribute 'demo' and any attribute with a callable 'launch'."
        f" Available top-level names: {available}"
    )

if __name__ == "__main__":
    print("🚀 Setting up BytePlus Image Generation Studio with root directories + symlinks...")
    
    # Setup cache directory
    cache_dir = setup_cache_directory()
    print(f"📁 Cache directory: {cache_dir}")

    # Ensure root directories exist first (as requested)
    ensure_root_directories(Path(__file__).parent.resolve())
    
    # Download the space
    if download_space(cache_dir):
        # Create symbolic links in cache to real folders in main workspace
        create_symlinks_to_real_folders(cache_dir)
        
        # Modify BytePlus app to ensure directory creation
        modify_byteplus_for_directories(cache_dir)
        
        # Load and launch the app
        try:
            demo = load_app(cache_dir)
            print("🎉 Launching BytePlus Image Generation Studio...")
            demo.launch()
        except Exception as e:
            print(f"❌ Error launching app: {e}")
            sys.exit(1)
    else:
        print("❌ Failed to download space. Please check your HF_TOKEN and REPO_ID environment variables.")
        sys.exit(1)