data_store.py

#!/usr/bin/env python3
"""Runtime coordinator: Postgres handle, MQTT→Discord event queue, periodic enrichment."""

import asyncio
from datetime import datetime, timedelta
import logging
from zoneinfo import ZoneInfo
import aiohttp
from fastapi.encoders import jsonable_encoder

from broadcaster import Broadcaster
from models.node import Node
from storage.db.postgres import PostgresStorage
import utils

logger = logging.getLogger(__name__)

# Named provider presets (none shipped — configure explicitly).
_PROVIDER_PRESETS: dict = {}

# Dead legacy names dropped with a warning so stale configs still start.
_DEAD_LEGACY_PROVIDERS: dict = {
    "world.meshinfo.network": "upstream service is gone — use a Meshview entry",
    "bayme": "data.bayme.sh was never valid — use kind=\"meshview\" pointed at "
             "https://meshview.bayme.sh/api/nodes",
}

_PER_REQUEST_DELAY_SEC = 0.1
_REQUEST_TIMEOUT_SEC = 10.0
# Safety cap on a Meshview bulk /api/nodes body so one huge response can't blow up memory.
_MESHVIEW_BULK_MAX_BYTES = 50 * 1024 * 1024
# Stop hammering a provider after N consecutive failures in one cycle; resets next cycle.
_PROVIDER_FAIL_THRESHOLD = 10


def _resolve_providers(config) -> list:
  """Resolve config.server.enrich.providers (or legacy 'provider') into a runtime list.

  Accepts dicts (Meshview bulk or MeshInfo URL template) and bare URL templates.
  Dead legacy names are dropped with a warning.
  """
  cfg = config.get('server', {}).get('enrich', {}) or {}
  raw = cfg.get('providers')
  if raw is None and cfg.get('provider') is not None:
    raw = [cfg['provider']]
  if not raw:
    return []
  result = []
  for entry in raw:
    if isinstance(entry, str):
      if entry in _DEAD_LEGACY_PROVIDERS:
        logger.warning(
          "server.enrich provider %r is no longer reachable; %s",
          entry, _DEAD_LEGACY_PROVIDERS[entry],
        )
        continue
      preset = _PROVIDER_PRESETS.get(entry)
      if preset is None:
        # Bare URLs with an {ids} placeholder are generic templates.
        if entry.startswith(('http://', 'https://')) and '{ids}' in entry:
          result.append({
            "name": entry,
            "url_template": entry,
            "single_id_per_request": False,
          })
        else:
          logger.warning("Unknown enrichment provider %r; skipping", entry)
      else:
        result.append({"name": entry, **preset})
    elif isinstance(entry, dict) and entry.get("url"):
      kind = entry.get("kind", "meshinfo")
      url = entry["url"]
      if kind == "meshview":
        # Bulk fetch + local filter; no {ids} placeholder needed.
        try:
          days_active = int(entry.get("days_active", 7))
        except (TypeError, ValueError):
          logger.warning("Enrichment provider %r has invalid days_active %r; using 7",
                         entry.get("name", url), entry.get("days_active"))
          days_active = 7
        result.append({
          "name": entry.get("name", url),
          "kind": "meshview",
          "url": url,
          "days_active": days_active,
        })
      elif kind == "meshinfo":
        if "{ids}" not in url:
          logger.warning("Enrichment provider %r is missing {ids} placeholder; skipping", entry.get("name", url))
          continue
        result.append({
          "name": entry.get("name", url),
          "url_template": url,
          "single_id_per_request": bool(entry.get("single_id_per_request", False)),
        })
      else:
        logger.warning("Unknown enrichment provider kind %r; skipping", kind)
    else:
      logger.warning("Invalid enrichment provider entry %r; skipping", entry)
  return result


