import curl_cffi
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import os

import uvicorn

# Load Gemini API key from secret
GEM_API_KEY = os.getenv('GEMINI_API_KEY')

# Init app
app = FastAPI()

# CORS (allow frontend)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Optional: restrict in prod
    allow_methods=["*"],
    allow_headers=["*"],
)

# Request model
class PromptRequest(BaseModel):
    structure_name: str
    structure_desc: str
    blocks_allowed: list[str]  # Enforce lowercase block names client-side

@app.post("/generate_structure")
async def generate_with_gemini(req: PromptRequest):
    try:
        # Compose the extremely strict prompt
        prompt = f"""You are a Minecraft-style structure planner. You have the curiosity to build almost anything you could think of. {req.structure_name}

{req.structure_desc}

Only output a JSON object describing a 3D structure using this format:

{{
  "width": <int>,
  "height": <int>,
  "depth": <int>,
  "layers": [
    [
      ["stone", "stone", "stone"],
      ["stone", "air",   "stone"],
      ["stone", "stone", "stone"]
    ],
    ...
  ]
}}

Only use lowercase Minecraft block IDs (e.g. "stone", "air", "glass", "planks").
You could only build this structure using {" ".join(req.blocks_allowed)}
Do not include any natural language or explanation.
Output strictly valid JSON only."""

        # Call Gemini API using curl_cffi
        session = curl_cffi.Session(impersonate="chrome110")  # Optional impersonation
        response = session.post(
            url="https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent",
            headers={
                "x-goog-api-key": GEM_API_KEY,
                "Content-Type": "application/json"
            },
            json={
                "contents": [{
                    "parts": [{
                        "text": prompt
                    }]
                }]
            }
        )

        return {"response": response.text}

    except Exception as e:
        return {"error": str(e)}

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