import base64
import io
import torch
import spaces
from PIL import Image
from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel

_pipe = None


def _load_pipeline():
    global _pipe
    if _pipe is not None:
        return _pipe

    print("Loading ControlNet...")
    controlnet = ControlNetModel.from_pretrained(
        "xinsir/controlnet-scribble-sdxl-1.0",
        torch_dtype=torch.float16,
    )

    print("Loading SDXL...")
    _pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
        "stabilityai/stable-diffusion-xl-base-1.0",
        controlnet=controlnet,
        torch_dtype=torch.float16,
        variant="fp16",
        use_safetensors=True,
    ).to("cuda")

    _pipe.enable_attention_slicing()
    _pipe.enable_vae_slicing()

    print("Pipeline ready.")
    return _pipe


def _b64_to_pil(b64: str) -> Image.Image:
    # strip data URI prefix if present
    if "," in b64:
        b64 = b64.split(",", 1)[1]
    return Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")


def _pil_to_b64(img: Image.Image) -> str:
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()


@spaces.GPU(duration=120)
def generate(sketch_b64: str, description: str, strength: float = 0.7) -> str:
    pipe = _load_pipeline()

    sketch = _b64_to_pil(sketch_b64).resize((1024, 1024))

    negative_prompt = (
        "ugly, blurry, low quality, distorted, deformed, watermark, text, signature"
    )

    result = pipe(
        prompt=description,
        negative_prompt=negative_prompt,
        image=sketch,
        controlnet_conditioning_scale=strength,
        num_inference_steps=30,
        guidance_scale=7.5,
        height=1024,
        width=1024,
    ).images[0]

    return _pil_to_b64(result)