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/comfyanonymous/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 ---")
    controlnet_aux_path = os.path.join(APP_DIR, "custom_nodes", "comfyui_controlnet_aux")
    if not os.path.exists(controlnet_aux_path):
        os.system(f"git clone https://github.com/Fannovel16/comfyui_controlnet_aux.git {controlnet_aux_path}")
        print("✅ comfyui_controlnet_aux extension cloned.")
    else:
        print("✅ comfyui_controlnet_aux extension already exists.")

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

    print(f"✅ Current working directory is: {os.getcwd()}")
    
    import comfy.model_management
    print("--- Environment Ready ---")

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

    os.makedirs(os.path.join(APP_DIR, CHECKPOINT_DIR), exist_ok=True)
    os.makedirs(os.path.join(APP_DIR, LORA_DIR), exist_ok=True)
    os.makedirs(os.path.join(APP_DIR, EMBEDDING_DIR), exist_ok=True)
    os.makedirs(os.path.join(APP_DIR, CONTROLNET_DIR), exist_ok=True)
    os.makedirs(os.path.join(APP_DIR, DIFFUSION_MODELS_DIR), exist_ok=True)
    os.makedirs(os.path.join(APP_DIR, VAE_DIR), exist_ok=True)
    os.makedirs(os.path.join(APP_DIR, TEXT_ENCODERS_DIR), exist_ok=True)
    os.makedirs(os.path.join(APP_DIR, INPUT_DIR), exist_ok=True)
    print("✅ All required model directories are present.")