1 files changed, 103 insertions, 42 deletions
diff --git a/kennel/main.py b/kennel/main.py
index 183fb64..122e40a 100644
--- a/kennel/main.py
+++ b/kennel/main.py
@@ -1,57 +1,118 @@
+from fastapi import (
+    FastAPI,
+    Request,
+    Response,
+    WebSocket,
+    WebSocketException,
+    status,
+    Cookie,
+    Depends,
+)
+from fastapi.staticfiles import StaticFiles
+from kennel.engine.systems.system import SystemType
+from kennel.engine.systems.network import Event, Publishable, EventType
+from typing import Annotated, Optional
+from .kennel import (
+    kennel,
+    system_manager,
+    entity_manager,
+    create_session_controllable_entities,
+)
+from .app import app, templates, logger
+from .kennelcats import KennelCatService
+from .middleware import logger_middleware
+from .config import config
+import asyncio
 import uuid
 
-import structlog
-from fastapi import FastAPI, Request, Response
-from fastapi.staticfiles import StaticFiles
-from fastapi.templating import Jinja2Templates
+app.mount("/static", StaticFiles(directory="static"), name="static")
 
-app = FastAPI(
-    servers = [
-        {"url": "https://kennel.hatecomputers.club", "description": "prod"}
-    ]
-)
-logger = structlog.get_logger()
+loop = asyncio.get_event_loop()
 
 
-@app.middleware("http")
-async def logger_middleware(request: Request, call_next):
-    structlog.contextvars.clear_contextvars()
-    structlog.contextvars.bind_contextvars(
-        path=request.url.path,
-        method=request.method,
-        client_host=request.client.host,
-        request_id=str(uuid.uuid4()),
-    )
-    response = await call_next(request)
+@app.on_event("startup")
+async def startup_event():
+    logger.info("Starting Kennel...")
+    loop.create_task(kennel.run())
 
-    structlog.contextvars.bind_contextvars(
-        status_code=response.status_code,
-    )
 
-    # Exclude /healthcheck endpoint from producing logs
-    if request.url.path != "/healthcheck":
-        if 400 <= response.status_code < 500:
-            logger.warn("Client error")
-        elif response.status_code >= 500:
-            logger.error("Server error")
-        else:
-            logger.info("OK")
+@app.on_event("shutdown")
+async def shutdown_event():
+    logger.info("Stopping Kennel...")
+    kennel.stop()
+    loop.stop()
+    logger.info("Kennel stopped")
 
-    return response
 
-templates = Jinja2Templates(directory="templates")
+@app.get("/")
+async def index(request: Request):
+    return templates.TemplateResponse(request=request, name="index.html")
 
 
-@app.get("/healthcheck")
-async def healthcheck():
-    return Response("hello")
+@app.get("/assign")
+async def assign(
+    response: Response,
+    session: Annotated[Optional[str], Cookie()] = None,
+):
+    if session is None:
+        session = str(uuid.uuid4().hex)
+        response.set_cookie(key="session", value=session, max_age=config.COOKIE_MAX_AGE)
+    return {"session": session}
 
 
-@app.get("/")
-async def read_main(request: Request):
-    return templates.TemplateResponse(
-        request=request, name="index.html"
-    )
+async def get_cookie_or_token(
+    websocket: WebSocket,
+    session: Annotated[str | None, Cookie()] = None,
+):
+    if session is None:
+        raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)
+    return session
+
 
-app.mount("/static", StaticFiles(directory = "static"), name = "static")
+class WebSocketClient(Publishable):
+    def __init__(self, websocket: WebSocket):
+        self.websocket = websocket
 
+    async def publish(self, event: Event):
+        await self.websocket.send_json(event.dict())
+
+
+@app.websocket("/ws")
+async def websocket_endpoint(
+    websocket: WebSocket,
+    session: Annotated[str, Depends(get_cookie_or_token)],
+):
+    await websocket.accept()
+    logger.info(f"Websocket connection established for session {session}")
+
+    session_entities = create_session_controllable_entities(session)
+    try:
+        network_system = system_manager.get_system(SystemType.NETWORK)
+        if network_system is None:
+            raise "Network system not found"
+
+        for entity in session_entities:
+            logger.info(f"Adding entity {entity.id} for session {session}")
+            entity_manager.add_entity(entity)
+            network_system.add_event(Event(EventType.ENTITY_BORN, {"id": entity.id}))
+
+        network_system.add_client(session, WebSocketClient(websocket))
+        while True:
+            message = await websocket.receive_json()
+            network_system.client_event(session, Event.from_dict(message))
+    except Exception as e:
+        logger.error(f"WebSocket exception {e}")
+    finally:
+        logger.info("Websocket connection closed")
+        for entity in session_entities:
+            logger.info(f"Removing entity {entity.id} for session {session}")
+            entity_manager.remove_entity(entity.id)
+            network_system.add_event(Event(EventType.ENTITY_DEATH, {"id": entity.id}))
+
+        network_system.remove_client(session)
+        await websocket.close()
+
+
+@app.get("/healthcheck")
+async def healthcheck():
+    return Response("healthy")