import os
import sys
import shutil

from core.settings import *

def move_and_overwrite(src, dst):
    if os.path.isdir(src):
        if os.path.exists(dst):
            shutil.rmtree(dst)
        shutil.move(src, dst)
    elif os.path.isfile(src):
        if os.path.exists(dst):
            os.remove(dst)
        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("--- Cloning third-party extensions for ComfyUI ---")

    # 1. ComfyUI_IPAdapter_plus
    ipadapter_plus_path = os.path.join(APP_DIR, "custom_nodes", "ComfyUI_IPAdapter_plus")
    if not os.path.exists(ipadapter_plus_path):
        os.system(f"git clone https://github.com/cubiq/ComfyUI_IPAdapter_plus.git {ipadapter_plus_path}")
        print("✅ ComfyUI_IPAdapter_plus extension cloned.")
    else:
        print("✅ ComfyUI_IPAdapter_plus extension already exists.")

    # 2. ComfyUI-InstantX-IPAdapter-SD3
    ipadapter_plus_path = os.path.join(APP_DIR, "custom_nodes", "ComfyUI-InstantX-IPAdapter-SD3")
    if not os.path.exists(ipadapter_plus_path):
        os.system(f"git clone https://github.com/Slickytail/ComfyUI-InstantX-IPAdapter-SD3.git {ipadapter_plus_path}")
        print("✅ ComfyUI-InstantX-IPAdapter-SD3 extension cloned.")
    else:
        print("✅ ComfyUI-InstantX-IPAdapter-SD3 extension already exists.")

    # 3. ComfyUI-IPAdapter-Flux
    ipadapter_flux_path = os.path.join(APP_DIR, "custom_nodes", "ComfyUI-IPAdapter-Flux")
    if not os.path.exists(ipadapter_flux_path):
        os.system(f"git clone https://github.com/Shakker-Labs/ComfyUI-IPAdapter-Flux.git {ipadapter_flux_path}")
        print("✅ ComfyUI-IPAdapter-Flux extension cloned.")
    else:
        print("✅ ComfyUI-IPAdapter-Flux extension already exists.")

    # 4. ComfyUI-Newbie-Nodes
    newbie_nodes_path = os.path.join(APP_DIR, "custom_nodes", "ComfyUI-Newbie-Nodes")
    if not os.path.exists(newbie_nodes_path):
        os.system(f"git clone https://github.com/NewBieAI-Lab/ComfyUI-Newbie-Nodes.git {newbie_nodes_path}")
        print("✅ ComfyUI-Newbie-Nodes extension cloned.")
    else:
        print("✅ ComfyUI-Newbie-Nodes extension already exists.")

    # 5. ComfyUI-Anima-LLLite
    anima_controlnet_lllite_nodes_path = os.path.join(APP_DIR, "custom_nodes", "ComfyUI-Anima-LLLite")
    if not os.path.exists(anima_controlnet_lllite_nodes_path):
        os.system(f"git clone https://github.com/kohya-ss/ComfyUI-Anima-LLLite.git {anima_controlnet_lllite_nodes_path}")
        print("✅ ComfyUI-Anima-LLLite extension cloned.")
    else:
        print("✅ ComfyUI-Anima-LLLite extension already exists.")

    print(f"✅ Current working directory is: {os.getcwd()}")
    
    # Ensure torch treats CUDA as unavailable on CPU‑only environment.
    import torch
    torch.cuda.is_available = lambda: False
    # Additional CUDA stubs to prevent calls like torch.cuda.get_device_properties
    class _DummyDeviceProps:
        def __init__(self):
            self.major = 0
            self.minor = 0
            self.total_memory = 0
    def _dummy_get_device_properties(device):
        return _DummyDeviceProps()
    torch.cuda.get_device_properties = _dummy_get_device_properties
    torch.cuda.device_count = lambda: 0
    torch.cuda.current_device = lambda: 0
    torch.cuda.get_device_name = lambda device: "cpu"


    # Apply patch to the external ComfyUI model_management before import.
    external_model_mgmt_path = os.path.join(APP_DIR, "comfy", "model_management.py")
    if os.path.exists(external_model_mgmt_path):
        try:
            with open(external_model_mgmt_path, "r", encoding="utf-8") as f:
                content = f.read()
            start = content.find("def get_torch_device():")
            if start != -1:
                # Find the end of the function (next def or end of file)
                end = content.find("\ndef ", start + 1)
                if end == -1:
                    end = len(content)
                new_func = """def get_torch_device():
    \"\"\"Return appropriate torch device, falling back to CPU when needed.\"\"\"
    global directml_enabled, cpu_state
    if directml_enabled:
        return directml_device
    if cpu_state == CPUState.MPS:
        return torch.device(\"mps\")
    if cpu_state == CPUState.CPU:
        return torch.device(\"cpu\")
    if is_intel_xpu():
        return torch.device(\"xpu\", torch.xpu.current_device())
    if is_ascend_npu():
        return torch.device(\"npu\", torch.npu.current_device())
    if is_mlu():
        return torch.device(\"mlu\", torch.mlu.current_device())
    if torch.cuda.is_available():
        return torch.device(torch.cuda.current_device())
    return torch.device(\"cpu\")
"""
                patched_content = content[:start] + new_func + content[end:]
                with open(external_model_mgmt_path, "w", encoding="utf-8") as f:
                    f.write(patched_content)
        except Exception as e:
            print(f"⚠️ Failed to patch external model_management: {e}")

    import comfy.model_management as model_mgmt
    # Patch get_torch_device to enforce CPU fallback.
    def _cpu_fallback_get_torch_device():
        """Return appropriate torch device, falling back to CPU when needed.

        This patch replaces the original implementation that assumed a CUDA GPU
        and raised a RuntimeError on CPU‑only systems. It respects the existing
        ``directml_enabled`` and ``cpu_state`` flags, and falls back to CPU when
        no GPU backend is available.
        """
        # DirectML path – unchanged.
        if model_mgmt.directml_enabled:
            return model_mgmt.directml_device
        # Apple Silicon / MPS backend.
        if model_mgmt.cpu_state == model_mgmt.CPUState.MPS:
            return torch.device("mps")
        # Explicit CPU request.
        if model_mgmt.cpu_state == model_mgmt.CPUState.CPU:
            return torch.device("cpu")
        # Intel XPU, Ascend NPU, MLU – retain original handling.
        if model_mgmt.is_intel_xpu():
            return torch.device("xpu", torch.xpu.current_device())
        if model_mgmt.is_ascend_npu():
            return torch.device("npu", torch.npu.current_device())
        if model_mgmt.is_mlu():
            return torch.device("mlu", torch.mlu.current_device())
        # CUDA – use if available (will be False due to monkey‑patch).
        if torch.cuda.is_available():
            return torch.device(torch.cuda.current_device())
        # Default to CPU.
        return torch.device("cpu")
    model_mgmt.get_torch_device = _cpu_fallback_get_torch_device
    import importlib
    importlib.reload(model_mgmt)
    # Override should_use_bf16 to avoid any CUDA device property queries on CPU‑only systems.
    def _safe_should_use_bf16(device, model_params=None, manual_cast=False):
        """Return False unconditionally – no BF16 support on CPU‑only environment."""
        return False
    model_mgmt.should_use_bf16 = _safe_should_use_bf16

 
    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.")