import os
import torch
import folder_paths
from comfy.model_patcher import ModelPatcher

# --- Loader Logic ---
class QwenImageEditLoaderStandalone:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model_name": (folder_paths.get_filename_list("checkpoints"),),
            }
        }

    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_qwen"
    CATEGORY = "QwenEdit/Standalone"

    def load_qwen(self, model_name):
        model_path = folder_paths.get_full_path("checkpoints", model_name)
        # This part relies on ComfyUI's internal state to map the Qwen architecture
        # In a real-world scenario, you'd ensure the 'diffusers' or 'transformers' 
        # library is used here to map the Mamad8 specific keys.
        out = comfylib_load_checkpoint(model_path) # Placeholder for internal loader
        return out

# --- The Requested Encoder Node ---
class TextEncodeQwenImageEditPlusStandalone:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "clip": ("CLIP",),
                "prompt": ("STRING", {"multiline": True}),
                "latent_target_pixel": ("INT", {"default": 1024, "min": 64, "max": 4096, "step": 64}),
            },
            "optional": {
                "vae": ("VAE",),
                "image1": ("IMAGE",),
                "image2": ("IMAGE",),
                "image3": ("IMAGE",),
            }
        }

    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "encode"
    CATEGORY = "QwenEdit/Standalone"

    def encode(self, clip, prompt, latent_target_pixel, vae=None, image1=None, image2=None, image3=None):
        # We append the latent_target_pixel to the conditioning dictionary 
        # so the sampler knows the target resolution Qwen was trained for.
        tokens = clip.tokenize(prompt)
        cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
        
        return ([[cond, {
            "pooled_output": pooled, 
            "latent_target_pixel": latent_target_pixel,
            "visual_context": [image1, image2, image3] # Passing images into conditioning
        }]], )

NODE_CLASS_MAPPINGS = {
    "QwenImageEditLoaderStandalone": QwenImageEditLoaderStandalone,
    "TextEncodeQwenImageEditPlusStandalone": TextEncodeQwenImageEditPlusStandalone
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "QwenImageEditLoaderStandalone": "Qwen Image Edit Loader (Standalone)",
    "TextEncodeQwenImageEditPlusStandalone": "Text Encode Qwen Image Edit Plus (Mamad8 Standalone)"
}