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

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

# Init app
app = FastAPI()

# CORS (allow frontend)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Pydantic model for manual prompting
class PromptRequest(BaseModel):
    prompt: str

# Utility to pick a random structure and build a strict prompt
def pick_random_structure():
    with open("structure_list.json", "r") as f:
        return random.choice(json.load(f))

def make_prompt(structure):
    name = structure["name"]
    desc = structure["description"]
    blocks = " ".join(structure["blocks_allowed"])
    return f"""You are a Minecraft-style structure planner. You have the curiosity to build almost anything you could think of. {name}

Structure description: {desc}

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

{{
  "width": <int>,
  "height": <int>,
  "layout": [
    ["stone", "air", "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 {blocks}
Do not include any natural language or explanation.
Output strictly valid JSON only."""

# Manual prompting route
@app.post("/generate_structure")
async def generate_with_gemini(req: PromptRequest):
    try:
        response = curl_cffi.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": req.prompt}]
                    }
                ]
            }
        )
        return {"response": response.text}
    except Exception as e:
        return {"error": str(e)}

# Auto structure generation route
@app.get("/generate_random_structure")
async def generate_random_structure():
    try:
        structure = pick_random_structure()
        prompt = make_prompt(structure)

        response = curl_cffi.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.json()
    except Exception as e:
        return {"error": str(e)}

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