"""CLIP embedding API for TryON — Gradio Space (free CPU tier).

Gradio auto-exposes each function as a REST endpoint:
  POST /gradio_api/call/embed_text   data: ["red floral dress"]
  POST /gradio_api/call/embed_image  data: ["<base64 jpeg/png>"]

Both return a JSON string containing a 512-dim CLIP vector, from the same
clip-ViT-B-32 model used to build the catalog index, so text, screenshots,
and product images all land in one vector space.
"""

import base64
import io
import json

import gradio as gr
import spaces
from PIL import Image
from sentence_transformers import SentenceTransformer

# device must be pinned: ZeroGPU fakes cuda-availability in the main process,
# and encoding on the fake device silently returns all-zero vectors
model = SentenceTransformer("clip-ViT-B-32", device="cpu")


# ZeroGPU hardware refuses to start unless at least one @spaces.GPU function
# exists. The real endpoints stay CPU-only so API calls never hit GPU quotas.
@spaces.GPU(duration=5)
def gpu_warmup() -> str:
    return "ok"


def embed_text(text: str) -> str:
    text = (text or "").strip()
    if not text:
        return json.dumps({"error": "text is empty"})
    vec = model.encode(text[:300])
    return json.dumps({"embedding": vec.tolist()})


def embed_image(image_base64: str) -> str:
    try:
        raw = base64.b64decode((image_base64 or "").split(",")[-1])
        img = Image.open(io.BytesIO(raw)).convert("RGB")
    except Exception:
        return json.dumps({"error": "could not decode image"})
    img.thumbnail((512, 512))
    vec = model.encode(img)
    return json.dumps({"embedding": vec.tolist()})


with gr.Blocks(title="TryON CLIP API") as demo:
    gr.Markdown("# TryON CLIP Embedding API\n512-dim clip-ViT-B-32 vectors for fashion search.")

    with gr.Tab("Text"):
        t_in = gr.Textbox(label="Text", placeholder="red floral summer dress")
        t_out = gr.Textbox(label="Embedding JSON")
        t_btn = gr.Button("Embed")
        t_btn.click(embed_text, inputs=t_in, outputs=t_out, api_name="embed_text")

    with gr.Tab("Image"):
        i_in = gr.Textbox(label="Image as base64", placeholder="paste base64 of a jpeg/png")
        i_out = gr.Textbox(label="Embedding JSON")
        i_btn = gr.Button("Embed")
        i_btn.click(embed_image, inputs=i_in, outputs=i_out, api_name="embed_image")

demo.launch()