# ───── No7, https://wharib-microsoftbiomedvlp.hf.space/embed_image - https://wharib-microsoftbiomedvlp.hf.space/embed_text - MicrosoftBiomedVLP

import io, torch, tempfile
from pathlib import Path
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import JSONResponse, HTMLResponse
from transformers import AutoTokenizer, AutoModel
from PIL import Image

MODEL_ID = "microsoft/BiomedVLP-BioViL-T"
DEVICE = torch.device("cpu")

# this version relies on hi-ml-multimodal==0.2.1
from health_multimodal.image.utils import get_image_inference

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
text_model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True).eval()
image_engine = get_image_inference("biovil_t")

# single-process shared buffer (requires uvicorn --workers 1)
buffer = {"image": None}

@torch.no_grad()
def _text_emb(sentence: str) -> torch.Tensor:
    toks = tokenizer(sentence, return_tensors="pt", truncation=True, max_length=128)
    # Some tokenisers add token_type_ids that the model doesn't expect
    if "token_type_ids" in toks:
        toks.pop("token_type_ids")
    # BioViL-T with trust_remote_code=True exposes this helper:
    return text_model.get_projected_text_embeddings(**toks).squeeze(0)

@torch.no_grad()
def _image_emb(pil_img: Image.Image) -> torch.Tensor:
    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=True) as tmp:
        pil_img.save(tmp.name)
        emb = image_engine.get_projected_global_embedding(Path(tmp.name))
    return emb

app = FastAPI(docs_url="/docs")

@app.get("/", response_class=HTMLResponse)
async def root() -> str:
    return (
        "<h2>BioViL-T multimodal embedding API</h2>"
        "<ul>"
        "<li><code>POST /embed_image</code> – image (stores embedding; waits for text)</li>"
        "<li><code>POST /embed_text</code> – text (returns cosine_similarity if image is set, then resets)</li>"
        "</ul>"
        "<p>Requires: hi-ml-multimodal==0.2.1, uvicorn --workers 1</p>"
    )

@app.post("/embed_image")
async def embed_image(file: UploadFile = File(...)):
    try:
        pil = Image.open(io.BytesIO(await file.read())).convert("RGB")
    except Exception as e:
        raise HTTPException(400, f"Bad image: {e}")
    buffer["image"] = _image_emb(pil)
    return JSONResponse({"status": "image received, waiting for text"})

@app.post("/embed_text")
async def embed_text(text: str = Form(...)):
    text = (text or "").strip()
    if not text:
        raise HTTPException(400, "Empty text prompt.")
    if buffer["image"] is None:
        return JSONResponse({"status": "waiting for image"})
    text_vec = _text_emb(text)
    score = torch.nn.functional.cosine_similarity(buffer["image"], text_vec, dim=0).item()
    buffer["image"] = None  # reset
    return JSONResponse({"cosine_similarity": score})

@app.get("/health")
async def health():
    return {"ok": True, "has_image": buffer["image"] is not None}