prop/lib/microwaveprop/weather/grid_cache.ex

defmodule Microwaveprop.Weather.GridCache do
  @moduledoc """
  Cross-pod cache of derived HRRR grid rows keyed by `valid_time`. When
  `Microwaveprop.Valkey` is configured, the chunked map for each
  valid_time is held in Valkey (one key per 5°×5° spatial chunk plus a
  ZSET index of cached valid_times). Otherwise the cache stays
  pod-local in ETS — preserved so dev / test work without Valkey.

  Each valid_time worth of CONUS HRRR is ~92k cells × 22 fields ≈
  32 MiB compressed. Holding that in pod-local ETS meant N replicas
  paid the full memory cost per pod. With the Valkey backend a single
  cluster-wide copy serves every reader.

  ## Storage layout (Valkey)

    * `prop:wg:vts` — ZSET of cached valid_time iso8601 strings, score
      = unix epoch seconds. `latest_valid_time/0` does ZREVRANGE 0 0.
    * `prop:wg:<vt_iso>:<lat_band>:<lon_band>` — string key per chunk,
      value is `:erlang.term_to_binary(%{{lat, lon} => row, ...})`.
      `fetch_bounds/2` MGETs only the chunks intersecting the viewport.

  Per-key TTL is `@ttl_seconds`, so cache eviction is automatic — no
  prune worker needed when running on Valkey. The ETS path keeps the
  legacy `prune_keep_latest`/`prune_older_than` helpers since they're
  still called by `Microwaveprop.Weather.fill_grid_cache_from_scalar/1`
  on every fill.

  ## Fill coordination

  Cache fills run `ScalarFile.read_bounds/2` which gunzips and decodes
  per-chunk files (~15 s for a full CONUS valid_time). The
  `:weather_grid_fill_locks` ETS table (kept regardless of backend)
  stops N concurrent /weather mounts on a single pod from each firing
  the slow read. Multi-pod duplicate fills are tolerated: the second
  pod's writes overwrite identical bytes in Valkey, no data corruption.
  """
  use GenServer

  alias Microwaveprop.Valkey
  alias Phoenix.PubSub

  require Logger

  @table :weather_grid_cache
  @lock_table :weather_grid_fill_locks
  @topic "weather:cache"
  @pubsub Microwaveprop.PubSub
  @chunk_step 5

  # Per-chunk TTL in Valkey. 8 hours covers analysis + 7 forecast hours
  # (matches the ETS valid_time cap). Beyond 8 h, fall back to disk.
  @ttl_seconds 8 * 3600

  @valkey_index_key "prop:wg:vts"
  @valkey_chunk_prefix "prop:wg:"

  @type row :: %{required(:lat) => float(), required(:lon) => float(), optional(atom()) => any()}
  @type bounds :: %{optional(String.t()) => float()}

  @spec start_link(keyword()) :: GenServer.on_start()
  def start_link(opts), do: GenServer.start_link(__MODULE__, opts, name: __MODULE__)

  # ---------- Public API ----------

  @spec fetch(DateTime.t()) :: {:ok, [row()]} | :miss
  def fetch(valid_time) do
    if valkey_backend?() do
      valkey_fetch(valid_time)
    else
      ets_fetch(valid_time)
    end
  end

  @spec fetch_bounds(DateTime.t(), bounds() | nil) :: {:ok, [row()]} | :miss
  def fetch_bounds(valid_time, bounds) do
    if valkey_backend?() do
      valkey_fetch_bounds(valid_time, bounds)
    else
      ets_fetch_bounds(valid_time, bounds)
    end
  end

  @spec fetch_point(DateTime.t(), float(), float()) :: {:ok, row()} | :miss
  def fetch_point(valid_time, lat, lon) do
    if valkey_backend?() do
      valkey_fetch_point(valid_time, lat, lon)
    else
      ets_fetch_point(valid_time, lat, lon)
    end
  end

  @spec put(DateTime.t(), [row()]) :: :ok
  def put(valid_time, rows) do
    chunked = list_to_chunks(rows)

    if valkey_backend?() do
      valkey_put(valid_time, chunked)
    else
      :ets.insert(@table, {valid_time, chunked})
      :ok
    end
  end

  @doc """
  Insert (Valkey path) or insert + broadcast (ETS path). Kept as a single
  public name so legacy callers don't have to branch.
  """
  @spec broadcast_put(DateTime.t(), [row()]) :: :ok
  def broadcast_put(valid_time, rows) do
    if valkey_backend?() do
      # Valkey is the single source of truth — peers will read the same
      # data on their next fetch. Broadcasting becomes redundant.
      put(valid_time, rows)
    else
      _ = PubSub.broadcast(@pubsub, @topic, {:weather_cache_refresh, valid_time, rows})
      :ok
    end
  end

  @spec latest_valid_time() :: DateTime.t() | nil
  def latest_valid_time do
    if valkey_backend?() do
      valkey_latest_valid_time()
    else
      ets_latest_valid_time()
    end
  end

  @spec prune_older_than(DateTime.t()) :: non_neg_integer()
  def prune_older_than(cutoff) do
    if valkey_backend?() do
      valkey_prune_older_than(cutoff)
    else
      match_spec = [{{:"$1", :_}, [{:<, :"$1", {:const, cutoff}}], [true]}]
      :ets.select_delete(@table, match_spec)
    end
  end

  @doc """
  Keep only the `keep` most-recent valid_times, dropping the rest. On
  Valkey this is a no-op because per-key TTL handles eviction; the ETS
  path is unchanged.
  """
  @spec prune_keep_latest(non_neg_integer()) :: non_neg_integer()
  def prune_keep_latest(keep) when is_integer(keep) and keep >= 0 do
    if valkey_backend?() do
      0
    else
      ets_prune_keep_latest(keep)
    end
  end

  @spec clear() :: :ok
  def clear, do: GenServer.call(__MODULE__, :clear)

  @spec claim_fill(term()) :: boolean()
  def claim_fill(valid_time) do
    GenServer.call(__MODULE__, {:claim_fill, valid_time, self()})
  end

  @spec release_fill(term()) :: :ok
  def release_fill(valid_time) do
    GenServer.call(__MODULE__, {:release_fill, valid_time})
  end

  @spec sync() :: :ok
  def sync, do: GenServer.call(__MODULE__, :sync)

  # ---------- GenServer ----------

  @impl true
  def init(_opts) do
    # ETS table is created unconditionally so the legacy backend (and
    # tests that bypass Valkey) work without separate setup. The lock
    # table is also unconditional — fill coordination still happens
    # per-pod regardless of where the data lives.
    _ = :ets.new(@table, [:set, :named_table, :public, :compressed, read_concurrency: true])
    _ = :ets.new(@lock_table, [:set, :named_table, :protected])
    :ok = PubSub.subscribe(@pubsub, @topic)
    {:ok, %{monitors: %{}}}
  end

  @impl true
  def handle_call(:sync, _from, state), do: {:reply, :ok, state}

  def handle_call(:clear, _from, state) do
    if valkey_backend?(), do: valkey_clear(), else: :ets.delete_all_objects(@table)

    for {ref, _vt} <- state.monitors, do: Process.demonitor(ref, [:flush])
    :ets.delete_all_objects(@lock_table)

    {:reply, :ok, %{state | monitors: %{}}}
  end

  def handle_call({:claim_fill, valid_time, caller}, _from, state) do
    if :ets.insert_new(@lock_table, {valid_time, caller}) do
      ref = Process.monitor(caller)
      {:reply, true, %{state | monitors: Map.put(state.monitors, ref, valid_time)}}
    else
      {:reply, false, state}
    end
  end

  def handle_call({:release_fill, valid_time}, _from, state) do
    {monitors, _matched} = pop_monitor_for(state.monitors, valid_time)
    :ets.delete(@lock_table, valid_time)
    {:reply, :ok, %{state | monitors: monitors}}
  end

  @impl true
  def handle_info({:DOWN, ref, :process, _pid, _reason}, state) do
    case Map.pop(state.monitors, ref) do
      {nil, _} ->
        {:noreply, state}

      {valid_time, monitors} ->
        :ets.delete(@lock_table, valid_time)
        {:noreply, %{state | monitors: monitors}}
    end
  end

  def handle_info({:weather_cache_refresh, valid_time, rows}, state) do
    # Only meaningful on the ETS backend. With Valkey, broadcasting was
    # disabled at the source and this handler will simply not fire.
    if !valkey_backend?(), do: put(valid_time, rows)
    {:noreply, state}
  end

  def handle_info(_msg, state), do: {:noreply, state}

  defp pop_monitor_for(monitors, valid_time) do
    case Enum.find(monitors, fn {_ref, vt} -> vt == valid_time end) do
      nil ->
        {monitors, nil}

      {ref, ^valid_time} ->
        Process.demonitor(ref, [:flush])
        {Map.delete(monitors, ref), ref}
    end
  end

  # ---------- Backend selector ----------

  defp valkey_backend?, do: Valkey.configured?()

  # ---------- ETS backend ----------

  defp ets_fetch(valid_time) do
    case :ets.lookup(@table, valid_time) do
      [{_, chunked}] ->
        emit_lookup(true)
        {:ok, chunks_to_list(chunked)}

      [] ->
        emit_lookup(false)
        :miss
    end
  end

  defp ets_fetch_bounds(valid_time, bounds) do
    case :ets.lookup(@table, valid_time) do
      [{_, chunked}] ->
        emit_lookup(true)
        {:ok, chunks_filtered_to_list(chunked, bounds)}

      [] ->
        emit_lookup(false)
        :miss
    end
  end

  defp ets_fetch_point(valid_time, lat, lon) do
    case :ets.lookup(@table, valid_time) do
      [{_, chunked}] ->
        chunk_key = chunk_key_for(lat, lon)

        case chunked |> Map.get(chunk_key, %{}) |> Map.get({lat, lon}) do
          nil ->
            emit_lookup(false)
            :miss

          row ->
            emit_lookup(true)
            {:ok, row}
        end

      [] ->
        emit_lookup(false)
        :miss
    end
  end

  defp ets_latest_valid_time do
    match_spec = [{{:"$1", :_}, [], [:"$1"]}]

    case :ets.select(@table, match_spec) do
      [] -> nil
      times -> Enum.max(times, DateTime)
    end
  end

  defp ets_prune_keep_latest(keep) do
    case :ets.select(@table, [{{:"$1", :_}, [], [:"$1"]}]) do
      [] ->
        0

      times ->
        sorted = Enum.sort(times, {:desc, DateTime})
        to_drop = Enum.drop(sorted, keep)
        Enum.each(to_drop, &:ets.delete(@table, &1))
        length(to_drop)
    end
  end

  # ---------- Valkey backend ----------

  defp valkey_fetch(valid_time) do
    case Valkey.scan_match(chunk_key_pattern(valid_time)) do
      {:ok, []} ->
        emit_lookup(false)
        :miss

      {:ok, keys} ->
        merge_fetched_chunks(keys)

      {:error, reason} ->
        log_valkey_error(:scan, reason)
        :miss
    end
  end

  defp merge_fetched_chunks(keys) do
    case Valkey.mget(keys) do
      {:ok, vals} ->
        chunks =
          vals
          |> Enum.reject(&is_nil/1)
          |> Enum.reduce(%{}, &Map.merge(&2, decode_chunk_value(&1)))

        if map_size(chunks) == 0 do
          emit_lookup(false)
          :miss
        else
          emit_lookup(true)
          {:ok, chunks_to_list(chunks)}
        end

      {:error, reason} ->
        log_valkey_error(:fetch, reason)
        :miss
    end
  end

  defp valkey_fetch_bounds(valid_time, nil), do: valkey_fetch(valid_time)

  defp valkey_fetch_bounds(valid_time, %{"south" => _, "north" => _, "west" => _, "east" => _} = bounds) do
    keys = chunk_keys_in_bounds(valid_time, bounds)

    case Valkey.mget(keys) do
      {:ok, vals} ->
        rows = collect_rows_in_bounds(vals, bounds)

        if rows == [] do
          emit_lookup(false)
          :miss
        else
          emit_lookup(true)
          {:ok, rows}
        end

      {:error, reason} ->
        log_valkey_error(:mget, reason)
        :miss
    end
  end

  defp collect_rows_in_bounds(vals, bounds) do
    vals
    |> Enum.reject(&is_nil/1)
    |> Enum.flat_map(fn cells ->
      cells |> decode_chunk_value() |> Enum.flat_map(&cells_for_chunk_in_bounds(&1, bounds))
    end)
  end

  defp cells_for_chunk_in_bounds({_chunk_key, cell_map}, bounds) do
    cells_in_bounds(cell_map, bounds)
  end

  defp valkey_fetch_point(valid_time, lat, lon) do
    chunk_key = chunk_key_for(lat, lon)
    redis_key = chunk_redis_key(valid_time, chunk_key)

    case Valkey.get(redis_key) do
      {:ok, %{} = chunk_map} ->
        case chunk_map |> Map.get(chunk_key, %{}) |> Map.get({lat, lon}) do
          nil ->
            emit_lookup(false)
            :miss

          row ->
            emit_lookup(true)
            {:ok, row}
        end

      :miss ->
        emit_lookup(false)
        :miss

      {:error, reason} ->
        log_valkey_error(:get, reason)
        :miss
    end
  end

  defp valkey_put(valid_time, chunked) do
    pairs =
      Enum.map(chunked, fn {chunk_key, cells} ->
        {chunk_redis_key(valid_time, chunk_key), %{chunk_key => cells}}
      end)

    case Valkey.mset_with_ttl(pairs, @ttl_seconds) do
      :ok ->
        _ = Valkey.zadd(@valkey_index_key, DateTime.to_unix(valid_time), DateTime.to_iso8601(valid_time))
        :ok

      {:error, reason} ->
        log_valkey_error(:mset, reason)
        :ok
    end
  end

  defp valkey_latest_valid_time do
    case Valkey.zrevrange(@valkey_index_key, 0, 0) do
      {:ok, [iso | _]} ->
        case DateTime.from_iso8601(iso) do
          {:ok, dt, _} -> dt
          _ -> nil
        end

      _ ->
        nil
    end
  end

  defp valkey_prune_older_than(cutoff) do
    cutoff_score = DateTime.to_unix(cutoff)

    case Valkey.zrangebyscore(@valkey_index_key, "-inf", "(#{cutoff_score}") do
      {:ok, []} ->
        0

      {:ok, members} ->
        keys = Enum.flat_map(members, &chunk_keys_for_iso/1)
        _ = Valkey.del(keys)
        _ = Valkey.zrem(@valkey_index_key, members)
        length(members)

      _ ->
        0
    end
  end

  defp chunk_keys_for_iso(iso) do
    case DateTime.from_iso8601(iso) do
      {:ok, dt, _} ->
        case Valkey.scan_match(chunk_key_pattern(dt)) do
          {:ok, ks} -> ks
          _ -> []
        end

      _ ->
        []
    end
  end

  defp valkey_clear do
    case Valkey.scan_match("#{@valkey_chunk_prefix}*") do
      {:ok, keys} -> Valkey.del(keys)
      _ -> :ok
    end

    Valkey.del([@valkey_index_key])
  end

  defp chunk_key_pattern(valid_time) do
    "#{@valkey_chunk_prefix}#{DateTime.to_iso8601(valid_time)}:*"
  end

  defp chunk_redis_key(valid_time, {lat_band, lon_band}) do
    "#{@valkey_chunk_prefix}#{DateTime.to_iso8601(valid_time)}:#{lat_band}:#{lon_band}"
  end

  defp chunk_keys_in_bounds(valid_time, %{"south" => s, "north" => n, "west" => w, "east" => e}) do
    s_band = chunk_band(s * 1.0)
    n_band = chunk_band(n * 1.0)
    w_band = chunk_band(w * 1.0)
    e_band = chunk_band(e * 1.0)

    for lat <- s_band..n_band, lon <- w_band..e_band do
      chunk_redis_key(valid_time, {lat, lon})
    end
  end

  defp cells_in_bounds(cell_map, %{"south" => s, "north" => n, "west" => w, "east" => e}) do
    for {{lat, lon}, row} <- cell_map, lat >= s, lat <= n, lon >= w, lon <= e, do: row
  end

  # The chunk value persisted to Valkey is `%{chunk_key => cells_map}`
  # so the decoder can return the cells to a uniform shape regardless of
  # whether we read 1 chunk (fetch_point/fetch_bounds) or many (fetch).
  defp decode_chunk_value(%{} = single_chunk_map), do: single_chunk_map
  defp decode_chunk_value(_), do: %{}

  defp log_valkey_error(op, reason) do
    Logger.warning("GridCache Valkey #{op} failed (#{inspect(reason)}); falling back to disk.")
  end

  defp emit_lookup(hit) do
    :telemetry.execute([:microwaveprop, :weather, :grid_cache, :lookup], %{}, %{hit: hit})
  end

  # ---------- Chunk helpers (shared) ----------

  defp list_to_chunks(rows) do
    Enum.reduce(rows, %{}, fn %{lat: lat, lon: lon} = row, acc ->
      key = chunk_key_for(lat, lon)
      chunk = Map.get(acc, key, %{})
      Map.put(acc, key, Map.put(chunk, {lat, lon}, row))
    end)
  end

  defp chunks_to_list(chunked) do
    Enum.flat_map(chunked, fn {_chunk_key, cells} ->
      Enum.map(cells, fn {_, row} -> row end)
    end)
  end

  defp chunks_filtered_to_list(chunked, nil), do: chunks_to_list(chunked)

  defp chunks_filtered_to_list(chunked, %{"south" => s, "north" => n, "west" => w, "east" => e} = bounds) do
    chunked
    |> Enum.filter(fn {chunk_key, _} -> chunk_intersects_bounds?(chunk_key, bounds) end)
    |> Enum.flat_map(fn {_, cells} ->
      for {{lat, lon}, row} <- cells, lat >= s, lat <= n, lon >= w, lon <= e, do: row
    end)
  end

  defp chunk_key_for(lat, lon) do
    {chunk_band(lat * 1.0), chunk_band(lon * 1.0)}
  end

  defp chunk_band(value) when is_float(value) do
    (value / @chunk_step) |> Float.floor() |> trunc()
  end

  defp chunk_intersects_bounds?({lat_band, lon_band}, %{"south" => s, "north" => n, "west" => w, "east" => e}) do
    chunk_south = lat_band * @chunk_step
    chunk_north = chunk_south + @chunk_step
    chunk_west = lon_band * @chunk_step
    chunk_east = chunk_west + @chunk_step

    chunk_north >= s and chunk_south <= n and chunk_east >= w and chunk_west <= e
  end
end