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 ---")
    custom_nodes_dir = os.path.join(APP_DIR, "custom_nodes")
    os.makedirs(custom_nodes_dir, exist_ok=True)

    newbie_nodes_path = os.path.join(custom_nodes_dir, "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.")

    easy_use_path = os.path.join(custom_nodes_dir, "ComfyUI-Easy-Use")
    if not os.path.exists(easy_use_path):
        os.system(f"git clone https://github.com/yolain/ComfyUI-Easy-Use.git {easy_use_path}")
        print("✅ ComfyUI-Easy-Use extension cloned.")
        easy_use_req_path = os.path.join(easy_use_path, "requirements.txt")
        if os.path.exists(easy_use_req_path):
            print("--- Installing requirements for ComfyUI-Easy-Use ---")
            os.system(f"pip install -r {easy_use_req_path}")
            print("✅ ComfyUI-Easy-Use requirements installed.")
    else:
        print("✅ ComfyUI-Easy-Use 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, MODEL_PATCHES_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.")