import json
import random
from typing import Dict

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse

app = FastAPI()
clients: Dict[str, WebSocket] = {}
player_states: Dict[str, dict] = {}
colors = ["#ff2244", "#00ffaa", "#ffcc00", "#00ccff", "#ff66ff", "#66ff66"]


def get_state_broadcast() -> dict:
    return {
        "t": "state",
        "players": [
            {
                "id": pid,
                "pos": s["pos"],
                "rot": s["rot"],
                "name": s.get("name", "Agent"),
                "color": s.get("color", "#ff2244"),
            }
            for pid, s in player_states.items()
        ],
    }


async def broadcast(sender: str, payload: dict) -> None:
    message = json.dumps(payload)
    dead_clients = []
    for pid, ws in clients.items():
        if pid == sender:
            continue
        try:
            await ws.send_text(message)
        except Exception:
            dead_clients.append(pid)
    for pid in dead_clients:
        clients.pop(pid, None)
        player_states.pop(pid, None)


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket) -> None:
    await websocket.accept()
    player_id = f"player_{id(websocket)}"
    color = random.choice(colors)

    clients[player_id] = websocket
    player_states[player_id] = {
        "pos": [0.0, 1.7, 0.0],
        "rot": [0.0, 0.0, 0.0, 1.0],
        "color": color,
        "name": f"Agent-{random.randint(100, 999)}",
    }

    await websocket.send_text(json.dumps({"t": "id", "id": player_id, "color": color}))
    await broadcast(player_id, get_state_broadcast())

    try:
        while True:
            data = await websocket.receive_text()
            msg = json.loads(data)
            msg_type = msg.get("t")

            if msg_type == "pos" and player_id in player_states:
                player_states[player_id]["pos"] = msg.get("pos", [0.0, 1.7, 0.0])
                player_states[player_id]["rot"] = msg.get("rot", [0.0, 0.0, 0.0, 1.0])
                await broadcast(player_id, get_state_broadcast())
            elif msg_type == "shoot":
                await broadcast(
                    player_id,
                    {
                        "t": "shoot",
                        "origin": msg.get("origin"),
                        "dir": msg.get("dir"),
                        "weapon": msg.get("weapon", 0),
                    },
                )
            elif msg_type == "kill":
                await broadcast(
                    player_id,
                    {
                        "t": "kill",
                        "killer": msg.get("killer", "Unknown"),
                        "victim": msg.get("victim", "Bot"),
                    },
                )
            elif msg_type == "name" and player_id in player_states:
                player_states[player_id]["name"] = msg.get("name", "Agent")

    except WebSocketDisconnect:
        clients.pop(player_id, None)
        player_states.pop(player_id, None)
        await broadcast(player_id, get_state_broadcast())


@app.get("/")
async def root() -> FileResponse:
    return FileResponse("index.html")


@app.get("/{path:path}")
async def catchall(path: str) -> FileResponse:  # noqa: ARG001
    return FileResponse("index.html")