#!/usr/bin/env python3
"""
MIMO - Fast Startup Version for HuggingFace Spaces
Minimal imports to prevent timeout, full features loaded on demand
"""

import os
import gradio as gr

# Optional: small warmup function so Spaces runtime detects a GPU task and removes
# the startup warning "No @spaces.GPU function detected". This does NOT import
# heavy ML libs; it only checks environment lazily at call. If spaces package
# isn't available the decorator import will fail silently.
try:  # keep ultra-safe
    import spaces

    @spaces.GPU
    def warmup_gpu():  # lightweight, returns availability flag
        try:
            # defer torch import until after user installs heavy deps
            import importlib
            torch_spec = importlib.util.find_spec("torch")
            if torch_spec is None:
                return {"cuda": False, "detail": "torch not installed yet"}
            import torch  # type: ignore
            return {"cuda": torch.cuda.is_available()}
        except Exception as _e:  # noqa: N806
            return {"cuda": False, "detail": str(_e)}
except Exception:
    # spaces not present; ignore – minimal build still works
    pass

def create_simple_interface():
    """Create a simple interface that loads quickly"""

    def setup_and_load():
        """Force-clean and install modern stack, stub missing functorch symbol early, then validate.

        Steps:
        1. Uninstall conflicting packages (torch, torchvision, diffusers, transformers, peft, accelerate, safetensors).
        2. Install torch/torchvision first (CPU build to reduce risk) then other libs pinned.
        3. Pre-create functorch eager_transforms.grad_and_value stub if absent BEFORE importing transformers/diffusers.
        4. Validate imports.
        """
        try:
            import subprocess, sys, importlib, traceback, types

            def run(cmd):
                try:
                    subprocess.check_call(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
                    return True
                except Exception:
                    return False

            def pip_install(spec):
                ok = run([sys.executable, '-m', 'pip', 'install', '--no-cache-dir', spec])
                return ok, (f"Installed {spec}" if ok else f"Failed {spec}")

            messages = []
            # 1. Force uninstall
            uninstall_list = [
                'diffusers', 'transformers', 'torchvision', 'torch', 'peft', 'accelerate', 'safetensors'
            ]
            for pkg in uninstall_list:
                run([sys.executable, '-m', 'pip', 'uninstall', '-y', pkg])
            messages.append("Forced uninstall of prior core packages (best-effort)")

            # 2. Install core (CPU torch to avoid GPU wheel delays; pipeline mainly uses GPU later if available)
            core_specs = [ 'torch==2.0.1', 'torchvision==0.15.2' ]
            for spec in core_specs:
                ok, msg = pip_install(spec)
                messages.append(msg)

            # 3. Pre-stub functorch symbol before any heavy imports
            try:
                import importlib
                fx_mod = importlib.import_module('torch._functorch.eager_transforms')
                if not hasattr(fx_mod, 'grad_and_value'):
                    # Create lightweight placeholder using autograd backward pass simulation
                    def grad_and_value(f):
                        def wrapper(*a, **kw):
                            import torch
                            x = f(*a, **kw)
                            try:
                                if isinstance(x, torch.Tensor) and x.requires_grad:
                                    g = torch.autograd.grad(x, [t for t in a if isinstance(t, torch.Tensor) and t.requires_grad], allow_unused=True)
                                else:
                                    g = None
                            except Exception:
                                g = None
                            return g, x
                        return wrapper
                    setattr(fx_mod, 'grad_and_value', grad_and_value)
                    messages.append('Stubbed functorch.grad_and_value')
                else:
                    messages.append('functorch.grad_and_value present')
            except Exception as e:
                messages.append(f'Could not prepare functorch stub: {e}')

            # 4. Install remainder
            # Phase 1: Core ML libs (force clean versions)
            stack_specs_phase1 = [
                "huggingface_hub==0.23.0",
                "safetensors==0.4.5",
                "diffusers==0.21.4",
                "transformers==4.35.2",
                "peft==0.7.1",
                "accelerate==0.25.0",
            ]
            for spec in stack_specs_phase1:
                ok, msg = pip_install(spec)
                messages.append(msg)

            # Phase 2: Utility libs needed by app_hf_spaces.py
            stack_specs_phase2 = [
                "einops==0.7.0",
                "opencv-python-headless==4.8.1.78",
                "imageio==2.31.6",
                "imageio-ffmpeg==0.4.8",
                "tqdm==4.66.1",
            ]
            for spec in stack_specs_phase2:
                ok, msg = pip_install(spec)
                messages.append(msg)

            # Patch diffusers to disable ONNX (avoid _CAFFE2_ATEN_FALLBACK errors)
            try:
                import sys
                if 'diffusers' not in sys.modules:
                    import diffusers.utils.import_utils as diff_imports
                    diff_imports.is_onnx_available = lambda: False
                    messages.append('Patched diffusers.is_onnx_available = False')
            except Exception as e:
                messages.append(f'ONNX patch failed (non-critical): {e}')

            # Defer tensorflow until after core validation to reduce failure surface
            deferred_tensorflow = 'tensorflow-cpu==2.13.0'
            # 5. Validate imports with diffusers fallback chain
            def try_import(autoencoder_strict=False):
                import importlib
                import torch  # noqa: F401
                import diffusers  # noqa: F401
                import transformers  # noqa: F401
                if autoencoder_strict:
                    # direct AutoencoderKL import path changed in some versions
                    from diffusers import AutoencoderKL  # noqa: F401
                return True

            # Try import with fallback: 0.21.4 → 0.20.2
            diffusers_versions = ["0.21.4", "0.20.2"]
            last_error = None
            for idx, ver in enumerate(diffusers_versions):
                try:
                    # Reinstall target diffusers version fresh each attempt
                    run([sys.executable, '-m', 'pip', 'uninstall', '-y', 'diffusers'])
                    ok, msg = pip_install(f'diffusers=={ver}')
                    messages.append(msg)
                    if not ok:
                        last_error = msg
                        continue
                    # Relax autoencoder import for first attempts (some versions restructure)
                    strict = (ver == diffusers_versions[-1])
                    try_import(autoencoder_strict=strict)
                    messages.append(f'diffusers import OK at {ver} (strict={strict})')
                    last_error = None
                    break
                except Exception as e:
                    last_error = str(e)
                    messages.append(f'diffusers version {ver} failed: {e}')

            if last_error:
                messages.append(f'Final diffusers import failure after fallbacks: {last_error}')
                return '❌ Setup failed during import validation\n' + '\n'.join(messages)

            # Install deferred tensorflow optionally
            ok_tf, msg_tf = pip_install(deferred_tensorflow)
            messages.append(msg_tf)

            # Secondary optional: attempt AutoencoderKL explicit import to ensure availability (soft)
            try:
                from diffusers import AutoencoderKL  # noqa: F401
            except Exception as e:
                messages.append(f'Warning: AutoencoderKL direct import not required but failed: {e}')

            # 6. Try app import
            try:
                from app_hf_spaces import CompleteMIMO, gradio_interface  # noqa: F401
            except Exception as e:
                tb = traceback.format_exc(limit=2)
                messages.append(f'App import partial failure: {e}\n{tb}')
                return '⚠️ Core libs installed but app import failed\n' + '\n'.join(messages)

            return '✅ Clean stack installed! Please refresh to load full MIMO.\n' + '\n'.join(messages)

        except Exception as e:
            return f'❌ Setup failed: {e}'

    with gr.Blocks(title="MIMO - Loading...", theme=gr.themes.Soft()) as demo:
        gr.HTML("""
        <div style="text-align: center; padding: 2rem;">
            <h1>🎭 MIMO - Character Video Synthesis</h1>
            <p>Loading complete implementation...</p>
            <p>Click the button below to install remaining dependencies and activate full features.</p>
        </div>
        """)

        setup_btn = gr.Button("� Install Dependencies & Activate MIMO", variant="primary", size="lg")
        status = gr.Textbox(label="Status", interactive=False, lines=3)

        setup_btn.click(fn=setup_and_load, outputs=[status])

        gr.HTML("""
        <div style="margin-top: 2rem; padding: 1rem; background: #f0f0f0; border-radius: 8px;">
            <h4>Why this approach?</h4>
            <p>To prevent HuggingFace Spaces build timeout, we use minimal dependencies at startup.</p>
            <p>Full MIMO features (Character Animation + Video Editing) will be available after setup.</p>
        </div>
        """)

    return demo

"""
We do NOT attempt to import the full heavy implementation during build/startup.
The previous version tried a best-effort import inside a try/except. Even though it
failed fast, it still triggered Python to resolve heavy modules (torch/diffusers)
which aren't installed in the minimal build image. That adds noise and (in some
cases) delays. We now always start with the light interface; the user explicitly
chooses to install heavy dependencies.

Keeping changes minimal per user request: no extra files or new features, just a
safer lazy-loading path.
"""

# Always start with minimal interface (no premature heavy imports)
app = create_simple_interface()

if __name__ == "__main__":
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )