import os
import torch
from diffusers import FluxTransformer2DModel, FluxFillPipeline
from diffusers.utils import load_image
from torchvision import transforms
from PIL import Image
import numpy as np


device_list = ['cuda', 'cpu']
node_dir = os.path.dirname(os.path.abspath(__file__))
comfy_dir = os.path.abspath(os.path.join(node_dir, '..', '..'))
models_dir = os.path.abspath(os.path.join(comfy_dir, 'models'))
checkpoints_dir = os.path.abspath(os.path.join(models_dir, 'checkpoints'))

# TryOffModel Node
class TryOffModelNode:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model_name": (["xiaozaa/cat-tryoff-flux"],),
                "device": (device_list,),
            }
        }

    CATEGORY = "Models"
    RETURN_TYPES = ("MODEL",)
    FUNCTION = "load_model"

    def load_model(self, model_name, device):
        model = FluxTransformer2DModel.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(device)
        return (model,)


# FluxFillModel Node
class TryOffFluxFillModelNode:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "transformer": ("MODEL",),
                "model_name": (["FLUX.1-dev"],),
                "device": (device_list,),
            }
        }

    CATEGORY = "Models"
    RETURN_TYPES = ("MODEL",)
    FUNCTION = "load_pipeline"

    def load_pipeline(self, transformer, model_name, device):
        model_path = os.path.join(checkpoints_dir, model_name)
        
        pipeline = FluxFillPipeline.from_pretrained(
            model_path,
            transformer=transformer,
            torch_dtype=torch.bfloat16,
        ).to(device)
        pipeline.enable_model_cpu_offload()
        return (pipeline,)


# TryOffRun Node
class TryOffRunNode:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image_in": ("IMAGE",),
                "mask_in": ("MASK",),
                "pipe": ("MODEL",),
                "width": ("INT", {"default": 576, "min": 128, "max": 1024, "step": 16}),
                "height": ("INT", {"default": 768, "min": 128, "max": 1024, "step": 16}),
                "num_steps": ("INT", {"default": 50, "min": 1, "max": 100}),
                "guidance_scale": ("FLOAT", {"default": 30.0, "min": 1.0, "max": 100.0, "step": 0.5}),
                "seed": ("INT", {"default": 42}),
                "prompt": ("STRING", {"multiline": True, "default": 
                        "The pair of images highlights clothing and its styling on a model, high resolution, 4K, 8K; "
                        "[IMAGE1] Detailed product shot of clothing "
                        "[IMAGE2] The same clothing is worn by a model in a lifestyle setting."}),
                "device": (device_list,),
            }
        }

    RETURN_TYPES = ("IMAGE", "IMAGE")
    RETURN_NAMES = ("garment_image", "tryoff_image")
    CATEGORY = "Processing"
    FUNCTION = "run_inference"

    def run_inference(self, image_in, mask_in, pipe, width, height, num_steps, guidance_scale, seed, prompt, device):
        pipe.transformer.to(torch.bfloat16)

        # Preprocessing transforms
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5]),
        ])
        mask_transform = transforms.Compose([
            transforms.ToTensor(),
        ])

        # Resize and preprocess
        def convert_image(tnsr):
            return Image.fromarray(np.clip(255.0 * tnsr.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)).convert("RGB")

        image = convert_image(image_in).resize((width, height))
        mask = convert_image(mask_in).resize((width, height))

        image_tensor = transform(image)
        mask_tensor = mask_transform(mask)[:1]  # Use only the first channel

        garment_tensor = torch.zeros_like(image_tensor)
        image_tensor = image_tensor * mask_tensor

        # Concatenate inputs for FluxFillPipeline
        inpaint_image = torch.cat([garment_tensor, image_tensor], dim=2)
        garment_mask = torch.zeros_like(mask_tensor)
        extended_mask = torch.cat([1 - garment_mask, garment_mask], dim=2)

        # Set random seed for reproducibility
        generator = torch.Generator(device=device).manual_seed(seed)

        # Run pipeline
        result = pipe(
            height=height,
            width=width * 2,
            image=inpaint_image,
            mask_image=extended_mask,
            num_inference_steps=num_steps,
            generator=generator,
            max_sequence_length=512,
            guidance_scale=guidance_scale,
            prompt=prompt,
        ).images[0]
        

        # Split result into garment and try-on images
        garment_result = result.crop((0, 0, width, height))
        tryoff_result = result.crop((width, 0, width * 2, height))

        tryoff_result = torch.tensor(
                    np.array(tryoff_result) / 255.0, dtype=torch.float32
                ).unsqueeze(0)
        garment_result = torch.tensor(
                    np.array(garment_result) / 255.0, dtype=torch.float32
                ).unsqueeze(0)

        return  (tryoff_result, garment_result,)