import os
import base64
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from google import genai
import uvicorn

GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")
if not GOOGLE_API_KEY:
  raise RuntimeError("GOOGLE_API_KEY environment variable is not set")

client = genai.Client(api_key=GOOGLE_API_KEY)

app = FastAPI()

class Prompt(BaseModel):
    prompt: str

@app.post("/generate")
def generate_image(data: Prompt):
    try:
        result = client.images.generate(
            model="gemini-2.0-flash-exp-image",
            prompt=data.prompt,
            size="1024x1024"
        )

        # Adjust this depending on the exact response shape of google-genai
        image_bytes = result.images[0].data  # example; adapt if needed
        b64 = base64.b64encode(image_bytes).decode("utf-8")

        return {"image": b64}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)