import torch
import base64
import io
import math
import gc
from PIL import Image, ImageOps
from diffusers import StableDiffusionXLPipeline, StableDiffusionXLInpaintPipeline, DPMSolverMultistepScheduler
from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor
import os

class EndpointHandler():
    def __init__(self, path=""):
        print("Loading ARX Elite Engine...")
        model_path = os.path.join(path, "biglust.safetensors")
        
        # 1. Load Base Model first to prevent the 4-vs-9 channel crash
        print("Loading Base Model...")
        base_pipe = StableDiffusionXLPipeline.from_single_file(
            model_path,
            torch_dtype=torch.float16,
            use_safetensors=True,
            safety_checker=None,
            low_cpu_mem_usage=True
        )
        
        # 2. Dynamically convert to Inpainting Architecture
        print("Converting Base Model to Inpainting Architecture...")
        components = base_pipe.components
        self.pipe = StableDiffusionXLInpaintPipeline(**components)
        
        # Flush RAM to prevent Out-Of-Memory crashes
        del base_pipe
        gc.collect()
        torch.cuda.empty_cache()
        
        # 3. Load the CLIP Vision Encoder (The AI's Eyes)
        print("Loading CLIP Vision Encoder...")
        self.pipe.image_encoder = CLIPVisionModelWithProjection.from_pretrained(
            "h94/IP-Adapter",
            subfolder="models/image_encoder",
            torch_dtype=torch.float16,
            low_cpu_mem_usage=True
        ).to("cuda")
        self.pipe.feature_extractor = CLIPImageProcessor()
        
        # 4. Load standard IP-Adapter for Subject Lock
        print("Loading IP-Adapter weights...")
        self.pipe.load_ip_adapter(
            "h94/IP-Adapter", 
            subfolder="sdxl_models", 
            weight_name="ip-adapter-plus_sdxl_vit-h.safetensors"
        )

        self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
            self.pipe.scheduler.config, 
            use_karras_sigmas=True,
            algorithm_type="sde-dpmsolver++"
        )
        
        self.pipe.to("cuda")
        print("ARX Inpainting + IP-Adapter Ready.")

    def decode_base64_image(self, image_string):
        if "," in image_string:
            image_string = image_string.split(",")[1]
        image_bytes = base64.b64decode(image_string)
        return Image.open(io.BytesIO(image_bytes)).convert("RGB")

    def encode_image_base64(self, image):
        buffered = io.BytesIO()
        image.save(buffered, format="PNG")
        return base64.b64encode(buffered.getvalue()).decode('utf-8')

    def __call__(self, data):
        inputs = data.pop("inputs", data)
        
        prompt = inputs.get("prompt", "masterpiece, best quality")
        negative_prompt = inputs.get("negative_prompt", "blurry, lowres, bad anatomy, worst quality, ugly")
        
        strength = float(inputs.get("strength", 0.85))
        guidance_scale = float(inputs.get("guidance_scale", 7.5))
        num_inference_steps = int(inputs.get("steps", 30)) 
        
        ip_scale = float(inputs.get("ip_scale", 0.6))
        
        init_image_b64 = inputs.get("init_image")
        mask_image_b64 = inputs.get("mask_image")
        reference_image_b64 = inputs.get("reference_image")

        if not init_image_b64 or not mask_image_b64:
            return {"error": "Missing init_image or mask_image."}

        raw_init = self.decode_base64_image(init_image_b64)
        raw_mask = self.decode_base64_image(mask_image_b64).convert("L")
        
        orig_w, orig_h = raw_init.size

        # 1. SMART SCALING: Keep aspect ratio perfectly intact while fitting SDXL's limits
        max_size = 1024
        if orig_w > orig_h:
            new_w = max_size
            new_h = int(max_size * (orig_h / orig_w))
        else:
            new_h = max_size
            new_w = int(max_size * (orig_w / orig_h))
            
        scaled_init = raw_init.resize((new_w, new_h), Image.LANCZOS)
        scaled_mask = raw_mask.resize((new_w, new_h), Image.LANCZOS)

        # 2. THE INVISIBLE PAD: Nearest multiple of 64
        pad_w = math.ceil(new_w / 64) * 64
        pad_h = math.ceil(new_h / 64) * 64
        
        init_padded = Image.new("RGB", (pad_w, pad_h), (0, 0, 0))
        init_padded.paste(scaled_init, (0, 0))
        
        mask_padded = Image.new("L", (pad_w, pad_h), 0)
        mask_padded.paste(scaled_mask, (0, 0))

        # 3. EXPLICIT DIMENSIONS: This stops the AI from squashing your photo!
        kwargs = {
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "image": init_padded,
            "mask_image": mask_padded,
            "width": pad_w,  
            "height": pad_h, 
            "strength": strength,
            "guidance_scale": guidance_scale,
            "num_inference_steps": num_inference_steps
        }

        # 4. IP-ADAPTER LOGIC
        if reference_image_b64:
            raw_ref = self.decode_base64_image(reference_image_b64)
            ip_image = ImageOps.fit(raw_ref, (224, 224), method=Image.LANCZOS)
            kwargs["ip_adapter_image"] = ip_image
            self.pipe.set_ip_adapter_scale(ip_scale)
        else:
            blank_image = Image.new("RGB", (224, 224), (0, 0, 0))
            kwargs["ip_adapter_image"] = blank_image
            self.pipe.set_ip_adapter_scale(0.0)

        # Generate!
        result_padded = self.pipe(**kwargs).images[0]
        
        # 5. SLICE OFF THE PADDING
        cropped_result = result_padded.crop((0, 0, new_w, new_h))
        
        # 6. RESTORE EXACT ORIGINAL DIMENSIONS FOR THE UI
        final_result = cropped_result.resize((orig_w, orig_h), Image.LANCZOS)

        return {"image": self.encode_image_base64(final_result)}