app.py

import asyncio
import json
import time
from collections import deque
from typing import Dict, Set

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles

app = FastAPI()

# ── In-memory state ────────────────────────────────────────────────────────
MAX_HISTORY = 100
message_history: deque = deque(maxlen=MAX_HISTORY)

class Room:
    def __init__(self, name: str):
        self.name = name
        self.connections: Dict[WebSocket, str] = {}   # ws → username

    def add(self, ws: WebSocket, username: str):
        self.connections[ws] = username

    def remove(self, ws: WebSocket):
        self.connections.pop(ws, None)

    @property
    def users(self):
        return list(self.connections.values())

    async def broadcast(self, payload: dict, skip: WebSocket | None = None):
        dead = []
        for ws in list(self.connections):
            if ws is skip:
                continue
            try:
                await ws.send_json(payload)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.remove(ws)

rooms: Dict[str, Room] = {"general": Room("general")}

# ── Helpers ───────────────────────────────────────────────────────────────────────

def make_msg(kind: str, **kw) -> dict:
    return {"kind": kind, "ts": int(time.time() * 1000), **kw}

def sanitize(text: str, max_len: int = 2000) -> str:
    return text.strip()[:max_len]

# ── Routes ──────────────────────────────────────────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
async def index():
    with open("static/index.html") as f:
        return f.read()

@app.get("/rooms")
async def list_rooms():
    return [{"name": k, "users": len(v.connections)} for k, v in rooms.items()]

@app.websocket("/ws/{room_name}")
async def websocket_endpoint(ws: WebSocket, room_name: str):
    await ws.accept()

    # ── Handshake: expect { "username": "..." } ───────────────────────────────────────
    try:
        init = await asyncio.wait_for(ws.receive_json(), timeout=10)
        username = sanitize(init.get("username", ""), 32) or "anon"
    except Exception:
        await ws.close(code=4000)
        return

    # Ensure room exists
    if room_name not in rooms:
        rooms[room_name] = Room(room_name)
    room = rooms[room_name]

    # Reject duplicate usernames in the same room
    if username.lower() in [u.lower() for u in room.users]:
        await ws.send_json(make_msg("error", text="Username already taken in this room."))
        await ws.close(code=4001)
        return

    room.add(ws, username)

    # Send history + welcome
    await ws.send_json(make_msg("history", messages=list(message_history)))
    await ws.send_json(make_msg("system", text=f"Welcome, {username}! You joined #{room_name}"))
    await room.broadcast(
        make_msg("system", text=f"👋 {username} joined #{room_name}", users=room.users),
        skip=ws,
    )

    try:
        while True:
            data = await ws.receive_json()
            kind = data.get("kind", "chat")

            if kind == "chat":
                text = sanitize(data.get("text", ""))
                if not text:
                    continue
                msg = make_msg("chat", room=room_name, username=username, text=text)
                message_history.append(msg)
                await room.broadcast(msg)          # broadcast to everyone incl. sender

            elif kind == "ping":
                await ws.send_json(make_msg("pong"))

    except WebSocketDisconnect:
        pass
    finally:
        room.remove(ws)
        if not room.connections and room_name != "general":
            rooms.pop(room_name, None)
        await room.broadcast(
            make_msg("system", text=f"👋 {username} left #{room_name}", users=room.users)
        )