import os
import sys
import shutil

from core.settings import *

def move_and_overwrite(src, dst):
    if os.path.isdir(src):
        os.makedirs(dst, exist_ok=True)
        for item in os.listdir(src):
            s_child = os.path.join(src, item)
            d_child = os.path.join(dst, item)
            move_and_overwrite(s_child, d_child)
    elif os.path.isfile(src):
        if os.path.exists(dst):
            try:
                os.remove(dst)
            except OSError:
                pass
        shutil.move(src, dst)

def initialize_comfyui():
    APP_DIR = sys.path[0]
    COMFYUI_TEMP_DIR = "ComfyUI_temp"

    print("--- Cloning ComfyUI Repository ---")
    if not os.path.exists(COMFYUI_TEMP_DIR):
        os.system(f"git clone https://github.com/comfy-Org/ComfyUI {COMFYUI_TEMP_DIR}")
        print("✅ ComfyUI repository cloned.")
    else:
        print("✅ ComfyUI repository already exists.")

    print(f"--- Merging ComfyUI from '{COMFYUI_TEMP_DIR}' to '{APP_DIR}' ---")
    for item in os.listdir(COMFYUI_TEMP_DIR):
        src_path = os.path.join(COMFYUI_TEMP_DIR, item)
        dst_path = os.path.join(APP_DIR, item)
        if item == '.git':
            continue
        move_and_overwrite(src_path, dst_path)
    
    try:
        shutil.rmtree(COMFYUI_TEMP_DIR)
        print("✅ ComfyUI merged and temporary directory removed.")
    except OSError as e:
        print(f"⚠️ Could not remove temporary directory '{COMFYUI_TEMP_DIR}': {e}")




    print(f"✅ Current working directory is: {os.getcwd()}")
    
    # [ZeroGPU / Upstream ComfyUI Compatibility]
    # Guard against ComfyUI PR #16746 is_integrated attribute check crash under ZeroGPU SimpleNamespace
    try:
        import torch
        if hasattr(torch, "cuda") and hasattr(torch.cuda, "get_device_properties"):
            _orig_gdp = torch.cuda.get_device_properties
            def _safe_gdp(*args, **kwargs):
                p = _orig_gdp(*args, **kwargs)
                if not hasattr(p, "is_integrated"):
                    try:
                        p.is_integrated = False
                    except Exception:
                        pass
                return p
            torch.cuda.get_device_properties = _safe_gdp
    except Exception:
        pass

    mm_path = os.path.join(APP_DIR, "comfy", "model_management.py")
    if os.path.exists(mm_path):
        try:
            with open(mm_path, "r", encoding="utf-8") as f:
                mm_code = f.read()
            unsafe_check = "bool(torch.cuda.get_device_properties(device).is_integrated)"
            safe_check = "bool(getattr(torch.cuda.get_device_properties(device), 'is_integrated', False))"
            if unsafe_check in mm_code:
                with open(mm_path, "w", encoding="utf-8") as f:
                    f.write(mm_code.replace(unsafe_check, safe_check))
                print("🛡️ [Setup] Patched comfy/model_management.py for safe is_integrated attribute access.")
        except Exception as e:
            print(f"⚠️ [Setup] Failed to patch comfy/model_management.py: {e}")

    import comfy.model_management
    print("--- Environment Ready ---")

    print("✅ ComfyUI initialized with default attention mechanism.")

    for dir_path in CATEGORY_TO_DIR_MAP.values():
        os.makedirs(os.path.join(APP_DIR, dir_path), exist_ok=True)
    
    os.makedirs(os.path.join(APP_DIR, INPUT_DIR), exist_ok=True)
    os.makedirs(os.path.join(APP_DIR, OUTPUT_DIR), exist_ok=True)
    
    print("✅ All required model directories are present.")