import os
import io
from dotenv import load_dotenv
from fastapi import FastAPI, File, UploadFile, Form, Query
from fastapi.responses import JSONResponse, StreamingResponse
from PIL import Image
import moondream as md

# Load environment variables
load_dotenv()
api_key = os.getenv("MOON_DREAM_KEY")

# Initialize Moondream model
model = md.vl(api_key=api_key)

# FastAPI app
app = FastAPI()

@app.post("/caption")
async def generate_caption(
    image: UploadFile = File(...),
    length: str = Form("short")
):
    try:
        image_bytes = await image.read()
        img = Image.open(io.BytesIO(image_bytes))
        response = model.caption(img, length=length)
        return JSONResponse(content={"caption": response["caption"]})
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})

@app.post("/query")
async def query_image(
    image: UploadFile = File(...),
    question: str = Form(...),
    stream: bool = Form(False)
):
    try:
        image_bytes = await image.read()
        img = Image.open(io.BytesIO(image_bytes))

        if stream:
            def generate():
                result = model.query(img, question, stream=True)
                for chunk in result["chunk"]:
                    yield chunk

            return StreamingResponse(generate(), media_type="text/plain")
        else:
            result = model.query(img, question)
            return JSONResponse(content={
                "answer": result["answer"],
                "request_id": result["request_id"]
            })

    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})