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:
    # Can't find a way to resolve the circular dependency on types but apparently this is a way of managing that
    # while still having type safety in method definitions.
    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. These are all created by setup(), which must be called before use; they are None only in the window
        # between construction of this (global, single-instance) object and that call.
        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. Unlike the caches above, these genuinely may never be populated,
        # if the corresponding static data provider isn't configured, so callers must handle None.
        self.cq_zone_data: geopandas.GeoDataFrame | None = None
        self.itu_zone_data: geopandas.GeoDataFrame | None = None

    @property
    def alerts(self) -> LiveDataCache[Alert]:
        assert self._alerts is not None, "DataStore.setup() must be called before use"
        return self._alerts

    @property
    def spots(self) -> LiveDataCache[Spot]:
        assert self._spots is not None, "DataStore.setup() must be called before use"
        return self._spots

    @property
    def callsign_data_countryfiles(self) -> diskcache.Cache:
        assert self._callsign_data_countryfiles is not None, "DataStore.setup() must be called before use"
        return self._callsign_data_countryfiles

    @property
    def callsign_data_clublogxml(self) -> diskcache.Cache:
        assert self._callsign_data_clublogxml is not None, "DataStore.setup() must be called before use"
        return self._callsign_data_clublogxml

    @property
    def callsign_data_clublogapi(self) -> diskcache.Cache:
        assert self._callsign_data_clublogapi is not None, "DataStore.setup() must be called before use"
        return self._callsign_data_clublogapi

    @property
    def callsign_data_qrz(self) -> diskcache.Cache:
        assert self._callsign_data_qrz is not None, "DataStore.setup() must be called before use"
        return self._callsign_data_qrz

    @property
    def callsign_data_hamqth(self) -> diskcache.Cache:
        assert self._callsign_data_hamqth is not None, "DataStore.setup() must be called before use"
        return self._callsign_data_hamqth

    @property
    def dxcc_data(self) -> diskcache.Cache:
        assert self._dxcc_data is not None, "DataStore.setup() must be called before use"
        return self._dxcc_data

    @property
    def activity_refs(self) -> diskcache.Cache:
        assert self._activity_refs is not None, "DataStore.setup() must be called before use"
        return self._activity_refs

    @property
    def status(self) -> SingleObjectDataCache[dict[str, Any]]:
        assert self._status is not None, "DataStore.setup() must be called before use"
        return self._status

    @property
    def solar_conditions(self) -> SingleObjectDataCache[SolarConditions]:
        assert self._solar_conditions is not None, "DataStore.setup() must be called before use"
        return self._solar_conditions

    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()