import asyncio
import json

from microdot import Microdot
from microdot.websocket import WebSocketError, with_websocket

import espnow
import network
from util import format_mac, parse_mac


app = Microdot()
_esp = None
_known_peers = set()
_ws_clients = set()


def _init_espnow():
    global _esp
    sta = network.WLAN(network.STA_IF)
    sta.active(True)
    _esp = espnow.ESPNow()
    _esp.active(True)


def _validate_envelope(obj):
    if obj.get("v") != "1":
        raise ValueError("message.v must be '1'")
    devices = obj["devices"]
    for address in devices.keys():
        parse_mac(address)
    return obj


def _send_espnow(address, payload):
    if _esp is None:
        raise ValueError("espnow is not initialized")
    mac = parse_mac(address)
    msg = json.dumps(payload, separators=(",", ":")).encode("utf-8")
    if mac not in _known_peers:
        _esp.add_peer(mac)
        _known_peers.add(mac)
    _esp.send(mac, msg)
    return mac, len(msg)


async def _broadcast_ws(obj):
    text = json.dumps(obj)
    dead = []
    for client in list(_ws_clients):
        try:
            await client.send(text)
        except Exception:
            dead.append(client)
    for client in dead:
        _ws_clients.discard(client)


async def _espnow_receive_loop():
    while True:
        host, msg = _esp.recv(0)
        if not host:
            await asyncio.sleep(0.01)
            continue
        await _broadcast_ws(
            {
                "from": format_mac(host),
                "payload": msg.decode("utf-8"),
            }
        )


@app.route("/ws")
@with_websocket
async def ws(request, ws):
    _ws_clients.add(ws)
    while True:
        try:
            raw = await ws.receive()
        except WebSocketError:
            break

        if not raw:
            break

        try:
            parsed = json.loads(raw)
            env = _validate_envelope(parsed)
            sent = []
            for address, payload in env["devices"].items():
                mac, payload_size = _send_espnow(address, payload)
                sent.append(
                    {
                        "address": format_mac(mac),
                        "bytes": payload_size,
                    }
                )
        except (ValueError, TypeError) as e:
            await ws.send(json.dumps({"ok": False, "error": str(e)}))
            continue

        await ws.send(
            json.dumps(
                {
                    "ok": True,
                    "sent": sent,
                }
            )
        )
    _ws_clients.discard(ws)


async def main(port=80):
    _init_espnow()
    asyncio.create_task(_espnow_receive_loop())
    await app.start_server(host="0.0.0.0", port=port)


if __name__ == "__main__":
    asyncio.run(main(port=80))