# server.py
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
import httpx
import json

app = FastAPI()

HF_MODEL_URL = "https://api-inference.huggingface.co/models/facebook/opt-1.3b"

@app.post("/ai")
async def ai_endpoint(request: Request):
    """
    Receives the Geometry Dash editor state as JSON:
    { "state": "<description of editor>" }
    Returns a JSON action:
    { "mouse_x": float, "mouse_y": float, "click": bool }
    """
    try:
        data = await request.json()
        editor_state = data.get("state", "No state provided")

        # Build prompt to instruct the model to respond with JSON for mouse actions
        prompt = (
            f"You are an AI controlling the Geometry Dash editor.\n"
            f"Respond ONLY with JSON in this exact format:\n"
            f'{{"mouse_x": <float>, "mouse_y": <float>, "click": <true/false>}}\n'
            f"Example: {{\"mouse_x\": 100.0, \"mouse_y\": 200.0, \"click\": true}}\n"
            f"Editor state: {editor_state}"
        )

        # Send POST request to Hugging Face inference endpoint
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.post(
                HF_MODEL_URL,
                headers={"Content-Type": "application/json"},
                json={"inputs": prompt}
            )
            result = response.json()

        # Extract text output from Hugging Face response
        text_output = ""
        if isinstance(result, list) and len(result) > 0 and "generated_text" in result[0]:
            text_output = result[0]["generated_text"]

        # Attempt to parse JSON from model output
        try:
            # Extract first JSON object from text output
            json_start = text_output.find("{")
            json_end = text_output.rfind("}") + 1
            json_str = text_output[json_start:json_end]
            action = json.loads(json_str)
        except Exception:
            # Fallback if parsing fails
            action = {"mouse_x": 150.0, "mouse_y": 100.0, "click": True}

        return JSONResponse(content=action)

    except Exception as e:
        # Return default safe action on error
        return JSONResponse(content={"mouse_x": 0.0, "mouse_y": 0.0, "click": False, "error": str(e)})

# Run server locally for testing
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("server:app", host="0.0.0.0", port=8080, reload=True)