prop/lib/microwaveprop/weather/grib2/extractor.ex

defmodule Microwaveprop.Weather.Grib2.Extractor do
  @moduledoc false

  alias Microwaveprop.Weather.Grib2.ComplexPacking
  alias Microwaveprop.Weather.Grib2.LambertConformal
  alias Microwaveprop.Weather.Grib2.Section
  alias Microwaveprop.Weather.Grib2.SimplePacking

  @doc """
  Extract weather values from a GRIB2 binary blob at the given lat/lon.

  Returns `{:ok, %{"VAR:LEVEL" => float}}` or `{:error, term}`.
  """
  def extract_points(binary, lat, lon) do
    messages = split_messages(binary)

    Enum.reduce(messages, {:ok, %{}}, fn
      msg, {:ok, acc} ->
        case extract_single(msg, lat, lon) do
          {:ok, key, value} -> {:ok, Map.put(acc, key, value)}
          {:error, :outside_grid} -> {:error, :outside_grid}
          {:error, _reason} -> {:ok, acc}
        end

      _msg, error ->
        error
    end)
  end

  @doc """
  Extract weather values from a GRIB2 binary blob for multiple lat/lon points.

  Takes a binary and a list of `{lat, lon}` tuples.
  Returns `{:ok, %{{lat, lon} => %{"VAR:LEVEL" => float}}}` or `{:error, term}`.
  """
  def extract_grid(binary, points) do
    messages = split_messages(binary)

    result =
      Enum.reduce(messages, {:ok, init_grid(points)}, fn msg, {:ok, acc} ->
        case extract_single_grid(msg, points) do
          {:ok, point_values} ->
            merged =
              Enum.reduce(point_values, acc, fn {point, key, value}, grid ->
                Map.update!(grid, point, &Map.put(&1, key, value))
              end)

            {:ok, merged}

          {:error, _reason} ->
            # Skip messages with missing/corrupt data sections
            {:ok, acc}
        end
      end)

    case result do
      {:ok, grid} ->
        {:ok, Map.reject(grid, fn {_point, values} -> values == %{} end)}

      error ->
        error
    end
  end

  @doc """
  Split a binary blob into individual GRIB2 messages by scanning for "GRIB" magic bytes
  and reading the total length from the indicator section.
  """
  def split_messages(binary), do: split_messages(binary, [])

  @doc """
  Compute linear index from grid (i, j) and scan mode.

  For HRRR scan_mode=64 (j positive, i positive, i consecutive): index = j * nx + i
  """
  def linear_index({i, j}, nx, _scan_mode) do
    j * nx + i
  end

  # --- Private ---

  defp split_messages(<<>>, acc), do: Enum.reverse(acc)

  defp split_messages(<<"GRIB", _::32, total_length::64-big, _rest::binary>> = binary, acc) do
    if total_length > byte_size(binary) do
      Enum.reverse(acc)
    else
      msg = binary_part(binary, 0, total_length)
      remaining = binary_part(binary, total_length, byte_size(binary) - total_length)
      split_messages(remaining, [msg | acc])
    end
  end

  defp split_messages(<<_::8, rest::binary>>, acc) do
    # Skip non-GRIB bytes (padding between messages)
    split_messages(rest, acc)
  end

  defp init_grid(points) do
    Map.new(points, fn point -> {point, %{}} end)
  end

  defp extract_single_grid(msg, points) do
    with {:ok, parsed} <- Section.parse_message(msg) do
      %{grid_params: grid, product: prod, packing_params: packing, data: data} = parsed
      key = "#{prod.var}:#{prod.level}"

      results =
        Enum.reduce_while(points, {:ok, []}, fn {lat, lon} = point, {:ok, acc} ->
          case LambertConformal.to_grid_index(grid, lat, lon) do
            {:ok, {i, j}} ->
              index = linear_index({i, j}, grid.nx, grid.scan_mode)

              case unpack_value(packing, data, index) do
                {:ok, value} -> {:cont, {:ok, [{point, key, value} | acc]}}
                {:error, reason} -> {:halt, {:error, reason}}
              end

            {:error, :outside_grid} ->
              {:cont, {:ok, acc}}
          end
        end)

      case results do
        {:ok, point_values} -> {:ok, point_values}
        {:error, reason} -> {:error, reason}
      end
    end
  rescue
    e -> {:error, "GRIB2 grid extraction failed: #{inspect(e)}"}
  end

  defp extract_single(msg, lat, lon) do
    with {:ok, parsed} <- Section.parse_message(msg),
         %{grid_params: grid, product: prod, packing_params: packing, data: data} = parsed,
         key = "#{prod.var}:#{prod.level}",
         {:ok, {i, j}} <- LambertConformal.to_grid_index(grid, lat, lon),
         index = linear_index({i, j}, grid.nx, grid.scan_mode),
         {:ok, value} <- unpack_value(packing, data, index) do
      {:ok, key, value}
    end
  rescue
    e -> {:error, "GRIB2 extraction failed: #{inspect(e)}"}
  end

  defp unpack_value(%{template: 3} = params, data, index) do
    ComplexPacking.extract_value(params, data, index)
  end

  defp unpack_value(params, data, index) do
    SimplePacking.extract_value(params, data, index)
  end
end