import os, io, base64
import spaces
import torch
import gradio as gr
from PIL import Image
from huggingface_hub import hf_hub_download
from diffusers import StableDiffusionXLPipeline
from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor

from msdiffusion.models.projection import Resampler
from msdiffusion.models.model import MSAdapter
from msdiffusion.utils import get_phrase_idx, get_eot_idx

BASE = "stabilityai/stable-diffusion-xl-base-1.0"
IMG_ENC = "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"
NUM_TOKENS = 16
DTYPE = torch.float16

print("loading SDXL…")
pipe = StableDiffusionXLPipeline.from_pretrained(BASE, torch_dtype=DTYPE, add_watermarker=False)
print("loading CLIP image encoder…")
image_encoder = CLIPVisionModelWithProjection.from_pretrained(IMG_ENC, torch_dtype=DTYPE)
image_processor = CLIPImageProcessor()

print("building resampler + MSAdapter…")
image_proj_model = Resampler(
    dim=1280, depth=4, dim_head=64, heads=20, num_queries=NUM_TOKENS,
    embedding_dim=image_encoder.config.hidden_size,
    output_dim=pipe.unet.config.cross_attention_dim, ff_mult=4,
    latent_init_mode="grounding",
    phrase_embeddings_dim=pipe.text_encoder.config.projection_dim,
).to(dtype=DTYPE)
ms_ckpt = hf_hub_download("doge1516/MS-Diffusion", "ms_adapter.bin")
ms_model = MSAdapter(pipe.unet, image_proj_model, ckpt_path=ms_ckpt, device="cpu", num_tokens=NUM_TOKENS)
ms_model.to(dtype=DTYPE)
print("models ready (CPU); GPU attaches per-call.")


def _b64_to_pil(s):
    if not s:
        return None
    if "," in s and s.strip().startswith("data:"):
        s = s.split(",", 1)[1]
    return Image.open(io.BytesIO(base64.b64decode(s))).convert("RGB").resize((512, 512))


def _phrase_idxes(phrases, prompt):
    res, cnt = [], {}
    for ph in phrases:
        k = cnt.get(ph, 0); cnt[ph] = k + 1
        res.append(get_phrase_idx(pipe.tokenizer, ph, prompt, num=k)[0])
    return res


@spaces.GPU(duration=150)
def generate(prompt, img1_b64, img2_b64, phrase1, phrase2, box1, box2, scale, seed, steps):
    dev = "cuda"
    pipe.to(dev); image_encoder.to(dev, dtype=DTYPE); ms_model.to(dev, dtype=DTYPE)
    ms_model.device = dev  # generate() places tensors on self.device; must match cuda

    subs, phrases, boxes = [], [], []
    for b64, ph, bx in ((img1_b64, phrase1, box1), (img2_b64, phrase2, box2)):
        im = _b64_to_pil(b64)
        if im is None:
            continue
        subs.append(im); phrases.append(ph.strip())
        boxes.append([float(x) for x in bx.split(",")])
    if not subs:
        return ""

    phrase_idxes = [_phrase_idxes(phrases, prompt)]
    eot_idxes = [[get_eot_idx(pipe.tokenizer, prompt)] * len(phrases)]

    images = ms_model.generate(
        pipe=pipe, pil_images=[subs], num_samples=1,
        num_inference_steps=int(steps), seed=int(seed), prompt=[prompt], scale=float(scale),
        image_encoder=image_encoder, image_processor=image_processor, boxes=[boxes],
        image_proj_type="resampler", image_encoder_type="clip",
        phrases=[phrases], drop_grounding_tokens=[0],
        phrase_idxes=phrase_idxes, eot_idxes=eot_idxes,
        height=1024, width=1024, mask_threshold=0.5, start_step=5,
    )
    buf = io.BytesIO(); images[0].save(buf, "PNG")
    return base64.b64encode(buf.getvalue()).decode()


demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="prompt"),
        gr.Textbox(label="img1_b64"), gr.Textbox(label="img2_b64"),
        gr.Textbox(label="phrase1", value="a man"), gr.Textbox(label="phrase2", value="a man"),
        gr.Textbox(label="box1", value="0.0,0.25,0.45,0.95"),
        gr.Textbox(label="box2", value="0.55,0.25,1.0,0.95"),
        gr.Number(value=0.6, label="scale"), gr.Number(value=42, label="seed"),
        gr.Number(value=30, label="steps"),
    ],
    outputs=gr.Textbox(label="result_b64"),
    title="MS-Diffusion — layout-guided multi-subject (base64 API)",
)
demo.queue(max_size=6).launch(server_name="0.0.0.0", server_port=7860)