import os
from fastapi import FastAPI, Request, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from google import genai
from google.genai import types

# --------------------
# App setup
# --------------------
app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# --------------------
# Gemini client
# --------------------
# HF → Settings → Variables and secrets
# Name: GEMINI_API_KEY
client = genai.Client()  # auto-picks GEMINI_API_KEY

DEFAULT_MODEL = "gemini-2.5-flash"

# --------------------
# Routes
# --------------------
@app.post("/api/generate")
async def generate(req: Request):
    try:
        payload = await req.json()

        model = payload.get("model", DEFAULT_MODEL)
        config = payload.get("config", None)

        # Simple mode
        if "prompt" in payload:
            contents = payload["prompt"]
        # Full SDK-compatible payload
        elif "contents" in payload:
            contents = payload["contents"]
        else:
            raise HTTPException(
                status_code=400,
                detail="Expected `prompt` or `contents`",
            )

        response = client.models.generate_content(
            model=model,
            contents=contents,
            config=config,
        )

        # Return raw structured response
        return response.model_dump()

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))