import os
import sys
import site

# ================= 强力 CPU 猴子补丁 (Monkey Patch) =================
# 必须在导入任何 comfy 或 torch 模块之前执行！

# 1. 强行清理和重写启动参数
if "--use-sage-attention" in sys.argv:
    sys.argv.remove("--use-sage-attention")

if "--cpu" not in sys.argv:
    sys.argv.append("--cpu")

# 2. 伪造 CUDA 环境变量
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
os.environ["FORCE_CPU"] = "1"

# 3. 拦截并重写 torch.cuda 的核心方法，彻底废除其硬件检测能力
try:
    import torch
    
    # 强制让 torch 认为 cuda 不可用
    torch.cuda.is_available = lambda: False
    torch.cuda.device_count = lambda: 0
    
    # 劫持 current_device 和 get_device_properties，让其不触发底层 C++ 的 _cuda_init()
    def mock_current_device():
        raise AssertionError("CUDA is not available in CPU-only mode.")
    torch.cuda.current_device = mock_current_device
    
    # 关键：直接重写 ComfyUI 报错的那一行代码所依赖的底层逻辑
    # 模拟一个假的设备对象或直接让 get_torch_device 返回 cpu
    print("🤖 [CPU Patch] Successfully mocked torch.cuda to prevent driver crash.")
except Exception as e:
    print(f"⚠️ [CPU Patch] Failed to pre-patch torch: {e}")

# ====================================================================

# 尝试导入 spaces（ZeroGPU 专用，CPU 环境下保留即可）
try:
    import spaces
except ImportError:
    spaces = None

APP_DIR = os.path.dirname(os.path.abspath(__file__))
if APP_DIR not in sys.path:
    sys.path.insert(0, APP_DIR)

# 动态劫持即将被导入的 comfy.model_management 的硬件获取函数
# 这是一个更激进的拦截：直接干涉 Python 的模块系统
class ComfyManagementHacker:
    def __init__(self):
        self.patched = False
    
    def find_spec(self, fullname, path, target=None):
        # 当系统尝试导入 comfy.model_management 时，我们提前介入
        if fullname == "comfy.model_management" and not self.patched:
            self.patched = True
            print("🎯 [Interceptor] Caught comfy.model_management loading! Injecting CPU device detour...")
            # 提前把 torch.cuda 再次阉割
            import torch
            torch.device = self.make_cpu_device_wrapper(torch.device)
        return None

    def make_cpu_device_wrapper(self, original_device):
        def wrapped_device(*args, **kwargs):
            # 任何时候如果尝试创建 cuda 设备，强行扭转为 cpu
            if args and ('cuda' in str(args[0])):
                return original_device('cpu')
            return original_device(*args, **kwargs)
        return wrapped_device

sys.meta_path.insert(0, ComfyManagementHacker())


def apply_sage_attention_patch():
    print("--- [Runtime Patch] 🤖 CPU mode. SageAttention disabled. ---")
    return "Skipped (CPU Mode)"

def dummy_gpu_decorator(func):
    return func  # 完全废除 @spaces.GPU 带来的潜在干扰

@dummy_gpu_decorator
def dummy_gpu_for_startup():
    print("--- [GPU Startup] Bypassed for CPU mode. ---")
    return "Bypassed"


def main():
    # 在这里，当 setup_comfyui 内部导入 comfy.model_management 时，
    # 我们的拦截器或上面的 torch 补丁会确保它不会去触发 C++ 的 _cuda_init()
    from comfy_integration import setup as setup_comfyui
    from utils.app_utils import load_ipadapter_presets
    
    print("--- [Setup] Starting ComfyUI initialization ---")
    try:
        setup_comfyui.initialize_comfyui()
    except Exception as e:
        print(f"🚨 [Init Error] ComfyUI initialization failed: {e}")
        print("💡 If it still mentions 'RuntimeError: Found no NVIDIA driver', ComfyUI's specific script bypassed the patch. You may need a GPU instance.")
        raise e
    
    print("--- [Setup] Reloading site-packages... ---")
    try:
        site.main()
    except Exception:
        pass

    print("--- Starting Application Setup ---")
    load_ipadapter_presets()

    from ui.layout import build_ui
    from ui.events import attach_event_handlers

    demo = build_ui(attach_event_handlers)
    print("--- Launching Gradio Interface ---")
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)


if __name__ == "__main__":
    main()