class DataStore:
  def __init__(self, config):
    self.config = config
    self.mqtt_connect_time: datetime = self.config['server']['start_time']
    # Bounded so MQTT producers shed load if the Discord consumer stalls.
    self.discord_event_queue: asyncio.Queue = asyncio.Queue(maxsize=1000)
    # Live SSE fan-out (node + chat). Lives here (not on MQTT) so it survives
    # MQTT supervise-restarts and is reachable from both ingest and the API.
    self.broadcaster = Broadcaster()
    self.pg_storage = PostgresStorage(config)

  def update(self, key, value):
    self.__dict__[key] = value

  async def update_node(self, id: str, node):
    """Apply geocoding + freshness fields, persist to Postgres, refresh the cache."""
    n = node.copy()
    # Isolate position from the caller's dict — geocoding mutates pos in place
    # and we'd otherwise leak enriched fields into the raw mqtt_messages log.
    if n.get('position') is None:
      n['position'] = {}
    else:
      n['position'] = dict(n['position'])

    # Geocode positions at most once an hour per node.
    if self.config['integrations']['geocoding']['enabled']:
      pos = n['position']
      if 'geocoded' not in pos:
        pos['geocoded'] = None
      lat_i = pos.get('latitude_i')
      lon_i = pos.get('longitude_i')
      last_geo = pos.get('last_geocoding')
      # DB rows arrive as ISO strings; normalize for comparison.
      if isinstance(last_geo, str):
        try:
          last_geo = datetime.fromisoformat(last_geo).astimezone(ZoneInfo(self.config['server']['timezone']))
        except ValueError:
          last_geo = None
      if lat_i is not None and lon_i is not None:
        if pos['geocoded'] is None or last_geo is None or last_geo < datetime.now().astimezone(ZoneInfo(self.config['server']['timezone'])) - timedelta(minutes=60):
          try:
            # geocode_position is sync (requests.get); offload so it doesn't block the event loop.
            geocoded = await asyncio.to_thread(
              utils.geocode_position,
              self.config['integrations']['geocoding']['geocode.maps.co']['api_key'],
              lat_i / 10000000,
              lon_i / 10000000,
            )
            if geocoded is not None:
              pos['geocoded'] = geocoded
              pos['last_geocoding'] = datetime.now().astimezone(ZoneInfo(self.config['server']['timezone']))
          except Exception as e:
            logger.warning("Failed to geocode position: %s", e)

    # Any packet reactivates a pruned node.
    n['active'] = True
    if isinstance(n.get('last_seen'), str):
      n['last_seen'] = datetime.fromisoformat(n['last_seen']).astimezone(ZoneInfo(self.config['server']['timezone']))
    now_local = datetime.now().astimezone(ZoneInfo(self.config['server']['timezone']))
    if isinstance(n.get('last_seen'), datetime):
      n['since'] = now_local - n['last_seen']
    n['last_seen'] = now_local

    if self.pg_storage is not None and getattr(self.pg_storage, "enabled", False) and getattr(self.pg_storage, "pool", None) is not None:
      try:
        await self.pg_storage.write_node(id, n)
        # Refresh cache only after a confirmed write.
        self.pg_storage.cache_node_set(id, n)
        # Live push: only after a confirmed persist, and only when someone is
        # listening. jsonable_encoder snapshots `n` (datetimes→ISO, dicts
        # copied) so a later handler mutating the live cache ref can't race the
        # serialization. Matches the per-node shape /v1/nodes emits.
        if self.broadcaster.subscriber_count:
          try:
            self.broadcaster.publish("node", jsonable_encoder(n))
          except Exception as e:
            logger.debug("node broadcast failed for %s: %s", id, e)
      except Exception as e:
        logger.error("Failed to write node %s to Postgres: %s", id, e)

  async def load(self):
    logger.info("Loading data from PostgreSQL")
    await self._load_from_postgres()

  async def _load_from_postgres(self):
    """Open the pool, run migrations, and seed the local + broadcast node rows."""
    try:
      ok = await self.pg_storage.connect()
      if not ok:
        raise RuntimeError("PostgreSQL connect failed")

      await self.pg_storage.ensure_schema()

      default_id = self.config['server']['node_id']
      default_node = Node.default_node(default_id)
      await self.pg_storage.write_node(default_id, default_node)
      self.pg_storage.cache_node_set(default_id, default_node)

      broadcast_node = Node.default_node('ffffffff')
      await self.pg_storage.write_node('ffffffff', broadcast_node)
      self.pg_storage.cache_node_set('ffffffff', broadcast_node)

      logger.info("PostgreSQL mode: data is queried directly from the database")

    except Exception as e:
      logger.exception("Failed to initialize PostgreSQL connection: %s", e)
      raise

  ### helpers

  async def backfill_node_infos(self):
    """Query every configured enrichment provider for the current set of unknown nodes."""
    ids_needing_enrichment = await self.pg_storage.find_nodes_needing_enrichment()
    if not ids_needing_enrichment:
      return
    logger.info("Nodes needing enrichment: %d", len(ids_needing_enrichment))
    await self.enrich_nodes(ids_needing_enrichment)

  async def enrich_nodes(self, node_ids_to_enrich):
    """Iterate providers in order, asking each for the still-unknown ids."""
    providers = _resolve_providers(self.config)
    if not providers:
      logger.debug("No enrichment providers configured")
      return
    pending: set = set(node_ids_to_enrich)
    timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT_SEC)
    async with aiohttp.ClientSession(timeout=timeout) as session:
      for prov in providers:
        if not pending:
          break
        before_count = len(pending)
        if prov.get("kind") == "meshview":
          result = await self._enrich_via_meshview_bulk(session, prov, pending)
        else:
          result = await self._enrich_via_provider(session, prov, sorted(pending))
        # One log line per provider per cycle. WARNING when the provider looks
        # unreachable (0/N succeeded or circuit-breaker fired), INFO otherwise.
        attempted = result["attempted"]
        succeeded = result["succeeded"]
        named = result["named"]
        aborted = result["aborted"]
        if attempted > 0 and succeeded == 0:
          logger.warning(
            "%s: 0/%d requests succeeded — upstream may be unreachable",
            prov["name"], attempted,
          )
        elif aborted:
          logger.warning(
            "%s: aborted after %d consecutive failures; named %d of %d (%d/%d requests sent succeeded)",
            prov["name"], _PROVIDER_FAIL_THRESHOLD, len(named), before_count, succeeded, attempted,
          )
        elif succeeded < attempted:
          logger.info(
            "%s: named %d of %d node(s); %d/%d requests succeeded",
            prov["name"], len(named), before_count, succeeded, attempted,
          )
        elif attempted > 0:
          logger.info(
            "%s: named %d of %d node(s)",
            prov["name"], len(named), before_count,
          )
        pending -= named
    if pending:
      logger.debug("After all providers, %d node(s) remain unnamed", len(pending))

  async def _enrich_via_provider(self, session, prov: dict, node_ids: list) -> dict:
    """Query one provider for the given ids.

    Returns {named, attempted, succeeded, aborted}. `aborted` is set when the
    circuit breaker trips so the caller can distinguish "unreachable" from
    "responded but didn't know any of our nodes".
    """
    named: set = set()
    attempted = 0
    succeeded = 0
    consecutive_failures = 0
    aborted = False
    if prov["single_id_per_request"]:
      # Provider rejects multi-id requests; iterate, paced so we don't hammer it.
      for node_id in node_ids:
        attempted += 1
        data = await self._fetch_provider(session, prov, [node_id])
        if data is not None:
          succeeded += 1
          consecutive_failures = 0
          for nid, info in data.items():
            if await self._apply_enrichment(nid, info):
              named.add(nid)
        else:
          consecutive_failures += 1
          if consecutive_failures >= _PROVIDER_FAIL_THRESHOLD:
            aborted = True
            break
        await asyncio.sleep(_PER_REQUEST_DELAY_SEC)
    else:
      # Generic MeshInfo template — comma-join into one call.
      attempted = 1
      data = await self._fetch_provider(session, prov, node_ids)
      if data is not None:
        succeeded = 1
        for nid, info in data.items():
          if await self._apply_enrichment(nid, info):
            named.add(nid)
    return {"named": named, "attempted": attempted, "succeeded": succeeded, "aborted": aborted}

  async def _fetch_provider(self, session, prov: dict, ids: list) -> dict | None:
    """One GET against `prov`. Returns the parsed JSON dict, or None on any failure."""
    # str.replace rather than .format — operator-supplied URLs may have stray
    # `{...}` braces that would trip str.format's placeholder parser.
    url = prov["url_template"].replace("{ids}", ",".join(ids))
    try:
      async with session.get(url) as response:
        if response.status == 200:
          return await response.json()
        logger.debug("%s: HTTP %d for %d id(s)", prov["name"], response.status, len(ids))
    except Exception as e:
      logger.debug("%s: request failed for %d id(s): %s", prov["name"], len(ids), e)
    return None

  async def _enrich_via_meshview_bulk(self, session, prov: dict, pending: set) -> dict:
    """One bulk GET against a Meshview /api/nodes; filter locally.
    Cheaper than per-id requests when many unknowns are covered by one response."""
    named: set = set()
    sep = "&" if "?" in prov["url"] else "?"
    url = f"{prov['url']}{sep}days_active={prov['days_active']}"
    try:
      async with session.get(url) as response:
        if response.status != 200:
          logger.debug("%s: HTTP %d", prov["name"], response.status)
          return {"named": named, "attempted": 1, "succeeded": 0, "aborted": False}
        # Bound memory: skip a pathologically large bulk body rather than load it.
        clen = getattr(response, "content_length", None)
        if clen is not None and clen > _MESHVIEW_BULK_MAX_BYTES:
          logger.warning("%s: bulk response too large (%d bytes); skipping", prov["name"], clen)
          return {"named": named, "attempted": 1, "succeeded": 0, "aborted": False}
        data = await response.json()
    except Exception as e:
      logger.debug("%s: bulk fetch failed: %s", prov["name"], e)
      return {"named": named, "attempted": 1, "succeeded": 0, "aborted": False}
    # Meshview gives node_id as uint32 decimal + snake_case names; adapt to our shape.
    lookup: dict = {}
    for node in (data.get("nodes") or []):
      nid = utils.normalize_node_id(node.get("node_id"))
      if nid is None:
        continue
      long_ = (node.get("long_name") or "").strip()
      short = (node.get("short_name") or "").strip()
      if long_ or short:
        lookup[nid] = {"longName": long_, "shortName": short}
    for nid in list(pending):
      info = lookup.get(nid)
      if info is None:
        continue
      if await self._apply_enrichment(nid, info):
        named.add(nid)
    return {"named": named, "attempted": 1, "succeeded": 1, "aborted": False}

  async def _apply_enrichment(self, node_id: str, info: dict) -> bool:
    """Write enriched name fields direct to Postgres + cache; returns True if applied.
    Bypasses update_node so a name lookup doesn't bump last_seen or active."""
    short = info.get('shortName')
    long_ = info.get('longName')
    if not short and not long_:
      return False
    node = await self.pg_storage.get_node_cached(node_id)
    if node is None:
      return False
    if short:
      node['shortname'] = short
    if long_:
      node['longname'] = long_
    try:
      await self.pg_storage.write_node(node_id, node)
      self.pg_storage.cache_node_set(node_id, node)
    except Exception as e:
      logger.error("Failed to write enrichment for %s: %s", node_id, e)
      return False
    return True