MeshDD-Bot/meshbot/webserver.py

import asyncio
import json
import logging
import os

from aiohttp import web

from meshbot import config
from meshbot.database import Database

logger = logging.getLogger(__name__)

STATIC_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "static")


class WebSocketManager:
    def __init__(self):
        self.clients: set[web.WebSocketResponse] = set()

    async def broadcast(self, msg_type: str, data: dict | list):
        message = json.dumps({"type": msg_type, "data": data})
        closed = set()
        for ws in self.clients:
            try:
                await ws.send_str(message)
            except Exception:
                closed.add(ws)
        self.clients -= closed


class WebServer:
    def __init__(self, db: Database, ws_manager: WebSocketManager, bot=None, scheduler=None):
        self.db = db
        self.ws_manager = ws_manager
        self.bot = bot
        self.scheduler = scheduler
        self.app = web.Application()
        self._setup_routes()

    def _setup_routes(self):
        self.app.router.add_get("/ws", self._ws_handler)
        self.app.router.add_get("/api/nodes", self._api_nodes)
        self.app.router.add_get("/api/messages", self._api_messages)
        self.app.router.add_get("/api/stats", self._api_stats)
        self.app.router.add_get("/api/scheduler/jobs", self._api_scheduler_get)
        self.app.router.add_post("/api/scheduler/jobs", self._api_scheduler_add)
        self.app.router.add_put("/api/scheduler/jobs/{name}", self._api_scheduler_update)
        self.app.router.add_delete("/api/scheduler/jobs/{name}", self._api_scheduler_delete)
        self.app.router.add_post("/api/send", self._api_send)
        self.app.router.add_get("/scheduler", self._serve_scheduler)
        self.app.router.add_get("/map", self._serve_map)
        self.app.router.add_get("/", self._serve_index)
        self.app.router.add_static("/static", STATIC_DIR)

    async def _ws_handler(self, request: web.Request) -> web.WebSocketResponse:
        ws = web.WebSocketResponse()
        await ws.prepare(request)
        self.ws_manager.clients.add(ws)
        logger.info("WebSocket client connected (%d total)", len(self.ws_manager.clients))

        try:
            # Send initial data
            nodes = await self.db.get_all_nodes()
            await ws.send_str(json.dumps({"type": "initial", "data": nodes}))

            stats = await self.db.get_stats()
            stats["version"] = config.get("version", "0.0.0")
            await ws.send_str(json.dumps({"type": "stats_update", "data": stats}))

            if self.bot:
                channels = self.bot.get_channels()
                await ws.send_str(json.dumps({"type": "channels", "data": channels}))

            messages = await self.db.get_recent_messages(50)
            await ws.send_str(json.dumps({"type": "initial_messages", "data": messages}))

            async for msg in ws:
                pass  # We only send, not receive
        finally:
            self.ws_manager.clients.discard(ws)
            logger.info("WebSocket client disconnected (%d remaining)", len(self.ws_manager.clients))

        return ws

    async def _api_nodes(self, request: web.Request) -> web.Response:
        nodes = await self.db.get_all_nodes()
        return web.json_response(nodes)

    async def _api_messages(self, request: web.Request) -> web.Response:
        limit = int(request.query.get("limit", "50"))
        messages = await self.db.get_recent_messages(limit)
        return web.json_response(messages)

    async def _api_stats(self, request: web.Request) -> web.Response:
        stats = await self.db.get_stats()
        stats["version"] = config.get("version", "0.0.0")
        return web.json_response(stats)

    async def _api_send(self, request: web.Request) -> web.Response:
        if not self.bot:
            return web.json_response({"error": "Bot not available"}, status=503)
        data = await request.json()
        text = data.get("text", "").strip()
        channel = int(data.get("channel", 0))
        if not text:
            return web.json_response({"error": "Text is required"}, status=400)
        await self.bot.send_message(text, channel)
        return web.json_response({"ok": True})

    async def _serve_index(self, request: web.Request) -> web.Response:
        return web.FileResponse(os.path.join(STATIC_DIR, "index.html"))

    async def _serve_map(self, request: web.Request) -> web.Response:
        return web.FileResponse(os.path.join(STATIC_DIR, "map.html"))

    async def _serve_scheduler(self, request: web.Request) -> web.Response:
        return web.FileResponse(os.path.join(STATIC_DIR, "scheduler.html"))

    async def _api_scheduler_get(self, request: web.Request) -> web.Response:
        if not self.scheduler:
            return web.json_response([], status=200)
        return web.json_response(self.scheduler.get_jobs())

    async def _api_scheduler_add(self, request: web.Request) -> web.Response:
        if not self.scheduler:
            return web.json_response({"error": "Scheduler not available"}, status=503)
        job = await request.json()
        jobs = self.scheduler.add_job(job)
        if self.ws_manager:
            await self.ws_manager.broadcast("scheduler_update", jobs)
        return web.json_response(jobs, status=201)

    async def _api_scheduler_update(self, request: web.Request) -> web.Response:
        if not self.scheduler:
            return web.json_response({"error": "Scheduler not available"}, status=503)
        name = request.match_info["name"]
        updates = await request.json()
        jobs = self.scheduler.update_job(name, updates)
        if jobs is None:
            return web.json_response({"error": "Job not found"}, status=404)
        if self.ws_manager:
            await self.ws_manager.broadcast("scheduler_update", jobs)
        return web.json_response(jobs)

    async def _api_scheduler_delete(self, request: web.Request) -> web.Response:
        if not self.scheduler:
            return web.json_response({"error": "Scheduler not available"}, status=503)
        name = request.match_info["name"]
        jobs = self.scheduler.delete_job(name)
        if jobs is None:
            return web.json_response({"error": "Job not found"}, status=404)
        if self.ws_manager:
            await self.ws_manager.broadcast("scheduler_update", jobs)
        return web.json_response(jobs)

    async def start(self, host: str, port: int):
        runner = web.AppRunner(self.app)
        await runner.setup()
        site = web.TCPSite(runner, host, port)
        await site.start()
        logger.info("Webserver started at http://%s:%d", host, port)
        return runner