spothole/core/data_store.py

from __future__ import annotations

import logging
import re
from pathlib import Path
from typing import TYPE_CHECKING, Any

import diskcache
import geopandas

from core.config import MAX_ALERT_AGE, MAX_SPOT_AGE
from core.live_data_cache import LiveDataCache
from core.single_object_data_cache import SingleObjectDataCache
from data.solar_conditions import SolarConditions

if TYPE_CHECKING:
    # Deferred to avoid a circular import: data.alert and data.spot both import core.data_store at module level.
    from data.alert import Alert
    from data.spot import Spot

logger = logging.getLogger(__name__)

CACHE_DIR = "./cache/"


class DataStore:
    """Data caching/storage object. Handles storage of spots, alerts, solar conditions, activity reference data, and
    callsign lookup data using different caching strategies for each."""

    def __init__(self) -> None:
        # Constants
        self._MAX_SPOT_COUNT = 100000
        self._MAX_ALERT_COUNT = 100000
        self._SPOT_ALERT_SNAPSHOT_INTERVAL_SEC = 300
        self.CALLSIGN_DATA_TTL_SEC = 30 * 24 * 60 * 60
        # Caches
        self.alerts: LiveDataCache[Alert] | None = None
        self.spots: LiveDataCache[Spot] | None = None
        self.callsign_data_countryfiles: diskcache.Cache | None = None
        self.callsign_data_clublogxml: diskcache.Cache | None = None
        self.callsign_data_clublogapi: diskcache.Cache | None = None
        self.callsign_data_qrz: diskcache.Cache | None = None
        self.callsign_data_hamqth: diskcache.Cache | None = None
        self.dxcc_data: diskcache.Cache | None = None
        self.dxcc_lookup_by_call_regex: list[tuple[re.Pattern[str], Any]] = []
        self.activity_refs: diskcache.Cache | None = None
        self.status: SingleObjectDataCache[dict[str, Any]] | None = None
        self.solar_conditions: SingleObjectDataCache[SolarConditions] | None = None
        # ITU/CQ zone GeoJSON data is only ever loaded statically from a local file so these don't even need to be
        # caches, they can just be straight objects
        self.cq_zone_data: geopandas.GeoDataFrame | None = None
        self.itu_zone_data: geopandas.GeoDataFrame | None = None

    def setup(self) -> None:
        Path(CACHE_DIR).mkdir(parents=True, exist_ok=True)

        # For solar data and status data, we use a wrapper around disk cache where each cache contains only a single
        # object exposed to the wider application, and provides a store() method for callers to notify diskcache that
        # the object has changed and needs to be re-cached.
        self.solar_conditions = SingleObjectDataCache(f"{CACHE_DIR}solar", SolarConditions())
        self.status = SingleObjectDataCache(f"{CACHE_DIR}status", {})

        # Standard disk cache for static reference and activity ref data. Separate provider threads will repopulate
        # these on a regular basis but there's no need for a TTL since old data is better than no data.
        self.dxcc_data = diskcache.Cache(f"{CACHE_DIR}dxcc_data")
        self.regenerate_call_regex_to_dxcc_entity_map()

        # For activity reference data specifically, we need to key on both activity *and* reference, and trying to do
        # two layers of dict in diskcache absolutely destroys performance with unpickling huge dicts, so we have an
        # ugly "activity:ref" syntax for keys to keep it a single level.
        self.activity_refs = diskcache.Cache(f"{CACHE_DIR}activity_refs")
        logger.info(f"Loaded data for {len(self.activity_refs)} activity references.")

        # Standard disk cache for callsign data. This data does have a TTL to trigger an occasional re-lookup.
        # Old data *is* better than no data, but we can't have a background thread re-looking-up every callsign
        # we've seen, so we rely on them timing out and this triggering another lookup.
        self.callsign_data_countryfiles = diskcache.Cache(f"{CACHE_DIR}callsign_data_countryfiles")
        self.callsign_data_clublogxml = diskcache.Cache(f"{CACHE_DIR}callsign_data_clublogxml")
        self.callsign_data_clublogapi = diskcache.Cache(f"{CACHE_DIR}callsign_data_clublogapi")
        self.callsign_data_qrz = diskcache.Cache(f"{CACHE_DIR}callsign_data_qrz")
        self.callsign_data_hamqth = diskcache.Cache(f"{CACHE_DIR}callsign_data_hamqth")
        unique_keys = set()
        for c in [
            self.callsign_data_countryfiles,
            self.callsign_data_clublogxml,
            self.callsign_data_clublogapi,
            self.callsign_data_qrz,
            self.callsign_data_hamqth,
        ]:
            unique_keys.update(c)
        logger.info(f"Loaded data for {len(unique_keys)} callsigns.")

        # Special caches for spots and alerts, which have TTL and write snapshots to disk at an interval. We
        # specifically load these caches *last* so that any activity ref and callsign data is already loaded from disk
        # cache before the spots and alerts are live in the system.
        self.spots = LiveDataCache(
            maxsize=self._MAX_SPOT_COUNT,
            ttl=MAX_SPOT_AGE,
            snapshot_dir=f"{CACHE_DIR}spots",
            snapshot_interval_sec=self._SPOT_ALERT_SNAPSHOT_INTERVAL_SEC,
        )
        logger.info(f"Loaded {len(self.spots.keys())} spots from a previous run.")

        self.alerts = LiveDataCache(
            maxsize=self._MAX_ALERT_COUNT,
            ttl=MAX_ALERT_AGE,
            snapshot_dir=f"{CACHE_DIR}alerts",
            snapshot_interval_sec=self._SPOT_ALERT_SNAPSHOT_INTERVAL_SEC,
        )
        logger.info(f"Loaded {len(self.alerts.keys())} alerts from a previous run.")

    def regenerate_call_regex_to_dxcc_entity_map(self) -> None:
        """DXCC entity data from K0SWE includes a regex which we can use to match a callsign, and determine which DXCC
        entity it belongs to. But getting every DXCC entity data object out of DiskCache, iterating, compiling its regex
        and testing the callsign every time is expensive. So instead we build a separate in-memory lookup of compiled
        regex against DXCC entity code, as a list of tuples we can iterate through."""

        self.dxcc_lookup_by_call_regex = []
        for entry in [DATA_STORE.dxcc_data[key] for key in DATA_STORE.dxcc_data]:
            self.dxcc_lookup_by_call_regex.append((re.compile(entry["prefixRegex"]), entry["entityCode"]))

    def close(self) -> None:
        self.spots.close()
        self.alerts.close()
        self.solar_conditions.close()
        self.status.close()
        self.dxcc_data.close()
        self.activity_refs.close()
        self.callsign_data_countryfiles.close()
        self.callsign_data_clublogxml.close()
        self.callsign_data_clublogapi.close()
        self.callsign_data_qrz.close()
        self.callsign_data_hamqth.close()


# Global object
DATA_STORE = DataStore()