"""
Implements an alarm aggregator.

Alarms consist of a condition object, a path, and some data.
They can be raised, updated, or lowered.

The aggregator's job is to collect the alarms so that the server can poll
a single object.

Alarm conditions are subclasses of exceptions.
"""

from __future__ import annotations

from moat.util import merge
from moat.lib.micro import AC_use, Event, Queue, TaskGroup, WouldBlock
from moat.lib.path import Path
from moat.lib.rpc import BaseCmd

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from moat.lib.micro import _TaskGroupProto
    from moat.lib.rpc import Msg

    from collections.abc import Iterator


__all__ = [
    "Alert",
    "AlertHandler",
]


class Alert(Exception):
    """Alert wrapper. The path isn't stored in it."""

    def __init__(self, data):
        self.data = data


class AlertIter:
    q: Queue
    xal: Iterator[tuple[tuple[type[Alert], Path], Alert]]

    def __init__(self, ah: AlertHandler, s: bool | None):
        self.ah = ah
        self.s = s

    async def __aenter__(self):
        if hasattr(self, "q"):
            raise RuntimeError("already on")
        self.q = q = Queue(10)
        self.ah.q.add(q)
        self.xal = iter(() if self.s is False else list(self.ah.alarms.items()))
        return self

    async def __aexit__(self, *tb):
        self.ah.q.discard(self.q)
        del self.q

    def __aiter__(self):
        return self

    async def __anext__(self):
        a = None
        try:
            k, al = next(self.xal)
        except AttributeError:
            pass
        except StopIteration:
            del self.xal
            if self.s:
                raise StopAsyncIteration from None
        else:
            a, p = k
            d = al.data

        if a is None:
            try:
                q = self.q
                if q is None:
                    raise RuntimeError("not started")
                a, p, d = await q.get()
            except EOFError:
                raise StopAsyncIteration  # noqa:B904,RUF100
        res = {"a": a, "p": p}
        if d is not None:
            res["d"] = d
        return res


class AlertHandler(BaseCmd):
    """
    Collect open alerts.

    This helper class keeps track of open alerts and lets multiple clients
    receive them.

    This command can monitor + redistribute alerts from satellites::

        apps:
          al: link.Alert
          r: …
        al:
          mon:
          - rem: !P r
            al: !P a
        r:
          cfg:
            a: link.Alert

    The rem/al split is used because incoming alert paths are prefixed with
    the "rem" path so that the destination is known and unique.
    """

    q: set[Queue]
    tg: _TaskGroupProto

    def __init__(self, cfg):
        super().__init__(cfg)
        self.alarms: dict[tuple[type[Alert], Path], Alert] = {}
        self.q = set()
        self._mon = {}

    async def setup(self) -> None:  # noqa:D102
        await super().setup()
        await self._start_mon()

    async def _rdr(self, p: dict, evt: Event) -> None:
        rem = p["rem"]
        al = p["al"]
        async with self.root.cmd(Path.build(rem + al + ("r",))).stream_in() as it:
            evt.set()
            async for res in it:
                await self.cmd_w(a=res["a"], p=rem + res["p"], d=res.get("d", None))

    async def _start_mon(self) -> None:
        if (pl := self.cfg.get("mon", None)) is not None:
            if getattr(self, "tg", None) is None:
                self.tg = await AC_use(self, TaskGroup())
            tg = self.tg

            evs = []
            for k, v in pl.items():
                evt = Event()

                async def _run(v=v, evt=evt) -> None:
                    await self._rdr(v, evt)

                self._mon[k] = await tg.spawn(_run)
                evs.append(evt)
            for e in evs:
                await e.wait()

        else:
            pl = {}

        for k, v in list(self._mon.items()):
            if k not in pl:
                v.cancel()
                del self._mon[k]

    async def reload(self) -> None:  # noqa:D102
        await self._start_mon()

    doc_r = dict(_d="Stream curren alerts", s="bool:stop(default:wait)", _o="alert")

    async def stream_r(self, msg: Msg) -> None:
        """read open alarms.

        If @s ("static") is True, send a snapshot.
        Otherwise iterate on new alerts.
        """
        async with msg.stream_out(), AlertIter(self, msg.get("s", False)) as alit:
            async for al in alit:
                await msg.send(**al)

    doc_w = dict(_d="set alert", _0="type:class", _1="path", d="any:data, clears if missing")

    async def cmd_w(self, a: type[Alert], p: Path, d: dict | None = None) -> None:
        """
        Set an alert.

        @a: alarm class (a type, *not* an object)
        @p: path to the faulting object
        @d: data describing the current state

        If @d is `None`, the alert is cleared.
        """
        k = (a, p)
        al = self.alarms.get(k, None)
        if al is None:
            if d is None:
                return  # nothing to do
            self.alarms[k] = a(d)
        else:
            if d is not None:
                merge(al.data, d)
            else:
                del self.alarms[k]

        for q in list(self.q):
            try:
                q.put_nowait((a, p, d))
            except WouldBlock:
                self.q.discard(q)
                q.close_sender()

    doc_cl = dict(_d="close iters")

    async def cmd_cl(self):
        """
        Close all iterators of this alert handler.

        Used mainly for testing.
        """
        q, self.q = self.q, set()
        while q:
            q.pop().close_sender()
