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)

    results =
      Enum.reduce_while(messages, {:ok, %{}}, fn msg, {:ok, acc} ->
        case extract_single(msg, lat, lon) do
          {:ok, key, value} -> {:cont, {:ok, Map.put(acc, key, value)}}
          {:error, :outside_grid} -> {:halt, {:error, :outside_grid}}
          {:error, reason} -> {:halt, {:error, reason}}
        end
      end)

    results
  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 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