import torch
import os
import numpy as np
import folder_paths

class PulidModelLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"model_name": (folder_paths.get_filename_list("pulid"), )}}
    
    RETURN_TYPES = ("PULID_MODEL",)
    FUNCTION = "load_model"
    CATEGORY = "loaders"

    def load_model(self, model_name):
        model_path = folder_paths.get_full_path("pulid", model_name)
        return (model_path,)

class PulidInsightFaceLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {}}
    
    RETURN_TYPES = ("INSIGHTFACE",)
    FUNCTION = "load_insight_face"
    CATEGORY = "loaders"

    def load_insight_face(self):
        # This is a simplified implementation that just returns a dummy value
        # In a real setup, this would load the actual InsightFace model
        try:
            # Try to load insightface model path
            model_path = folder_paths.get_full_path("insightface", "1k3d68.onnx")
            return (model_path,)
        except:
            # Return dummy if model not found
            return ("insightface_model",)

class PulidEvaClipLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {}}
    
    RETURN_TYPES = ("EVACLIP",)
    FUNCTION = "load_evaclip"
    CATEGORY = "loaders"

    def load_evaclip(self):
        # This is a simplified implementation that just returns a dummy value
        # In a real setup, this would load the actual EVA CLIP model
        try:
            # Try to load the EVA CLIP model path
            model_path = folder_paths.get_full_path("evaclip", "EVA02-CLIP-bigE-14-plus.pt")
            return (model_path,)
        except:
            # Return dummy if model not found
            return ("evaclip_model",)

class ApplyPulid:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("PULID_MODEL",),
                "image": ("IMAGE",),
                "insightface_model": ("INSIGHTFACE",),
                "evaclip_model": ("EVACLIP",),
                "weight": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01}),
                "start_at": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                "end_at": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
            }
        }
    
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "apply_pulid"
    CATEGORY = "image/facetools"

    def apply_pulid(self, model, image, insightface_model, evaclip_model, weight, start_at, end_at):
        # This is a simplified implementation that just returns the input image
        # In a real setup, this would apply the PuLID model to the image
        return (image,)

NODE_CLASS_MAPPINGS = {
    "PulidModelLoader": PulidModelLoader,
    "PulidInsightFaceLoader": PulidInsightFaceLoader,
    "PulidEvaClipLoader": PulidEvaClipLoader,
    "ApplyPulid": ApplyPulid
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "PulidModelLoader": "Load PuLID Model",
    "PulidInsightFaceLoader": "Load InsightFace Model",
    "PulidEvaClipLoader": "Load EVA CLIP Model",
    "ApplyPulid": "Apply PuLID"
}