import os
import io
import base64
from dotenv import load_dotenv
from fastapi import FastAPI, Form
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()

def decode_base64_image(base64_str: str) -> Image.Image:
    try:
        image_data = base64.b64decode(base64_str.split(",")[-1])  # strip data URL prefix if present
        return Image.open(io.BytesIO(image_data))
    except Exception as e:
        raise ValueError("Invalid base64 image") from e

@app.post("/caption")
async def generate_caption(
    base64_image: str = Form(...),
    length: str = Form("short")
):
    try:
        img = decode_base64_image(base64_image)
        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(
    base64_image: str = Form(...),
    question: str = Form(...),
    stream: bool = Form(False)
):
    try:
        img = decode_base64_image(base64_image)

        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"]})

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