#!/usr/bin/env python3
"""
HuggingFace Spaces main launcher for BytePlus Image Generation Studio

Alternative directory strategy for Spaces where the repository may not contain
the required folders (Generated/, static/, view_session/):

1) Create these folders under a persistent/writable base (prefer /data, fallback /tmp)
2) Symlink them into the working directory so the app can use relative paths

This works even when the repo itself doesn't track those directories.
"""

import os
import sys
from pathlib import Path
import importlib.util

def _is_writable(p: Path) -> bool:
    try:
        p.mkdir(parents=True, exist_ok=True)
        test_file = p / ".__writable_test__"
        test_file.write_text("ok")
        test_file.unlink(missing_ok=True)
        return True
    except Exception:
        return False

def _pick_persistent_base() -> Path:
    """Pick a persistent base directory on Spaces.
    - Prefer /data if available/writable (persisted across restarts)
    - Fallback to /tmp if /data is unavailable
    """
    candidates = [Path("/data/byteplus"), Path("/data"), Path("/tmp/byteplus")]
    for c in candidates:
        base = c if c.name == "byteplus" else c / "byteplus"
        if _is_writable(base):
            return base
    # Final fallback
    return Path("/tmp/byteplus")

def _ensure_dir(path: Path):
    path.mkdir(parents=True, exist_ok=True)
    try:
        path.chmod(0o755)
    except Exception:
        pass

def _safe_symlink(target: Path, link_path: Path):
    """Create/repair a symlink from link_path -> target.
    If link_path exists as a real dir, attempt to migrate contents, then replace with symlink.
    """
    try:
        if link_path.is_symlink() or link_path.exists():
            # If it's already the correct symlink, keep it
            try:
                if link_path.is_symlink() and link_path.resolve() == target.resolve():
                    return
            except Exception:
                pass

            if link_path.is_symlink():
                link_path.unlink()
            elif link_path.is_dir():
                # Migrate existing contents to target
                target.mkdir(parents=True, exist_ok=True)
                for item in link_path.iterdir():
                    dest = target / item.name
                    try:
                        if item.is_dir():
                            if dest.exists():
                                # Merge: move children
                                for sub in item.iterdir():
                                    sub.rename(dest / sub.name)
                                item.rmdir()
                            else:
                                item.rename(dest)
                        else:
                            if dest.exists():
                                dest.unlink()
                            item.rename(dest)
                    except Exception as e:
                        print(f"⚠️  Migration warning for {item}: {e}")
                # Now remove the original dir
                try:
                    link_path.rmdir()
                except Exception:
                    pass
            else:
                # Regular file: remove to replace with symlink
                link_path.unlink(missing_ok=True)

        # Ensure target exists before linking
        target.mkdir(parents=True, exist_ok=True)
        link_path.symlink_to(target, target_is_directory=True)
        print(f"🔗 Linked {link_path} -> {target}")
    except Exception as e:
        print(f"⚠️  Could not create symlink {link_path} -> {target}: {e}")
        # Fallback: ensure local directory exists
        try:
            _ensure_dir(link_path)
            print(f"📁 Fallback: created local dir {link_path}")
        except Exception as ie:
            print(f"❌ Fallback failed for {link_path}: {ie}")

def ensure_persistent_dirs_and_symlinks():
    """Ensure required dirs exist in persistent storage and are symlinked into CWD."""
    print("🏗️  Ensuring required directories via persistent storage and symlinks...")
    base = _pick_persistent_base()
    print(f"📦 Using storage base: {base}")

    # Create persistent dirs
    gen_p = base / "Generated"
    static_p = base / "static"
    view_p = base / "view_session"

    for d in [gen_p, static_p, view_p, static_p / "css", static_p / "js", static_p / "images"]:
        _ensure_dir(d)

    # Symlink into working directory
    cwd = Path.cwd()
    _safe_symlink(gen_p, cwd / "Generated")
    _safe_symlink(static_p, cwd / "static")
    _safe_symlink(view_p, cwd / "view_session")

    print("🎯 Symlinked working directories are ready!")

def load_byteplus_app():
    """Load and return the BytePlus application."""
    try:
        # Load the BytePlus app module
        spec = importlib.util.spec_from_file_location("byteplus_app", "byteplus_app.py")
        byteplus_module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(byteplus_module)
        
        # Try to find the Gradio app
        if hasattr(byteplus_module, 'demo'):
            return byteplus_module.demo
        elif hasattr(byteplus_module, 'app'):
            return byteplus_module.app
        elif hasattr(byteplus_module, 'interface'):
            return byteplus_module.interface
        else:
            # Look for any object with a launch method
            for attr_name in dir(byteplus_module):
                attr = getattr(byteplus_module, attr_name)
                if hasattr(attr, 'launch') and callable(attr.launch):
                    return attr
            
            raise ValueError("Could not find Gradio app in byteplus_app.py")
            
    except Exception as e:
        print(f"❌ Error loading BytePlus app: {e}")
        raise

def main():
    """Main function to setup and launch BytePlus."""
    print("🚀 Starting BytePlus Image Generation Studio for HuggingFace Spaces...")
    
    # Ensure persistent storage dirs exist and are symlinked into CWD
    ensure_persistent_dirs_and_symlinks()
    
    # Load and launch the BytePlus app
    try:
        print("📥 Loading BytePlus Image Generation Studio...")
        demo = load_byteplus_app()
        
        print("🎉 Launching BytePlus Image Generation Studio...")
        # Ensure proper binding for Spaces
        port = int(os.getenv("PORT", "7860"))
        demo.launch(server_name="0.0.0.0", server_port=port)
        
    except Exception as e:
        print(f"❌ Error launching BytePlus: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

if __name__ == "__main__":
    main()