import io
import threading

import torch
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from PIL import Image
from transformers import AutoModelForImageTextToText, AutoProcessor

MODEL_ID = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct"
MAX_IMAGES = 20
MAX_FILE_BYTES = 10 * 1024 * 1024

processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForImageTextToText.from_pretrained(MODEL_ID, dtype=torch.float32).eval()
# ponytail: one generation at a time; free CPU Space has 2 vCPUs, parallel requests would just thrash.
lock = threading.Lock()

app = FastAPI(title="SmolVLM2 API")


@app.get("/")
def health():
    return {"status": "ok", "model": MODEL_ID, "max_images": MAX_IMAGES}


@app.post("/generate")
def generate(
    prompt: str = Form(...),
    images: list[UploadFile] = File(default=[]),
    max_new_tokens: int = Form(256, ge=1, le=1024),
):
    if len(images) > MAX_IMAGES:
        raise HTTPException(400, f"Max {MAX_IMAGES} images per request")

    pil_images = []
    for f in images:
        data = f.file.read(MAX_FILE_BYTES + 1)
        if len(data) > MAX_FILE_BYTES:
            raise HTTPException(413, f"{f.filename} exceeds {MAX_FILE_BYTES // 2**20}MB")
        try:
            pil_images.append(Image.open(io.BytesIO(data)).convert("RGB"))
        except Exception:
            raise HTTPException(400, f"{f.filename} is not a valid image")

    content = [{"type": "image"} for _ in pil_images] + [{"type": "text", "text": prompt}]
    text = processor.apply_chat_template([{"role": "user", "content": content}], add_generation_prompt=True)
    with lock, torch.inference_mode():
        # Splitting turns each image into up to ~17 tiles: great for 1 image, but 20 images overflow the 8k context.
        # Set on the shared processor (not a call kwarg, which transformers ignores), hence inside the lock.
        processor.image_processor.do_image_splitting = len(pil_images) == 1
        inputs = processor(text=[text], images=[pil_images] if pil_images else None, return_tensors="pt")
        out = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)

    answer = processor.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    return {"response": answer.strip(), "num_images": len(pil_images)}