aprs.me/lib/aprsme/spatial_pubsub.ex

defmodule Aprsme.SpatialPubSub do
  @moduledoc """
  Spatial-aware PubSub system that broadcasts packets only to clients
  whose viewports contain the packet's location.
  """
  use GenServer

  require Logger

  # Grid size in degrees for spatial indexing
  @grid_size 1.0
  # Maximum number of concurrent clients
  @max_clients 10_000

  def start_link(opts \\ []) do
    GenServer.start_link(__MODULE__, opts, name: __MODULE__)
  end

  @doc """
  Register a client with their current viewport bounds.
  """
  def register_viewport(client_id, bounds) do
    with_server({:error, :not_running}, fn ->
      GenServer.call(__MODULE__, {:register_viewport, client_id, bounds})
    end)
  end

  @doc """
  Update a client's viewport bounds.
  """
  def update_viewport(client_id, bounds) do
    with_server({:error, :not_running}, fn ->
      GenServer.call(__MODULE__, {:update_viewport, client_id, bounds})
    end)
  end

  @doc """
  Unregister a client.
  """
  def unregister_client(client_id) do
    with_server(:ok, fn ->
      GenServer.cast(__MODULE__, {:unregister_client, client_id})
    end)
  end

  @doc """
  Broadcast a packet to all clients whose viewports contain the packet's location.
  """
  def broadcast_packet(packet) do
    with_server(:ok, fn ->
      GenServer.cast(__MODULE__, {:broadcast_packet, packet})
    end)
  end

  @doc """
  Get statistics about spatial filtering.
  """
  def get_stats do
    with_server(
      %{
        total_broadcasts: 0,
        filtered_broadcasts: 0,
        total_packets: 0,
        clients_count: 0,
        grid_cells: 0,
        avg_clients_per_cell: 0.0
      },
      fn ->
        GenServer.call(__MODULE__, :get_stats)
      end
    )
  end

  @doc """
  Start telemetry reporting for LiveDashboard integration.
  """
  def start_telemetry_reporting do
    with_server(:ok, fn ->
      GenServer.cast(__MODULE__, :start_telemetry_reporting)
    end)
  end

  # Server callbacks

  @impl true
  def init(_opts) do
    # Start telemetry reporting
    Process.send_after(self(), :report_telemetry, 5_000)

    state = %{
      # client_id => %{bounds: bounds, topic: topic, pid: pid}
      clients: %{},
      # grid_key => MapSet of client_ids
      spatial_index: %{},
      # Statistics
      stats: %{
        total_broadcasts: 0,
        filtered_broadcasts: 0,
        total_packets: 0,
        clients_count: 0
      }
    }

    {:ok, state}
  end

  @impl true
  def handle_call({:register_viewport, client_id, bounds}, {pid, _}, state) do
    if map_size(state.clients) >= @max_clients do
      {:reply, {:error, :client_limit_exceeded}, state}
    else
      topic = "spatial:#{client_id}"

      state = replace_existing_client(state, client_id)

      ref = Process.monitor(pid)

      client_info = %{
        bounds: normalize_bounds(bounds),
        topic: topic,
        pid: pid,
        monitor_ref: ref
      }

      new_state =
        state
        |> put_in([:clients, client_id], client_info)
        |> update_spatial_index(client_id, client_info.bounds)
        |> update_in([:stats, :clients_count], &(&1 + 1))

      {:reply, {:ok, topic}, new_state}
    end
  end

  @impl true
  def handle_call({:update_viewport, client_id, bounds}, _from, state) do
    case Map.get(state.clients, client_id) do
      nil ->
        {:reply, {:error, :not_registered}, state}

      client_info ->
        # Remove old spatial index entries
        old_bounds = client_info.bounds
        state = remove_from_spatial_index(state, client_id, old_bounds)

        # Update with new bounds
        normalized_bounds = normalize_bounds(bounds)
        updated_client = %{client_info | bounds: normalized_bounds}

        new_state =
          state
          |> put_in([:clients, client_id], updated_client)
          |> update_spatial_index(client_id, normalized_bounds)

        {:reply, :ok, new_state}
    end
  end

  @impl true
  def handle_call(:get_stats, _from, state) do
    stats =
      Map.merge(state.stats, %{
        grid_cells: map_size(state.spatial_index),
        avg_clients_per_cell: calculate_avg_clients_per_cell(state.spatial_index)
      })

    {:reply, stats, state}
  end

  @impl true
  def handle_cast({:unregister_client, client_id}, state) do
    new_state = remove_client(state, client_id)
    {:noreply, new_state}
  end

  @impl true
  def handle_cast({:broadcast_packet, packet}, state) do
    state = update_in(state, [:stats, :total_packets], &(&1 + 1))

    case extract_location(packet) do
      {lat, lon} when is_number(lat) and is_number(lon) ->
        # Find all clients whose viewports contain this location
        client_ids = find_clients_for_location(state, lat, lon)

        # Optimize: Batch collect topics then spawn async task for broadcasts
        topics =
          client_ids
          |> Enum.map(fn client_id -> Map.get(state.clients, client_id) end)
          |> Enum.reject(&is_nil/1)
          |> Enum.map(& &1.topic)

        # Use dedicated broadcast task supervisor for better performance
        _ =
          Aprsme.BroadcastTaskSupervisor.broadcast_async(
            topics,
            {:spatial_packet, packet}
          )

        # Update statistics (immediately return control to GenServer)
        state =
          state
          |> update_in([:stats, :total_broadcasts], &(&1 + length(client_ids)))
          |> update_in([:stats, :filtered_broadcasts], &(&1 + max(0, map_size(state.clients) - length(client_ids))))

        {:noreply, state}

      _ ->
        # No valid location, skip broadcasting
        {:noreply, state}
    end
  end

  @impl true
  def handle_info({:DOWN, _ref, :process, pid, _reason}, state) do
    # Find and remove the client associated with this pid
    client_id =
      Enum.find_value(state.clients, fn {id, %{pid: client_pid}} ->
        if client_pid == pid, do: id
      end)

    new_state =
      if client_id do
        remove_client(state, client_id)
      else
        state
      end

    {:noreply, new_state}
  end

  @impl true
  def handle_info(:report_telemetry, state) do
    # Emit telemetry metrics
    emit_telemetry_metrics(state)

    # Schedule next report
    Process.send_after(self(), :report_telemetry, 5_000)

    {:noreply, state}
  end

  @impl true
  def handle_info(_, state), do: {:noreply, state}

  # Private functions

  defp normalize_bounds(%{north: n, south: s, east: e, west: w}) do
    %{
      north: ensure_float(n),
      south: ensure_float(s),
      east: ensure_float(e),
      west: ensure_float(w)
    }
  end

  defp normalize_bounds(invalid) do
    Logger.warning("normalize_bounds called with invalid bounds map: #{inspect(invalid)}")
    %{north: 0.0, south: 0.0, east: 0.0, west: 0.0}
  end

  defp ensure_float(val) when is_binary(val) do
    case Float.parse(val) do
      {f, _} -> f
      :error -> 0.0
    end
  end

  defp ensure_float(val) when is_integer(val), do: val * 1.0
  defp ensure_float(val) when is_float(val), do: val
  defp ensure_float(_), do: 0.0

  defp update_spatial_index(state, client_id, bounds) do
    # Get all grid cells that intersect with the bounds
    grid_cells = get_intersecting_grid_cells(bounds)

    # Add client to each grid cell
    Enum.reduce(grid_cells, state, fn grid_key, acc_state ->
      update_in(acc_state, [:spatial_index, grid_key], fn
        nil -> MapSet.new([client_id])
        set -> MapSet.put(set, client_id)
      end)
    end)
  end

  defp remove_from_spatial_index(state, client_id, bounds) do
    grid_cells = get_intersecting_grid_cells(bounds)

    Enum.reduce(grid_cells, state, fn grid_key, acc_state ->
      remove_client_from_grid_cell(acc_state, grid_key, client_id)
    end)
  end

  defp remove_client_from_grid_cell(state, grid_key, client_id) do
    case Map.get(state.spatial_index, grid_key) do
      nil ->
        state

      set ->
        new_set = MapSet.delete(set, client_id)
        update_grid_cell_after_removal(state, grid_key, new_set)
    end
  end

  defp update_grid_cell_after_removal(state, grid_key, new_set) do
    if MapSet.size(new_set) == 0 do
      update_in(state, [:spatial_index], &Map.delete(&1, grid_key))
    else
      put_in(state, [:spatial_index, grid_key], new_set)
    end
  end

  defp get_intersecting_grid_cells(%{north: n, south: s, east: e, west: w}) do
    min_lat_cell = floor(s / @grid_size)
    max_lat_cell = floor(n / @grid_size)
    min_lon_cell = floor(w / @grid_size)
    max_lon_cell = floor(e / @grid_size)

    lon_cells =
      if min_lon_cell > max_lon_cell do
        max_positive = floor(180.0 / @grid_size)
        min_negative = floor(-180.0 / @grid_size)
        Enum.to_list(min_lon_cell..max_positive) ++ Enum.to_list(min_negative..max_lon_cell)
      else
        Enum.to_list(min_lon_cell..max_lon_cell)
      end

    for lat_cell <- min_lat_cell..max_lat_cell,
        lon_cell <- lon_cells do
      {lat_cell, lon_cell}
    end
  end

  defp find_clients_for_location(state, lat, lon) do
    grid_key = {floor(lat / @grid_size), floor(lon / @grid_size)}

    case Map.get(state.spatial_index, grid_key) do
      nil ->
        []

      client_set ->
        filter_clients_by_bounds(state, client_set, lat, lon)
    end
  end

  defp filter_clients_by_bounds(state, client_set, lat, lon) do
    Enum.filter(client_set, fn client_id ->
      client_contains_point?(state, client_id, lat, lon)
    end)
  end

  defp client_contains_point?(state, client_id, lat, lon) do
    case Map.get(state.clients, client_id) do
      %{bounds: bounds} -> point_in_bounds?(lat, lon, bounds)
      _ -> false
    end
  end

  defp point_in_bounds?(lat, lon, %{north: n, south: s, east: e, west: w}) do
    lat_in_bounds = lat >= s and lat <= n

    # Handle longitude wrap-around at international date line
    lon_in_bounds =
      if w > e do
        # Bounds cross the date line
        lon >= w or lon <= e
      else
        # Normal bounds
        lon >= w and lon <= e
      end

    lat_in_bounds and lon_in_bounds
  end

  defp extract_location(packet) do
    case packet do
      %{lat: lat, lon: lon} when not is_nil(lat) and not is_nil(lon) ->
        {ensure_float(lat), ensure_float(lon)}

      _ ->
        nil
    end
  end

  defp remove_client(state, client_id) do
    case Map.get(state.clients, client_id) do
      nil ->
        state

      %{bounds: bounds, monitor_ref: ref} ->
        Process.demonitor(ref, [:flush])

        state
        |> remove_from_spatial_index(client_id, bounds)
        |> update_in([:clients], &Map.delete(&1, client_id))
        |> update_in([:stats, :clients_count], &max(0, &1 - 1))

      %{bounds: bounds} ->
        state
        |> remove_from_spatial_index(client_id, bounds)
        |> update_in([:clients], &Map.delete(&1, client_id))
        |> update_in([:stats, :clients_count], &max(0, &1 - 1))
    end
  end

  defp replace_existing_client(state, client_id) do
    case Map.get(state.clients, client_id) do
      nil -> state
      _client_info -> remove_client(state, client_id)
    end
  end

  defp calculate_avg_clients_per_cell(spatial_index) do
    non_nil_cells =
      spatial_index
      |> Map.values()
      |> Enum.reject(&is_nil/1)

    if Enum.empty?(non_nil_cells) do
      0.0
    else
      total_clients =
        non_nil_cells
        |> Enum.map(&MapSet.size/1)
        |> Enum.sum()

      total_clients / length(non_nil_cells)
    end
  end

  defp emit_telemetry_metrics(state) do
    # Client metrics
    :telemetry.execute(
      [:aprsme, :spatial_pubsub, :clients],
      %{
        count: state.stats.clients_count,
        grid_cells: map_size(state.spatial_index),
        avg_clients_per_cell: calculate_avg_clients_per_cell(state.spatial_index)
      },
      %{}
    )

    # Broadcast metrics
    :telemetry.execute(
      [:aprsme, :spatial_pubsub, :broadcasts],
      %{
        total: state.stats.total_broadcasts,
        filtered: state.stats.filtered_broadcasts,
        packets: state.stats.total_packets
      },
      %{}
    )

    # Calculate and emit efficiency metrics
    total_potential_broadcasts = state.stats.total_packets * state.stats.clients_count

    efficiency =
      if total_potential_broadcasts > 0 do
        1.0 - state.stats.total_broadcasts / total_potential_broadcasts
      else
        0.0
      end

    :telemetry.execute(
      [:aprsme, :spatial_pubsub, :efficiency],
      %{
        ratio: efficiency,
        saved_broadcasts: state.stats.filtered_broadcasts
      },
      %{}
    )
  end

  defp with_server(default, fun) do
    if Process.whereis(__MODULE__) do
      fun.()
    else
      default
    end
  end
end