PolyWeather/web/redis_realtime_event_store.py

"""Redis Stream-backed realtime observation event store."""

from __future__ import annotations

import json
import os
import socket
import threading
import time
from typing import Any, Callable, Dict, Iterable, List, Optional, Set

from loguru import logger

from web.realtime_event_store import MAX_REPLAY_LIMIT, TIME_CONTRACT_KEYS
from web.realtime_patch_schema import EVENT_TYPE


DEFAULT_STREAM_KEY = "stream:city_observation"
DEFAULT_COUNTER_KEY = "counter:city_observation_revision"
DEFAULT_MAXLEN = 50000
DEFAULT_SOCKET_TIMEOUT_SECONDS = 15.0
DEFAULT_SOCKET_CONNECT_TIMEOUT_SECONDS = 5.0
DEFAULT_REPLAY_CHUNK_SIZE = 512
DEFAULT_REPLAY_SCAN_MAX = 5000

APPEND_EVENT_SCRIPT = """
local revision = redis.call('INCR', KEYS[2])
local stream_id = redis.call(
  'XADD', KEYS[1], 'MAXLEN', '~', ARGV[1], '*',
  'revision', revision,
  'type', ARGV[2],
  'schema_type', ARGV[3],
  'schema_version', ARGV[4],
  'city', ARGV[5],
  'source', ARGV[6],
  'obs_time', ARGV[7],
  'payload_json', ARGV[8],
  'created_at_ms', ARGV[9],
  'ts', ARGV[10],
  'producer_id', ARGV[11]
)
return {revision, stream_id}
"""


def _decode(value: Any) -> str:
    if isinstance(value, bytes):
        return value.decode("utf-8")
    return str(value or "")


def _normalize_city_set(cities: Optional[Iterable[str]]) -> Set[str]:
    return {str(city or "").strip().lower() for city in (cities or set()) if str(city or "").strip()}


def _time_contract_from_payload(payload: Dict[str, Any]) -> Dict[str, Any]:
    return {key: payload[key] for key in TIME_CONTRACT_KEYS if key in payload}


def _int_or_zero(value: Any) -> int:
    try:
        return int(value)
    except (TypeError, ValueError):
        return 0


def _float_env(name: str, default: float) -> float:
    try:
        return float(os.getenv(name) or default)
    except (TypeError, ValueError):
        return default


def _int_env(name: str, default: int, *, min_value: int, max_value: int) -> int:
    try:
        value = int(os.getenv(name) or default)
    except (TypeError, ValueError):
        value = default
    return max(min_value, min(max_value, int(value)))


