import torch
import base64
import io
from PIL import Image
from diffusers import StableDiffusionXLImg2ImgPipeline
import os

class EndpointHandler():
    def __init__(self, path=""):
        # 'path' is the folder where HF automatically loaded your repo files
        print("Loading ARX Pipeline...")
        
        # 1. Point directly to the model you uploaded to the repo
        model_path = os.path.join(path, "biglust.safetensors")
        
        # 2. Load the pipeline
        self.pipe = StableDiffusionXLImg2ImgPipeline.from_single_file(
            model_path,
            torch_dtype=torch.float16,
            use_safetensors=True,
            safety_checker=None
        )
        
        # 3. Load IP-Adapter (HF will download this from the public hub automatically)
        self.pipe.load_ip_adapter(
            "h94/IP-Adapter", 
            subfolder="sdxl_models", 
            weight_name="ip-adapter_sdxl.bin"
        )
        
        self.pipe.to("cuda")
        print("ARX Pipeline 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):
        """
        data param format: {"inputs": { "prompt": "...", "init_image": "..." }}
        """
        # HF wraps payloads in an "inputs" key automatically
        inputs = data.pop("inputs", data)

        prompt = inputs.get("prompt", "masterpiece, best quality")
        negative_prompt = inputs.get("negative_prompt", "lowres, bad anatomy, worst quality, ugly")
        strength = float(inputs.get("strength", 0.65))
        guidance_scale = float(inputs.get("guidance_scale", 7.0))
        num_inference_steps = int(inputs.get("steps", 25))
        ip_adapter_scale = float(inputs.get("ip_adapter_scale", 0.50))
        
        init_image_b64 = inputs.get("init_image")
        ip_adapter_image_b64 = inputs.get("ip_adapter_image")

        if not init_image_b64 or not ip_adapter_image_b64:
            return {"error": "Both init_image and ip_adapter_image must be provided."}

        init_image = self.decode_base64_image(init_image_b64).resize((1024, 1024))
        ip_image = self.decode_base64_image(ip_adapter_image_b64).resize((1024, 1024))

        self.pipe.set_ip_adapter_scale(ip_adapter_scale)

        # Generate!
        result = self.pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            image=init_image,
            ip_adapter_image=ip_image,
            strength=strength,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps
        ).images[0]

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