import os
import sys
import site

# ================= 强力 CPU 环境全面拟真补丁 =================
if "--use-sage-attention" in sys.argv:
    sys.argv.remove("--use-sage-attention")

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

os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
os.environ["FORCE_CPU"] = "1"

try:
    import torch
    
    # 1. 基础 CUDA 状态架空
    torch.cuda.is_available = lambda: False
    torch.cuda.device_count = lambda: 0
    torch.cuda.current_device = lambda: 0
    
    # 2. 伪造 GPU 内存统计数据与 C++ 运行时属性
    def mock_memory_stats(*args, **kwargs):
        return {
            'reserved_bytes.all.current': 0,
            'reserved_bytes.all.peak': 0,
            'allocated_bytes.all.current': 0,
            'allocated_bytes.all.peak': 0
        }
    torch.cuda.memory_stats = mock_memory_stats
    torch.cuda.memory_allocated = lambda *args, **kwargs: 0
    torch.cuda.max_memory_allocated = lambda *args, **kwargs: 0
    torch.cuda.memory_reserved = lambda *args, **kwargs: 0
    torch.cuda.max_memory_reserved = lambda *args, **kwargs: 0
    
    # 【核心修复】：拦截 mem_get_info，返回假的 (空闲显存, 总显存)，直接给它 16GB 虚拟显存
    torch.cuda.mem_get_info = lambda *args, **kwargs: (16 * 1024 * 1024 * 1024, 16 * 1024 * 1024 * 1024)
    
    # 额外拦截设备属性查询，防止接下来报错
    class MockDeviceProperties:
        name = "Kaggle/HuggingFace CPU-Fake GPU"
        major = 8
        minor = 0
        total_memory = 16 * 1024 * 1024 * 1024
    torch.cuda.get_device_properties = lambda *args, **kwargs: MockDeviceProperties()

    # 3. 构造设备代理类，保持 type 身份
    _orig_device = torch.device

    class TorchDeviceProxy:
        def __new__(cls, *args, **kwargs):
            if args and isinstance(args[0], str) and 'cuda' in args[0]:
                return _orig_device('cpu')
            try:
                return _orig_device(*args, **kwargs)
            except Exception:
                return _orig_device('cpu')

        @property
        def __class__(self):
            return _orig_device

    TorchDeviceProxy.__name__ = _orig_device.__name__
    TorchDeviceProxy.__module__ = _orig_device.__module__
    TorchDeviceProxy.__qualname__ = _orig_device.__qualname__

    torch.device = TorchDeviceProxy
    print("🤖 [CPU Patch] Advanced TorchCuda/MemGetInfo mock deployed.")
except Exception as e:
    print(f"⚠️ [CPU Patch] Error while generating proxy patch: {e}")

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

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)

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

def dummy_gpu_decorator(func):
    return func 

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

def main():
    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}")
        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()