admin / Syanpse-Vanguard

"""Broker abstraction over the raw-event queue.

The agent-ingest hot path does NOT parse events; it pushes the raw batch here
and acks the agent immediately. Normalizer workers consume via a consumer group
(at-least-once delivery + backpressure) and do the heavy lifting off the hot
path. Redis Streams in production; an in-process queue for single-process dev."""
from __future__ import annotations

import asyncio
from abc import ABC, abstractmethod
from dataclasses import dataclass

from sv.config import settings


@dataclass
class RawMessage:
    msg_id: str
    agent_id: str
    seq: int
    payload: bytes  # NDJSON of raw event lines


class Broker(ABC):
    @abstractmethod
    async def push_raw(self, agent_id: str, seq: int, payload: bytes) -> str: ...

    @abstractmethod
    async def ensure_group(self) -> None: ...

    @abstractmethod
    async def read_group(self, consumer: str, count: int, block_ms: int
                         ) -> list[RawMessage]: ...

    @abstractmethod
    async def ack(self, msg_id: str) -> None: ...

    @abstractmethod
    async def close(self) -> None: ...


class RedisBroker(Broker):
    def __init__(self, url: str, stream: str, group: str):
        import redis.asyncio as aioredis
        self._r = aioredis.from_url(url)
        self._stream = stream
        self._group = group

    async def ensure_group(self) -> None:
        try:
            await self._r.xgroup_create(self._stream, self._group,
                                        id="0", mkstream=True)
        except Exception as exc:  # BUSYGROUP = already exists, which is fine
            if "BUSYGROUP" not in str(exc):
                raise

    async def push_raw(self, agent_id: str, seq: int, payload: bytes) -> str:
        return await self._r.xadd(
            self._stream,
            {"agent_id": agent_id, "seq": str(seq), "payload": payload},
        )

    async def read_group(self, consumer: str, count: int, block_ms: int
                         ) -> list[RawMessage]:
        resp = await self._r.xreadgroup(
            self._group, consumer, {self._stream: ">"},
            count=count, block=block_ms,
        )
        out: list[RawMessage] = []
        for _stream, entries in resp or []:
            for msg_id, fields in entries:
                out.append(RawMessage(
                    msg_id=msg_id.decode() if isinstance(msg_id, bytes) else msg_id,
                    agent_id=_s(fields[b"agent_id"]),
                    seq=int(_s(fields[b"seq"])),
                    payload=fields[b"payload"],
                ))
        return out

    async def ack(self, msg_id: str) -> None:
        await self._r.xack(self._stream, self._group, msg_id)

    async def close(self) -> None:
        await self._r.aclose()


class MemoryBroker(Broker):
    """Single-process fallback for tests / `SV_BROKER=memory`."""
    _q: "asyncio.Queue[RawMessage]" = None  # type: ignore

    def __init__(self):
        self._q = asyncio.Queue()
        self._counter = 0

    async def ensure_group(self) -> None:
        return None

    async def push_raw(self, agent_id: str, seq: int, payload: bytes) -> str:
        self._counter += 1
        mid = f"mem-{self._counter}"
        await self._q.put(RawMessage(mid, agent_id, seq, payload))
        return mid

    async def read_group(self, consumer: str, count: int, block_ms: int
                         ) -> list[RawMessage]:
        try:
            first = await asyncio.wait_for(self._q.get(), timeout=block_ms / 1000)
        except asyncio.TimeoutError:
            return []
        batch = [first]
        while len(batch) < count and not self._q.empty():
            batch.append(self._q.get_nowait())
        return batch

    async def ack(self, msg_id: str) -> None:
        return None

    async def close(self) -> None:
        return None


def build_broker() -> Broker:
    if settings.broker == "memory":
        return MemoryBroker()
    if settings.broker == "redis":
        return RedisBroker(settings.redis_url, settings.raw_stream, settings.raw_group)
    raise ValueError(f"unknown SV_BROKER={settings.broker!r}")


def _s(v) -> str:
    return v.decode() if isinstance(v, (bytes, bytearray)) else str(v)