import spaces
import os
import sys
import subprocess
import gradio as gr

# ======= ВЫНОСИМ УСТАНОВКУ ДО ИНИЦИАЛИЗАЦИИ GPU =======

def _install_dependencies():
    """Устанавливает зависимости ЕДИНОРАЗОВО перед запуском GPU"""
    if not os.environ.get("DETECTRON2_INSTALLED"):
        print("🔧 Installing dependencies...")
        
        # 1. ОСНОВНЫЕ ЗАВИСИМОСТИ (без компиляции)
        print("📦 Installing core dependencies...")
        subprocess.check_call([
            sys.executable, "-m", "pip", "install",
            "mmcv>=1.4.0,<2.0.0", "cython", "shapely", "timm", "h5py", "scikit-image"
        ])
        
        # 2. pytorch3d - устанавливаем БЕЗ компиляции
        print("📦 Installing pytorch3d (no compilation)...")
        try:
            # Вариант 1: Предварительно собранные wheels
            subprocess.check_call([
                sys.executable, "-m", "pip", "install",
                "fvcore", "iopath", "ninja"
            ])
            
            # Скачиваем готовый wheel для pytorch3d
            subprocess.check_call([
                sys.executable, "-m", "pip", "install",
                "--find-links", "https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu118_pyt208/download.html",
                "pytorch3d==0.7.4"
            ])
        except:
            # Вариант 2: Упрощённая установка
            print("⚠️  Using simplified pytorch3d installation...")
            subprocess.check_call([
                sys.executable, "-m", "pip", "install",
                "git+https://github.com/facebookresearch/pytorch3d.git@stable",
                "--no-build-isolation", "--no-deps"
            ])
        
        # 3. Detectron2 - ТОЛЬКО установка, без компиляции
        print("📦 Installing detectron2...")
        try:
            # Если есть локальная копия, устанавливаем БЕЗ компиляции
            if os.path.exists("MaskClustering/third_party/detectron2"):
                print("📂 Found local detectron2, installing...")
                # Ключевое: устанавливаем без сборки C++/CUDA компонентов
                os.environ["FORCE_CUDA"] = "0"
                os.environ["BUILD_CPP_TESTS"] = "0"
                
                subprocess.check_call([
                    sys.executable, "-m", "pip", "install",
                    "-e", "MaskClustering/third_party/detectron2",
                    "--no-build-isolation"
                ], env={**os.environ, "FORCE_CUDA": "0"})
            else:
                # Устанавливаем готовую версию
                subprocess.check_call([
                    sys.executable, "-m", "pip", "install",
                    "detectron2==0.6", "--no-deps"
                ])
        except Exception as e:
            print(f"⚠️  Detectron2 installation warning: {e}")
        
        # 4. ВАЖНО: ПРОПУСКАЕМ ВСЮ КОМПИЛЯЦИЮ CUDA!
        print("⚠️  Skipping ALL CUDA compilation in ZeroGPU environment")
        print("ℹ️  MSDeformAttn and other CUDA ops will run in fallback CPU mode")
        
        # Создаём файл-заглушку для CUDA компонентов
        repo_root = os.path.dirname(os.path.abspath(__file__))
        cuda_ops_dir = os.path.join(repo_root, "MaskClustering/third_party/detectron2/projects/CropFormer/mask2former/modeling/pixel_decoder/ops")
        
        if os.path.exists(cuda_ops_dir):
            # Создаём простой __init__.py чтобы импорт не падал
            init_file = os.path.join(cuda_ops_dir, "__init__.py")
            with open(init_file, "w") as f:
                f.write("# CUDA ops disabled in ZeroGPU environment\n")
                f.write("print('CUDA ops running in CPU fallback mode')\n")
        
        os.environ["DETECTRON2_INSTALLED"] = "1"
        os.environ["FORCE_CUDA"] = "0"  # КРИТИЧЕСКИ ВАЖНО!
        print("✅ All dependencies installed (CUDA compilation SKIPPED)!")
        return True
    
    print("✅ Dependencies already installed.")
    return False

# ======= СНАЧАЛА УСТАНАВЛИВАЕМ ВСЁ =======

# Ключевые переменные окружения ДО импорта torch
os.environ["FORCE_CUDA"] = "0"
os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"

# Фикс для асинхронной очереди Gradio
os.environ["GRADIO_QUEUE_ENABLED"] = "False"

@spaces.GPU()
def run_app():
    import mvp
    port = int(os.getenv("PORT", "7860"))
    mvp.demo.launch(
        server_name="0.0.0.0",
        server_port=port,
        show_error=True,
        share=False,
        show_api=False,
    )

if __name__ == "__main__":
    # Первый запуск: установка зависимостей
    if not os.environ.get("DETECTRON2_INSTALLED"):
        print("🚀 First run: installing dependencies...")
        _install_dependencies()
        print("🔄 Restarting application...")
        
        # Перезапускаем приложение
        os.execv(sys.executable, [sys.executable] + sys.argv)
    else:
        # Второй запуск: запускаем GPU-приложение
        print("🚀 Starting GPU application...")
        
        # Патч для gradio_client (оставляем как есть)
        try:
            import gradio_client.utils as _gcu
            if hasattr(_gcu, "_json_schema_to_python_type"):
                _orig = _gcu._json_schema_to_python_type
                def _json_schema_to_python_type_patched(schema, defs=None):
                    if isinstance(schema, bool):
                        return "Any"
                    return _orig(schema, defs)
                _gcu._json_schema_to_python_type = _json_schema_to_python_type_patched
        except Exception:
            pass

        # ======= ТЕПЕРЬ GPU ИНИЦИАЛИЗИРУЕТСЯ ПРАВИЛЬНО =======
        
        # Запускаем приложение
        run_app()