import time
import copy
import torch
import comfy.model_management
import os
import importlib.util
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logging.info("MultiGPU: Initialization started")

current_device = comfy.model_management.get_torch_device()
logging.info(f"MultiGPU: Initial device {current_device}")

def get_torch_device_patched():
    if (
        not torch.cuda.is_available()
        or comfy.model_management.cpu_state == comfy.model_management.CPUState.CPU
        or "cpu" in str(current_device).lower()
    ):
        return torch.device("cpu")
    return torch.device(current_device)

comfy.model_management.get_torch_device = get_torch_device_patched

def override_class(cls):
    class NodeOverride(cls):
        @classmethod
        def INPUT_TYPES(s):
            inputs = copy.deepcopy(cls.INPUT_TYPES())
            devices = ["cpu"] + [f"cuda:{i}" for i in range(torch.cuda.device_count())]
            inputs["required"]["device"] = (devices,)
            return inputs

        CATEGORY = "multigpu"
        FUNCTION = "override"

        def override(self, *args, device, **kwargs):
            global current_device
            current_device = device
            fn = getattr(super(), cls.FUNCTION)
            return fn(*args, **kwargs)

    return NodeOverride

NODE_CLASS_MAPPINGS = {}

def register_module(module_path, target_nodes):
    try:
        # For core nodes, skip module loading and just register from the global mappings
        if not module_path:
            logging.info("MultiGPU: Starting core node registration")
            from nodes import NODE_CLASS_MAPPINGS as GLOBAL_NODE_CLASS_MAPPINGS
            for node in target_nodes:
                if node in GLOBAL_NODE_CLASS_MAPPINGS:
                    NODE_CLASS_MAPPINGS[f"{node}MultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS[node])
                    logging.info(f"MultiGPU: Registered core node {node}")
                else:
                    logging.info(f"MultiGPU: Core node {node} not found - this shouldn't happen!")
            return

        # For custom nodes, try to load the module first
        full_path = os.path.join("custom_nodes", module_path, "__init__.py")
        logging.info(f"MultiGPU: Checking for module at {full_path}")

        if not os.path.exists(full_path):
            logging.info(f"MultiGPU: Module {module_path} not found - skipping")
            return

        logging.info(f"MultiGPU: Found {module_path}, attempting to load")
        spec = importlib.util.spec_from_file_location(module_path, full_path)
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        logging.info(f"MultiGPU: Executed {module_path} initialization")

        # Use the module's local dictionary instead of the global one
        local_map_name = "NODE_CLASS_MAPPINGS"
        local_map = getattr(module, local_map_name, None)
        if not local_map:
            logging.info(f"MultiGPU: {module_path} has no '{local_map_name}' dictionary, skipping override.")
            return

        all_defined_nodes = list(local_map.keys())
        logging.info(f"MultiGPU: {module_path} local dict keys: {all_defined_nodes}")

        for node in target_nodes:
            if node in local_map:
                mgpu_class = override_class(local_map[node])
                NODE_CLASS_MAPPINGS[f"{node}MultiGPU"] = mgpu_class
                logging.info(f"MultiGPU: Successfully wrapped {node} from {module_path}")
            else:
                logging.info(f"MultiGPU: Node '{node}' not found in {module_path}'s local dictionary")

    except Exception as e:
        logging.info(f"MultiGPU: Error processing {module_path}: {str(e)}")

# Register desired nodes
register_module("",                         ["UNETLoader", "VAELoader", "CLIPLoader", "DualCLIPLoader", "TripleCLIPLoader", "CheckpointLoaderSimple", "ControlNetLoader"])
register_module("ComfyUI-GGUF",             ["UnetLoaderGGUF","UnetLoaderGGUFAdvanced","CLIPLoaderGGUF","DualCLIPLoaderGGUF","TripleCLIPLoaderGGUF"])
register_module("x-flux-comfyui",           ["LoadFluxControlNet"])
register_module("ComfyUI-Florence2",        ["Florence2ModelLoader","DownloadAndLoadFlorence2Model"])
register_module("ComfyUI-LTXVideo",         ["LTXVLoader"])
register_module("ComfyUI-MMAudio",          ["MMAudioFeatureUtilsLoader","MMAudioModelLoader","MMAudioSampler"])
register_module("ComfyUI_bitsandbytes_NF4", ["CheckpointLoaderNF4",])

logging.info(f"MultiGPU: Registration complete. Final mappings: {', '.join(NODE_CLASS_MAPPINGS.keys())}")