defmodule Microwaveprop.Weather.GridCache do
  @moduledoc """
  Node-local ETS cache of derived HRRR grid rows keyed by `valid_time`. Mirrors
  `Microwaveprop.Propagation.ScoreCache` but for the `/weather` map.

  The Weather map LiveView calls `latest_weather_grid/1` on mount and every
  pan/zoom. Each call otherwise hits the 42M-row partitioned `hrrr_profiles`
  table, runs per-row `derive_and_clean` transforms, and returns 3-10k rows.
  With this cache those calls become in-memory map iterations.

  ## Storage layout

  Each cache entry is bucketed into 5°×5° spatial chunks matching
  `Microwaveprop.Weather.ScalarFile`'s on-disk layout. The ETS value for
  `valid_time` is `%{{lat_band, lon_band} => %{{lat, lon} => row}}`. That
  way `fetch_bounds/2` only walks the chunks that intersect the requested
  viewport instead of the full 92k-cell CONUS map. `fetch_point/3` reads
  exactly one chunk.

  Populated by `Microwaveprop.Weather.warm_grid_cache/1` after the hourly
  worker upserts new HRRR data, fanned out across the cluster via the
  `"weather:cache"` PubSub topic so every node stays in sync.
  """
  use GenServer

  alias Phoenix.PubSub

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

  @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__)

  @spec fetch(DateTime.t()) :: {:ok, [row()]} | :miss
  def 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

  @spec fetch_bounds(DateTime.t(), bounds() | nil) :: {:ok, [row()]} | :miss
  def 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

  @spec fetch_point(DateTime.t(), float(), float()) :: {:ok, row()} | :miss
  def 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 emit_lookup(hit) do
    :telemetry.execute([:microwaveprop, :weather, :grid_cache, :lookup], %{}, %{hit: hit})
  end

  @spec put(DateTime.t(), [row()]) :: :ok
  def put(valid_time, rows) do
    chunked = list_to_chunks(rows)
    :ets.insert(@table, {valid_time, chunked})
    :ok
  end

  @doc "Insert locally AND broadcast to peer nodes via PubSub."
  @spec broadcast_put(DateTime.t(), [row()]) :: :ok
  def broadcast_put(valid_time, rows) do
    _ = PubSub.broadcast(@pubsub, @topic, {:weather_cache_refresh, valid_time, rows})
    :ok
  end

  @spec latest_valid_time() :: DateTime.t() | nil
  def latest_valid_time do
    match_spec = [{{:"$1", :_}, [], [:"$1"]}]

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

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

  @doc """
  Keep only the `keep` most-recent valid_times, dropping the rest. Returns
  the number of entries removed. Used to bound memory when forecast hours
  are cached on demand.
  """
  @spec prune_keep_latest(non_neg_integer()) :: non_neg_integer()
  def prune_keep_latest(keep) when is_integer(keep) and keep >= 0 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

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

  @doc """
  Atomically claim the right to fill the cache for `valid_time`. Returns
  `true` if this caller won the claim and should run the fill; `false` if
  another caller is already filling. Prevents N concurrent /weather mounts
  after a pod restart from each firing the 15-second cold-fill read and
  starving the Postgres connection pool.

  The GenServer `Process.monitor`s the caller: if the caller crashes
  before calling `release_fill/1`, the lock is released automatically.
  """
  @spec claim_fill(term()) :: boolean()
  def claim_fill(valid_time) do
    GenServer.call(__MODULE__, {:claim_fill, valid_time, self()})
  end

  @doc "Release a fill lock claimed via `claim_fill/1`."
  @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)
  end

  @impl true
  def init(_opts) do
    _ = :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
    :ets.delete_all_objects(@table)

    # Demonitor all tracked callers and drop every lock so tests start clean.
    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
    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

  # ---------- Chunk helpers ----------

  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