import json, asyncio
from aiohttp import web
from ..utils.auth import panel_auth

active_ws = set()

def setup_routes(app, prefix=''):
    app.router.add_get(f'{prefix}/api/logs/ws', panel_auth(ws_handler))

async def ws_handler(req):
    ws = web.WebSocketResponse(heartbeat=30.0)
    await ws.prepare(req)
    active_ws.add(ws)
    try:
        async for msg in ws:
            if msg.type == web.WSMsgType.TEXT:
                d = json.loads(msg.data)
                if d.get('action') == 'set_level':
                    ls = req.app.get('log_service')
                    if ls: ls.set_level(d.get('level','INFO'))
    finally: active_ws.discard(ws)
    return ws

def broadcast_log(log_record):
    if not active_ws: return
    payload = json.dumps({"type":"log", "level":log_record.get('level','INFO'), 
                          "message":log_record.get('simple_message',''), "timestamp":log_record.get('timestamp',0)})
    for ws in list(active_ws):
        if not ws.closed: asyncio.ensure_future(ws.send_str(payload))
        else: active_ws.discard(ws)