class RedisRealtimeEventStore:
    """Persist replayable observation patch events in a Redis Stream."""

    uses_external_live_fanout = True

    def __init__(
        self,
        *,
        redis_url: Optional[str] = None,
        redis_client: Any = None,
        stream_key: Optional[str] = None,
        counter_key: Optional[str] = None,
        maxlen: Optional[int] = None,
        producer_id: Optional[str] = None,
    ) -> None:
        self.stream_key = stream_key or os.getenv("POLYWEATHER_REDIS_STREAM_KEY") or DEFAULT_STREAM_KEY
        self.counter_key = counter_key or os.getenv("POLYWEATHER_REDIS_COUNTER_KEY") or DEFAULT_COUNTER_KEY
        self.maxlen = max(1, int(maxlen or os.getenv("POLYWEATHER_REDIS_STREAM_MAXLEN") or DEFAULT_MAXLEN))
        self.replay_chunk_size = _int_env(
            "POLYWEATHER_REDIS_REPLAY_CHUNK_SIZE",
            DEFAULT_REPLAY_CHUNK_SIZE,
            min_value=50,
            max_value=2000,
        )
        self.replay_scan_max = _int_env(
            "POLYWEATHER_REDIS_REPLAY_SCAN_MAX",
            DEFAULT_REPLAY_SCAN_MAX,
            min_value=500,
            max_value=50000,
        )
        self.producer_id = producer_id or os.getenv("POLYWEATHER_INSTANCE_ID") or socket.gethostname()
        self._client = redis_client or self._build_client(redis_url)
        self._subscriber_lock = threading.Lock()
        self._subscriber_thread: Optional[threading.Thread] = None
        self._subscriber_stop: Optional[threading.Event] = None

    @staticmethod
    def _build_client(redis_url: Optional[str]) -> Any:
        try:
            import redis  # type: ignore
        except ImportError as exc:
            raise RuntimeError("redis package is required for Redis realtime event store") from exc

        url = redis_url or os.getenv("POLYWEATHER_REDIS_URL") or "redis://127.0.0.1:6379/0"
        client = redis.Redis.from_url(
            url,
            socket_timeout=_float_env("POLYWEATHER_REDIS_SOCKET_TIMEOUT_SECONDS", DEFAULT_SOCKET_TIMEOUT_SECONDS),
            socket_connect_timeout=_float_env(
                "POLYWEATHER_REDIS_SOCKET_CONNECT_TIMEOUT_SECONDS",
                DEFAULT_SOCKET_CONNECT_TIMEOUT_SECONDS,
            ),
            health_check_interval=30,
        )
        client.ping()
        return client

    def append_event(self, event: Dict[str, Any]) -> Dict[str, Any]:
        if event.get("type") != EVENT_TYPE:
            raise ValueError("unsupported realtime event type")
        payload = event.get("payload")
        if not isinstance(payload, dict):
            raise ValueError("event payload must be an object")

        created_at_ms = int(time.time() * 1000)
        ts = int(event.get("ts") or created_at_ms)
        payload_json = json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
        result = self._client.eval(
            APPEND_EVENT_SCRIPT,
            2,
            self.stream_key,
            self.counter_key,
            self.maxlen,
            str(event["type"]),
            str(event["schema_type"]),
            int(event["schema_version"]),
            str(event["city"]),
            str(event["source"]),
            str(event.get("obs_time") or ""),
            payload_json,
            created_at_ms,
            ts,
            self.producer_id,
        )
        revision = int(_decode(result[0] if isinstance(result, (list, tuple)) else result))
        return {
            "type": event["type"],
            "revision": revision,
            "city": str(event["city"]),
            "source": str(event["source"]),
            "obs_time": event.get("obs_time"),
            **_time_contract_from_payload(payload),
            "ts": ts,
            "payload": payload,
        }

    def latest_revision(self) -> int:
        value = self._client.get(self.counter_key)
        revision = _int_or_zero(_decode(value))
        if revision:
            return revision
        rows = self._client.xrevrange(self.stream_key, max="+", min="-", count=1)
        events = self._rows_to_events(rows)
        return max((event["revision"] for event in events), default=0)

    def status(self) -> Dict[str, Any]:
        out: Dict[str, Any] = {
            "store": "redis",
            "redis_connected": False,
            "stream_key": self.stream_key,
            "latest_revision": 0,
            "stream_len": None,
            "oldest_revision": None,
            "subscriber_connected": bool(
                self._subscriber_thread and self._subscriber_thread.is_alive()
            ),
        }
        try:
            ping = getattr(self._client, "ping", None)
            if callable(ping):
                ping()
            out["redis_connected"] = True
            out["latest_revision"] = self.latest_revision()
            xlen = getattr(self._client, "xlen", None)
            if callable(xlen):
                out["stream_len"] = int(xlen(self.stream_key))
            out["oldest_revision"] = self._oldest_revision()
        except Exception as exc:
            out["error"] = str(exc)
        return out

    def replay_events(
        self,
        *,
        cities: Optional[Set[str]],
        since_revision: int,
        limit: int,
    ) -> List[Dict[str, Any]]:
        city_set = _normalize_city_set(cities)
        since = max(0, int(since_revision or 0))
        bounded_limit = max(1, min(MAX_REPLAY_LIMIT, int(limit or 1)))
        if since <= 0:
            return self._replay_events_forward(
                city_set=city_set,
                since=since,
                limit=bounded_limit,
            )[0]
        return self._replay_events_reverse(
            city_set=city_set,
            since=since,
            limit=bounded_limit,
        )[0]

    def replay_requires_resync(
        self,
        *,
        cities: Optional[Set[str]],
        since_revision: int,
        replay_count: int,
        limit: int,
    ) -> bool:
        city_set = _normalize_city_set(cities)
        since = max(0, int(since_revision or 0))
        oldest_revision = self._oldest_revision()
        if since > 0 and oldest_revision and since < oldest_revision - 1:
            return True
        bounded_limit = max(1, int(limit or 1))
        if int(replay_count or 0) < bounded_limit:
            return False
        probe_limit = min(MAX_REPLAY_LIMIT + 1, bounded_limit + 1)
        if since <= 0:
            probe_events, _, _ = self._replay_events_forward(
                city_set=city_set,
                since=since,
                limit=probe_limit,
            )
            return len(probe_events) > bounded_limit

        probe_events, hit_boundary, scanned = self._replay_events_reverse(
            city_set=city_set,
            since=since,
            limit=probe_limit,
        )
        if len(probe_events) > bounded_limit:
            return True
        return not hit_boundary and scanned >= self.replay_scan_max

    def start_live_subscription(self, callback: Callable[[Dict[str, Any]], None]) -> None:
        with self._subscriber_lock:
            if self._subscriber_thread and self._subscriber_thread.is_alive():
                return
            self._subscriber_stop = threading.Event()
            self._subscriber_thread = threading.Thread(
                target=self._live_subscription_loop,
                args=(callback, self._subscriber_stop),
                name="polyweather-redis-realtime-subscriber",
                daemon=True,
            )
            self._subscriber_thread.start()

    def stop_live_subscription(self) -> None:
        with self._subscriber_lock:
            if self._subscriber_stop:
                self._subscriber_stop.set()
            self._subscriber_thread = None
            self._subscriber_stop = None

    def _live_subscription_loop(
        self,
        callback: Callable[[Dict[str, Any]], None],
        stop_event: threading.Event,
    ) -> None:
        last_seen_id = "$"
        while not stop_event.is_set():
            try:
                rows = self._client.xread({self.stream_key: last_seen_id}, count=100, block=5000)
                for _stream_name, entries in rows or []:
                    for entry_id, fields in entries:
                        last_seen_id = _decode(entry_id)
                        callback(self._entry_to_event(entry_id, fields))
            except Exception as exc:
                logger.warning(f"Redis realtime subscriber disconnected: {exc}")
                stop_event.wait(2.0)

    def _all_events(self) -> List[Dict[str, Any]]:
        rows = self._client.xrange(self.stream_key, min="-", max="+")
        return self._rows_to_events(rows)

    def _oldest_revision(self) -> Optional[int]:
        rows = self._client.xrange(self.stream_key, min="-", max="+", count=1)
        events = self._rows_to_events(rows)
        if not events:
            return None
        return int(events[0].get("revision") or 0) or None

    def _rows_to_events(self, rows: Any) -> List[Dict[str, Any]]:
        return [self._entry_to_event(entry_id, fields) for entry_id, fields in rows or []]

    @staticmethod
    def _matches_city(event: Dict[str, Any], city_set: Set[str]) -> bool:
        return not city_set or str(event.get("city") or "").strip().lower() in city_set

    def _replay_events_forward(
        self,
        *,
        city_set: Set[str],
        since: int,
        limit: int,
    ) -> tuple[List[Dict[str, Any]], bool, int]:
        replay: List[Dict[str, Any]] = []
        scanned = 0
        min_id = "-"
        hit_boundary = False
        while len(replay) < limit and scanned < self.replay_scan_max:
            count = min(self.replay_chunk_size, self.replay_scan_max - scanned)
            rows = self._client.xrange(self.stream_key, min=min_id, max="+", count=count)
            if not rows:
                hit_boundary = True
                break
            scanned += len(rows)
            for entry_id, fields in rows:
                event = self._entry_to_event(entry_id, fields)
                if int(event.get("revision") or 0) <= since:
                    continue
                if self._matches_city(event, city_set):
                    replay.append(event)
                    if len(replay) >= limit:
                        break
            min_id = f"({_decode(rows[-1][0])}"
        return replay, hit_boundary, scanned

    def _replay_events_reverse(
        self,
        *,
        city_set: Set[str],
        since: int,
        limit: int,
    ) -> tuple[List[Dict[str, Any]], bool, int]:
        replay_desc: List[Dict[str, Any]] = []
        scanned = 0
        max_id = "+"
        hit_boundary = False
        while len(replay_desc) < limit and scanned < self.replay_scan_max:
            count = min(self.replay_chunk_size, self.replay_scan_max - scanned)
            rows = self._client.xrevrange(self.stream_key, max=max_id, min="-", count=count)
            if not rows:
                hit_boundary = True
                break
            scanned += len(rows)
            for entry_id, fields in rows:
                event = self._entry_to_event(entry_id, fields)
                if int(event.get("revision") or 0) <= since:
                    hit_boundary = True
                    return list(reversed(replay_desc)), hit_boundary, scanned
                if self._matches_city(event, city_set):
                    replay_desc.append(event)
                    if len(replay_desc) >= limit:
                        break
            max_id = f"({_decode(rows[-1][0])}"
        return list(reversed(replay_desc)), hit_boundary, scanned

    @staticmethod
    def _entry_to_event(entry_id: Any, fields: Dict[Any, Any]) -> Dict[str, Any]:
        normalized = {_decode(key): _decode(value) for key, value in dict(fields or {}).items()}
        payload = json.loads(normalized.get("payload_json") or "{}")
        schema_type = normalized.get("schema_type") or "city_observation_patch"
        schema_version = int(normalized.get("schema_version") or 1)
        created_at_ms = _int_or_zero(normalized.get("created_at_ms")) or int(time.time() * 1000)
        ts = _int_or_zero(normalized.get("ts")) or created_at_ms
        obs_time = normalized.get("obs_time") or None
        return {
            "type": normalized.get("type") or f"{schema_type}.v{schema_version}",
            "revision": int(normalized["revision"]),
            "city": normalized.get("city") or "",
            "source": normalized.get("source") or "",
            "obs_time": obs_time,
            **_time_contract_from_payload(payload if isinstance(payload, dict) else {}),
            "ts": ts,
            "payload": payload if isinstance(payload, dict) else {},
        }