import base64
import subprocess
from tempfile import NamedTemporaryFile

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

app = FastAPI()


# define request body
class RequestData(BaseModel):
    prompt: str
    image: str


@app.get("/")
def greet_json():
    return {"Hello": "World!"}


@app.post("/query")
def query(data: RequestData):
    prompt = data.prompt
    image = data.image

    try:
        # decode base64 to image
        image = base64.b64decode(image)

        with NamedTemporaryFile(delete=True, suffix=".png") as temp_image:
            temp_image.write(image)
            temp_image.flush()

            result = subprocess.run(
                ["./moondream", "--prompt", prompt, "--image", temp_image.name],
                capture_output=True,
                text=True,
            )

            if result.returncode != 0:
                raise Exception(result.stderr)

            return {"response": result.stdout}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))