import base64
import io
import os

import runpod
from PIL import Image

from pipeline import load_pipeline


PIPELINE = load_pipeline(os.getenv("LORA_PATH", "."))


def _decode_image(data: str) -> Image.Image:
    if data.startswith("data:image"):
        data = data.split(",", 1)[1]
    return Image.open(io.BytesIO(base64.b64decode(data))).convert("RGB")


def _encode_png(image: Image.Image) -> str:
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return base64.b64encode(buffer.getvalue()).decode("utf-8")


def handler(event):
    payload = event.get("input", {})
    image_b64 = payload.get("image")
    if not image_b64:
        return {"error": "Missing input.image base64 PNG/JPEG data."}

    seed = payload.get("seed")
    if seed is not None:
        seed = int(seed)

    output = PIPELINE(
        image=_decode_image(image_b64),
        num_inference_steps=int(payload.get("num_inference_steps", 50)),
        guidance_scale=float(payload.get("guidance_scale", 7.5)),
        controlnet_conditioning_scale=float(payload.get("controlnet_conditioning_scale", 0.8)),
        strength=float(payload.get("strength", 0.6)),
        seed=seed,
    )

    return {
        "image": _encode_png(output["image"]),
        "rembg_ok": output["rembg_ok"],
    }


runpod.serverless.start({"handler": handler})