# app.py
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from PIL import Image, ImageDraw
import base64
import io
import asyncio

app = FastAPI()

# Simple in-memory game state per-connection
class GameState:
    def __init__(self):
        self.player_pos = [50, 50]
        self.canvas_size = (200, 200)
        self.step = 10

    def apply_action(self, action):
        x, y = self.player_pos
        if action == 'w': y -= self.step
        elif action == 's': y += self.step
        elif action == 'a': x -= self.step
        elif action == 'd': x += self.step
        x = max(0, min(self.canvas_size[0], x))
        y = max(0, min(self.canvas_size[1], y))
        self.player_pos = [x, y]

    def render_frame_b64(self):
        img = Image.new('RGB', self.canvas_size, (255, 255, 255))
        draw = ImageDraw.Draw(img)
        x, y = self.player_pos
        r = 6
        draw.ellipse((x-r, y-r, x+r, y+r), fill=(40, 120, 255))
        # scale up to 400x400 for client
        img = img.resize((400, 400), Image.NEAREST)
        buff = io.BytesIO()
        img.save(buff, format='PNG')
        b64 = base64.b64encode(buff.getvalue()).decode('ascii')
        return b64

@app.websocket('/ws')
async def websocket_endpoint(ws: WebSocket):
    await ws.accept()
    state = GameState()
    try:
        # send initial frame
        await ws.send_json({ 'type': 'frame', 'image_b64': state.render_frame_b64(), 'player_pos': state.player_pos })
        while True:
            data = await ws.receive_text()
            msg = None
            try:
                import json
                msg = json.loads(data)
            except Exception:
                continue
            if not isinstance(msg, dict):
                continue
            mtype = msg.get('type')
            if mtype == 'action':
                action = msg.get('action')
                state.apply_action(action)
                # return next frame
                await ws.send_json({ 'type': 'frame', 'image_b64': state.render_frame_b64(), 'player_pos': state.player_pos })
            elif mtype == 'reset':
                state = GameState()
                await ws.send_json({ 'type': 'frame', 'image_b64': state.render_frame_b64(), 'player_pos': state.player_pos })
            elif mtype == 'noop':
                await ws.send_json({ 'type': 'status', 'text': 'noop-ok' })
    except WebSocketDisconnect:
        return