import os
import time
import mimetypes

from dotenv import load_dotenv
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles
from huggingface_hub import InferenceClient


load_dotenv()

HF_TOKEN = os.getenv("HF_TOKEN")
MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen2.5-1.5B-Instruct")
GAME_DIR = "game"

MAX_TOKENS = 80
TEMPERATURE = 0.5
TOP_P = 0.85

client = InferenceClient(token=HF_TOKEN) if HF_TOKEN else None

mimetypes.add_type("application/wasm", ".wasm")
mimetypes.add_type("application/octet-stream", ".pck")


def generate_reply(prompt: str) -> str:
    """Generate short NPC dialogue for the Godot game."""
    if client is None:
        time.sleep(0.2)
        return "The air is calm today. Bring what you find, and we will make it useful."

    start_time = time.time()
    completion = client.chat.completions.create(
        model=MODEL_ID,
        messages=[
            {
                "role": "system",
                "content": (
                    "You are an NPC in the post-nuclear RPG Shadows of Tomorrow. "
                    "Answer in plain text only, with 1 or 2 short natural sentences."
                ),
            },
            {"role": "user", "content": prompt},
        ],
        max_tokens=MAX_TOKENS,
        temperature=TEMPERATURE,
        top_p=TOP_P,
    )
    print(f"HF inference model={MODEL_ID} took={time.time() - start_time:.2f}s")
    return completion.choices[0].message.content.strip()


app = FastAPI()


@app.get("/", include_in_schema=False)
async def root():
    """Make the Space open straight into the fullscreen Godot web export."""
    return RedirectResponse(url="/game/index.html", status_code=307)


@app.get("/healthz")
async def healthz():
    return {"ok": True, "model": MODEL_ID}


@app.post("/ai-response")
async def ai_response(request: Request):
    """Contract used by Godot: {"prompt": str} -> {"output": str}."""
    try:
        data = await request.json()
    except Exception:
        return JSONResponse({"error": "Invalid JSON body"}, status_code=400)

    prompt = (data or {}).get("prompt")
    if not prompt:
        return JSONResponse({"error": "Missing 'prompt' in request body"}, status_code=400)

    try:
        return {"output": generate_reply(prompt)}
    except Exception as e:
        print(f"Inference error: {e}")
        return JSONResponse({"error": str(e)}, status_code=500)


if os.path.isdir(GAME_DIR):
    app.mount("/game", StaticFiles(directory=GAME_DIR, html=True), name="game")
else:
    print(f"WARNING: '{GAME_DIR}/' not found. Export the Godot Web build before deploying.")


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=7860)