"""
generation-space — SDXL + ControlNet-depth + IP-Adapter room renderer.

/generate endpoint accepts:
  prompt                        : str
  negative_prompt               : str
  depth_image_base64            : str  (grayscale PNG, base64)
  ip_adapter_images_json        : str  (JSON array of product image URLs)
  controlnet_conditioning_scale : float
  ip_adapter_scale              : float
  seed                          : int

Returns: {"image_base64": str}  (rendered room, JPEG, base64)

ZeroGPU pattern: load all models on CPU at startup (no device_map, no cpu_offload),
move to CUDA only inside @spaces.GPU. After inference, move back to CPU to free VRAM.

VRAM budget on A10G (24 GB):
  SDXL base          ~7 GB fp16
  ControlNet-depth   ~1.5 GB fp16
  IP-Adapter         ~0.5 GB (UNet patch)
  Total              ~9 GB  — comfortable on A10G
"""

import base64
import io
import json

import gradio as gr
import httpx
import numpy as np
import spaces
import torch
from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline
from PIL import Image

SDXL_MODEL  = "stabilityai/stable-diffusion-xl-base-1.0"
CTRL_MODEL  = "diffusers/controlnet-depth-sdxl-1.0"
IPA_REPO    = "h94/IP-Adapter"
IPA_SUBDIR  = "sdxl_models"
IPA_WEIGHTS = "ip-adapter_sdxl.bin"

# ---------------------------------------------------------------------------
# Load on CPU at startup — ZeroGPU moves to GPU inside @spaces.GPU
# ---------------------------------------------------------------------------
print(f"Loading ControlNet: {CTRL_MODEL}...")
controlnet = ControlNetModel.from_pretrained(CTRL_MODEL, torch_dtype=torch.float16)

print(f"Loading SDXL: {SDXL_MODEL}...")
pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
    SDXL_MODEL,
    controlnet=controlnet,
    torch_dtype=torch.float16,
    use_safetensors=True,
)

print(f"Loading IP-Adapter...")
pipe.load_ip_adapter(IPA_REPO, subfolder=IPA_SUBDIR, weight_name=IPA_WEIGHTS)
print("Generation pipeline ready.")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _decode_b64_image(b64: str) -> Image.Image:
    return Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")


def _fetch_image(url: str) -> Image.Image:
    resp = httpx.get(url, timeout=15, follow_redirects=True)
    resp.raise_for_status()
    return Image.open(io.BytesIO(resp.content)).convert("RGB")


OUTPUT_W = 1216  # standard SDXL landscape — 3:2, ~1MP, consistent across all inputs
OUTPUT_H = 832


def _image_to_b64_jpeg(image: Image.Image) -> str:
    buf = io.BytesIO()
    image.save(buf, format="JPEG", quality=90)
    return base64.b64encode(buf.getvalue()).decode()


# ---------------------------------------------------------------------------
# Generation — GPU allocated for this function only
# ---------------------------------------------------------------------------

@spaces.GPU(duration=90)
def generate(
    prompt: str,
    negative_prompt: str,
    depth_image_base64: str,
    ip_adapter_images_json: str,
    controlnet_conditioning_scale: float,
    ip_adapter_scale: float,
    seed: int,
) -> dict:
    import traceback
    try:
        device = "cuda"
        pipe.to(device)
        print(f"generate: prompt={prompt[:60]}")

        control_image = _decode_b64_image(depth_image_base64).resize((OUTPUT_W, OUTPUT_H))
        print(f"generate: control_image resized to {OUTPUT_W}x{OUTPUT_H}")

        image_urls: list[str] = json.loads(ip_adapter_images_json)
        print(f"generate: fetching {len(image_urls)} IP-Adapter images")
        ip_images: list[Image.Image] = []
        for url in image_urls:
            try:
                ip_images.append(_fetch_image(url).resize((224, 224)))
                print(f"  fetched: {url[:60]}")
            except Exception as e:
                print(f"  WARNING: could not fetch {url[:60]}: {e}")
        if not ip_images:
            print("  WARNING: no IP images fetched, using white fallback")
            ip_images = [Image.new("RGB", (224, 224), (255, 255, 255))]

        print(f"generate: {len(ip_images)} IP images ready, running SDXL...")
        pipe.set_ip_adapter_scale(float(ip_adapter_scale))
        generator = torch.Generator(device=device).manual_seed(int(seed))

        # width/height inferred from control_image size (OUTPUT_W x OUTPUT_H)
        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            image=control_image,
            ip_adapter_image=[ip_images],  # wrap in list: multiple refs for 1 adapter
            num_inference_steps=30,
            guidance_scale=7.5,
            controlnet_conditioning_scale=float(controlnet_conditioning_scale),
            generator=generator,
        ).images[0]
        print("generate: SDXL done")

        pipe.to("cpu")
        torch.cuda.empty_cache()

        return {"image_base64": _image_to_b64_jpeg(result)}

    except Exception as e:
        traceback.print_exc()
        try:
            pipe.to("cpu")
            torch.cuda.empty_cache()
        except Exception:
            pass
        raise ValueError(f"generate failed: {type(e).__name__}: {e}") from e


# ---------------------------------------------------------------------------
# Gradio interface
# ---------------------------------------------------------------------------

with gr.Blocks(title="Generation Space") as demo:
    gr.Markdown("## SDXL + ControlNet-Depth + IP-Adapter Room Generator")

    with gr.Row():
        with gr.Column():
            prompt_in     = gr.Textbox(label="prompt", value="A Japandi living room, warm tones")
            neg_in        = gr.Textbox(label="negative_prompt", value="ugly, blurry, unrealistic")
            depth_in      = gr.Textbox(label="depth_image_base64", lines=3)
            ip_in         = gr.Textbox(label="ip_adapter_images_json", value="[]")
            ctrl_scale_in = gr.Slider(0.0, 1.5, value=0.7, label="controlnet_conditioning_scale")
            ip_scale_in   = gr.Slider(0.0, 1.0, value=0.5, label="ip_adapter_scale")
            seed_in       = gr.Number(value=42, label="seed", precision=0)
            gr.Button("Generate").click(
                generate,
                inputs=[prompt_in, neg_in, depth_in, ip_in, ctrl_scale_in, ip_scale_in, seed_in],
                outputs=gr.JSON(label="result"),
                api_name="generate",
            )

demo.launch(show_error=True)