prop/lib/microwaveprop/weather/hrrr_client.ex

defmodule Microwaveprop.Weather.HrrrClient do
  @moduledoc false

  alias Microwaveprop.Propagation.Grid
  alias Microwaveprop.Weather.Grib2.Extractor
  alias Microwaveprop.Weather.Grib2.Wgrib2

  require Logger

  @hrrr_base_default "https://noaa-hrrr-bdp-pds.s3.amazonaws.com"

  defp hrrr_base, do: Application.get_env(:microwaveprop, :hrrr_base_url, @hrrr_base_default)

  # Per-point profile fetches (contact detail + skew-T): full troposphere +
  # lower stratosphere. 25 levels from 1000→100 mb — fine-grained near the
  # surface for duct detection, loosening above 700 mb for the skew-T plot.
  @profile_pressure_levels [
    1000,
    975,
    950,
    925,
    900,
    875,
    850,
    825,
    800,
    775,
    750,
    725,
    700,
    650,
    600,
    550,
    500,
    450,
    400,
    350,
    300,
    250,
    200,
    150,
    100
  ]

  # Grid-wide fetches (PropagationGridWorker hot path): narrow to the
  # near-surface band the scorer actually reads. `SoundingParams.derive`
  # pulls `min_refractivity_gradient` from the lowest ~1 km, and native
  # hybrid-sigma data (when present) takes priority for that metric
  # anyway. Anything above 700 mb is dead weight that doubled the grid
  # worker's peak memory when 92k points × levels × vars all decode
  # through wgrib2 — enough to OOM the pod at 4 Gi. Keep the hot path
  # on this slimmer list.
  @grid_pressure_levels [
    1000,
    975,
    950,
    925,
    900,
    875,
    850,
    825,
    800,
    775,
    750,
    725,
    700
  ]

  @surface_messages [
    %{var: "TMP", level: "2 m above ground"},
    %{var: "DPT", level: "2 m above ground"},
    %{var: "PRES", level: "surface"},
    %{var: "HPBL", level: "surface"},
    %{var: "PWAT", level: "entire atmosphere (considered as a single layer)"},
    %{var: "UGRD", level: "10 m above ground"},
    %{var: "VGRD", level: "10 m above ground"},
    %{var: "TCDC", level: "entire atmosphere"},
    %{var: "APCP", level: "surface"}
  ]

  # --- Public API ---

  @spec surface_messages() :: [%{var: String.t(), level: String.t()}]
  def surface_messages, do: @surface_messages

  @spec fetch_grid([{float(), float()}], DateTime.t(), keyword()) ::
          {:ok, %{{float(), float()} => map()}} | {:error, term()}
  def fetch_grid(points, run_time, opts \\ []) do
    Microwaveprop.Instrument.span(
      [:hrrr, :fetch_grid],
      %{point_count: length(points), forecast_hour: Keyword.get(opts, :forecast_hour, 0)},
      fn -> do_fetch_grid(points, run_time, opts) end
    )
  end

  defp do_fetch_grid(points, run_time, opts) do
    hour_dt = nearest_hrrr_hour(run_time)
    date = DateTime.to_date(hour_dt)
    hour = hour_dt.hour
    forecast_hour = Keyword.get(opts, :forecast_hour, 0)

    with {:ok, sfc_grid} <- fetch_product_grid(date, hour, :surface, points, forecast_hour) do
      # Pressure is optional — if it fails, continue with surface-only data
      prs_grid =
        case maybe_fetch_pressure_grid(date, hour, points, opts, forecast_hour) do
          {:ok, grid} ->
            grid

          {:error, reason} ->
            Logger.warning("HRRR pressure fetch failed (continuing with surface only): #{inspect(reason)}")
            %{}
        end

      # Merge and build profiles in one pass to avoid a full intermediate map copy.
      # Use Map.merge to include any pressure-only points, then map directly to profiles.
      profiles =
        sfc_grid
        |> Map.merge(prs_grid, fn _point, sfc_data, prs_data -> Map.merge(sfc_data, prs_data) end)
        |> Map.new(fn {point, data} ->
          profile = build_profile(data)
          {point, Map.put(profile, :run_time, hour_dt)}
        end)

      {:ok, profiles}
    end
  end

  @spec fetch_profile(float(), float(), DateTime.t()) :: {:ok, map()} | {:error, term()}
  def fetch_profile(lat, lon, valid_time) do
    Microwaveprop.Instrument.span(
      [:hrrr, :fetch_profile],
      %{lat: lat, lon: lon},
      fn -> do_fetch_profile(lat, lon, valid_time) end
    )
  end

  defp do_fetch_profile(lat, lon, valid_time) do
    hour_dt = nearest_hrrr_hour(valid_time)
    date = DateTime.to_date(hour_dt)
    hour = hour_dt.hour

    with {:ok, sfc_data} <- fetch_product(date, hour, :surface, lat, lon) do
      # Pressure is optional — old HRRR data may have corrupt messages
      prs_data =
        case fetch_product(date, hour, :pressure, lat, lon) do
          {:ok, data} ->
            data

          {:error, reason} ->
            Logger.warning(
              "HRRR pressure fetch failed for #{lat},#{lon} (continuing with surface only): #{inspect(reason)}"
            )

            %{}
        end

      merged = Map.merge(sfc_data, prs_data)
      result = build_profile(merged)
      {:ok, Map.put(result, :run_time, hour_dt)}
    end
  end

  @spec nearest_hrrr_hour(DateTime.t()) :: DateTime.t()
  def nearest_hrrr_hour(dt) do
    total_seconds = dt.minute * 60 + dt.second
    rounded_dt = DateTime.add(dt, -total_seconds, :second)

    if dt.minute >= 30 do
      DateTime.add(rounded_dt, 3600, :second)
    else
      rounded_dt
    end
  end

  @spec hrrr_url(Date.t(), non_neg_integer(), :surface | :pressure, non_neg_integer()) :: String.t()
  def hrrr_url(date, hour, product, forecast_hour \\ 0) do
    date_str = Calendar.strftime(date, "%Y%m%d")
    hour_str = hour |> Integer.to_string() |> String.pad_leading(2, "0")
    fh_str = forecast_hour |> Integer.to_string() |> String.pad_leading(2, "0")

    file =
      case product do
        :surface -> "hrrr.t#{hour_str}z.wrfsfcf#{fh_str}.grib2"
        :pressure -> "hrrr.t#{hour_str}z.wrfprsf#{fh_str}.grib2"
      end

    "#{hrrr_base()}/hrrr.#{date_str}/conus/#{file}"
  end

  @spec parse_idx(String.t()) :: [%{msg: integer(), offset: integer(), var: String.t(), level: String.t()}]
  def parse_idx(text) do
    text
    |> String.split("\n")
    |> Enum.reject(&(&1 == ""))
    |> Enum.map(fn line ->
      parts = String.split(line, ":", parts: 8)

      %{
        msg: String.to_integer(Enum.at(parts, 0)),
        offset: String.to_integer(Enum.at(parts, 1)),
        var: Enum.at(parts, 3),
        level: Enum.at(parts, 4)
      }
    end)
  end

  @spec byte_ranges_for_messages([map()], [%{var: String.t(), level: String.t()}]) :: [
          {non_neg_integer(), non_neg_integer()}
        ]
  def byte_ranges_for_messages(idx_entries, wanted) do
    Enum.flat_map(wanted, fn w ->
      idx_entries
      |> Enum.with_index()
      |> Enum.flat_map(fn {entry, idx} ->
        byte_range_for_entry(entry, idx, idx_entries, w)
      end)
    end)
  end

  defp byte_range_for_entry(entry, idx, idx_entries, wanted) do
    if entry.var == wanted.var && entry.level == wanted.level do
      end_offset =
        case Enum.at(idx_entries, idx + 1) do
          nil -> entry.offset + 10_000_000
          next -> next.offset - 1
        end

      [{entry.offset, end_offset}]
    else
      []
    end
  end

  @spec merge_grid_data(map(), map()) :: map()
  def merge_grid_data(sfc_grid, prs_grid) do
    Map.merge(sfc_grid, prs_grid, fn _point, sfc_data, prs_data ->
      Map.merge(sfc_data, prs_data)
    end)
  end

  @spec build_profile(map()) :: map()
  def build_profile(parsed) do
    sfc_temp_k = parsed["TMP:2 m above ground"]
    sfc_dpt_k = parsed["DPT:2 m above ground"]
    sfc_pres_pa = parsed["PRES:surface"]

    # Iterate the full profile list so per-contact fetches that include
    # upper-air data still populate the whole skew-T trace. Grid fetches
    # only populate the 1000→700 mb subset (see @grid_pressure_levels);
    # the other levels just produce empty slots that this loop skips.
    profile =
      Enum.flat_map(@profile_pressure_levels, fn level ->
        level_str = "#{level} mb"
        tmp = parsed["TMP:#{level_str}"]
        dpt = parsed["DPT:#{level_str}"]
        hgt = parsed["HGT:#{level_str}"]

        if tmp && dpt && hgt do
          [
            %{
              "pres" => level * 1.0,
              "tmpc" => tmp - 273.15,
              "dwpc" => dpt - 273.15,
              "hght" => hgt
            }
          ]
        else
          []
        end
      end)

    # Fall back to lowest pressure level if surface fields are missing (near-coast points)
    lowest_level = List.first(profile)

    sfc_temp_c =
      if sfc_temp_k, do: sfc_temp_k - 273.15, else: lowest_level && lowest_level["tmpc"]

    sfc_dpt_c =
      if sfc_dpt_k, do: sfc_dpt_k - 273.15, else: lowest_level && lowest_level["dwpc"]

    %{
      surface_temp_c: sfc_temp_c,
      surface_dewpoint_c: sfc_dpt_c,
      surface_pressure_mb: if(sfc_pres_pa, do: sfc_pres_pa / 100.0),
      hpbl_m: parsed["HPBL:surface"],
      pwat_mm: parsed["PWAT:entire atmosphere (considered as a single layer)"],
      wind_u: parsed["UGRD:10 m above ground"],
      wind_v: parsed["VGRD:10 m above ground"],
      cloud_cover_pct: parsed["TCDC:entire atmosphere"],
      precip_mm: parsed["APCP:surface"],
      profile: profile
    }
  end

  # --- Private ---

  defp fetch_product(date, hour, product, lat, lon) do
    url = hrrr_url(date, hour, product)
    idx_url = url <> ".idx"

    wanted =
      case product do
        :surface -> @surface_messages
        :pressure -> pressure_messages(:profile)
      end

    Logger.info("HRRR fetching #{product} idx from #{idx_url}")

    with {:ok, idx_text} <- fetch_idx(idx_url),
         idx_entries = parse_idx(idx_text),
         ranges = byte_ranges_for_messages(idx_entries, wanted),
         _ = Logger.info("HRRR downloading #{length(ranges)} GRIB ranges for #{product}"),
         {:ok, grib_binary} <- download_grib_ranges(url, ranges) do
      Logger.info("HRRR #{product} downloaded #{byte_size(grib_binary)} bytes, extracting")
      Extractor.extract_points(grib_binary, lat, lon)
    end
  end

  defp fetch_product_grid(date, hour, product, _points, forecast_hour) do
    url = hrrr_url(date, hour, product, forecast_hour)
    idx_url = url <> ".idx"

    wanted =
      case product do
        :surface -> @surface_messages
        :pressure -> pressure_messages(:grid)
      end

    Logger.info("HRRR grid fetching #{product} idx from #{idx_url}")

    with {:ok, idx_text} <- fetch_idx(idx_url),
         idx_entries = parse_idx(idx_text),
         ranges = byte_ranges_for_messages(idx_entries, wanted),
         _ = Logger.info("HRRR grid downloading #{length(ranges)} GRIB ranges for #{product}"),
         {:ok, grib_binary} <- download_grib_ranges(url, ranges) do
      Logger.info("HRRR grid #{product} downloaded #{byte_size(grib_binary)} bytes, extracting")
      extract_grid(grib_binary, wanted, product)
    end
  end

  defp extract_grid(grib_binary, wanted, product) do
    if Wgrib2.available?() do
      match_pattern = wgrib2_match_pattern(wanted)
      grid_spec = Grid.wgrib2_grid_spec()
      Logger.info("HRRR grid using wgrib2 for #{product} extraction")
      Wgrib2.extract_grid(grib_binary, match_pattern, grid_spec)
    else
      Logger.info("HRRR grid using Elixir decoder for #{product} extraction (slow)")
      points = Grid.conus_points()
      Extractor.extract_grid(grib_binary, points)
    end
  end

  defp wgrib2_match_pattern(wanted) do
    vars = wanted |> Enum.map(& &1.var) |> Enum.uniq() |> Enum.join("|")
    ":(#{vars}):"
  end

  defp maybe_fetch_pressure_grid(date, hour, points, opts, forecast_hour) do
    if Keyword.get(opts, :pressure, true) do
      fetch_product_grid(date, hour, :pressure, points, forecast_hour)
    else
      {:ok, %{}}
    end
  end

  @doc """
  Pressure-level GRIB message descriptors for a fetch variant.

    * `:grid` — narrow list of 13 levels (1000→700 mb) used by the
      propagation grid hot path. Covers the ducting band without
      ballooning the per-forecast-hour memory footprint.
    * `:profile` — full 25 levels (1000→100 mb) used for per-contact
      skew-T fetches.
  """
  @spec pressure_messages(:grid | :profile) :: [%{var: String.t(), level: String.t()}]
  def pressure_messages(:grid) do
    build_pressure_messages(@grid_pressure_levels)
  end

  def pressure_messages(:profile) do
    build_pressure_messages(@profile_pressure_levels)
  end

  defp build_pressure_messages(levels) do
    for level <- levels, var <- ["TMP", "DPT", "HGT"] do
      %{var: var, level: "#{level} mb"}
    end
  end

  defp fetch_idx(url) do
    Microwaveprop.Instrument.span([:hrrr, :fetch_idx], %{url: url}, fn ->
      case Req.get(url, req_options()) do
        {:ok, %{status: 200, body: body}} ->
          {:ok, body}

        {:ok, %{status: status}} ->
          {:error, "HRRR idx HTTP #{status}"}

        {:error, reason} ->
          {:error, reason}
      end
    end)
  end

  @doc """
  Downloads the given byte ranges from a GRIB2 URL and concatenates
  them in offset order. Merges adjacent/overlapping ranges and runs
  requests in parallel. Backed by a disk cache keyed on
  `(url, ranges)` in dev.

  Public so the native-level client can reuse it rather than
  duplicating the HTTP/parallelism/cache logic.
  """
  @spec download_grib_ranges(String.t(), [{non_neg_integer(), non_neg_integer()}]) ::
          {:ok, binary()} | {:error, term()}
  def download_grib_ranges(_url, []), do: {:ok, <<>>}

  def download_grib_ranges(url, ranges) do
    # Check local cache first (dev only)
    cache_key = url_to_cache_key(url, ranges)

    case read_cache(cache_key) do
      {:ok, binary} ->
        Logger.info("HRRR cache hit: #{cache_key}")
        {:ok, binary}

      :miss ->
        case download_grib_ranges_remote(url, ranges) do
          {:ok, binary} = ok ->
            write_cache(cache_key, binary)
            ok

          error ->
            error
        end
    end
  end

  @doc """
  Like `download_grib_ranges/2`, but writes the result to `dest_path`
  instead of returning a binary. Each chunk is written to disk as it
  arrives, so peak memory stays low even for ~530 MB native-level files.

  Returns `:ok` or `{:error, reason}`.
  """
  @spec download_grib_ranges_to_file(String.t(), [{non_neg_integer(), non_neg_integer()}], Path.t()) ::
          :ok | {:error, term()}
  def download_grib_ranges_to_file(_url, [], _dest_path), do: :ok

  def download_grib_ranges_to_file(url, ranges, dest_path) do
    merged = ranges |> merge_ranges() |> Enum.sort_by(&elem(&1, 0))

    file = File.open!(dest_path, [:write, :binary])

    try do
      Enum.reduce_while(merged, :ok, fn {start, stop}, :ok ->
        case Req.get(url, [{:headers, [{"Range", "bytes=#{start}-#{stop}"}]} | req_options()]) do
          {:ok, %{status: 206, body: body}} ->
            IO.binwrite(file, body)
            {:cont, :ok}

          {:ok, %{status: status}} ->
            {:halt, {:error, "HRRR grib HTTP #{status}"}}

          {:error, reason} ->
            {:halt, {:error, reason}}
        end
      end)
    after
      File.close(file)
    end
  end

  defp download_grib_ranges_remote(url, ranges) do
    # Merge adjacent/overlapping ranges to reduce HTTP requests
    merged = merge_ranges(ranges)

    # Download ranges in parallel
    results =
      merged
      |> Task.async_stream(
        fn {start, stop} ->
          case Req.get(url, [{:headers, [{"Range", "bytes=#{start}-#{stop}"}]} | req_options()]) do
            {:ok, %{status: 206, body: body}} -> {:ok, start, body}
            {:ok, %{status: status}} -> {:error, "HRRR grib HTTP #{status}"}
            {:error, reason} -> {:error, reason}
          end
        end,
        max_concurrency: 8,
        timeout: 120_000
      )
      |> Enum.map(fn
        {:ok, result} -> result
        {:exit, reason} -> {:error, reason}
      end)

    # Sort by offset and concatenate
    case Enum.find(results, &match?({:error, _}, &1)) do
      {:error, reason} ->
        {:error, reason}

      nil ->
        binary =
          results
          |> Enum.sort_by(fn {:ok, offset, _} -> offset end)
          |> Enum.map(fn {:ok, _, body} -> body end)
          |> IO.iodata_to_binary()

        {:ok, binary}
    end
  end

  defp hrrr_cache_dir, do: Application.get_env(:microwaveprop, :hrrr_cache_dir)

  defp url_to_cache_key(url, ranges) do
    # e.g. hrrr.20240922_conus_hrrr.t18z.wrfprsf00_r42.grib2
    uri = URI.parse(url)
    path_part = uri.path |> String.trim_leading("/") |> String.replace("/", "_")
    range_hash = ranges |> :erlang.phash2() |> Integer.to_string(16)
    "#{path_part}_r#{range_hash}.grib2"
  end

  defp read_cache(key) do
    case hrrr_cache_dir() do
      nil ->
        :miss

      dir ->
        path = Path.join(dir, key)

        case File.read(path) do
          {:ok, binary} -> {:ok, binary}
          {:error, _} -> :miss
        end
    end
  end

  defp write_cache(key, binary) do
    case hrrr_cache_dir() do
      nil ->
        :ok

      dir ->
        File.mkdir_p!(dir)
        path = Path.join(dir, key)
        File.write!(path, binary)
        Logger.info("HRRR cached: #{path} (#{byte_size(binary)} bytes)")
    end
  end

  defp merge_ranges(ranges) do
    ranges
    |> Enum.sort_by(&elem(&1, 0))
    |> Enum.reduce([], fn {s, e}, acc ->
      case acc do
        [{ps, pe} | rest] when s <= pe + 1 -> [{ps, max(pe, e)} | rest]
        _ -> [{s, e} | acc]
      end
    end)
    |> Enum.reverse()
  end

  defp req_options do
    defaults = [receive_timeout: 120_000, retry: &retry?/2, max_retries: 5, retry_delay: &retry_delay/1]
    overrides = Application.get_env(:microwaveprop, :hrrr_req_options, [])
    Keyword.merge(defaults, overrides)
  end

  defp retry?(_request, response) do
    case response do
      %Req.Response{status: status} when status in [429, 500, 502, 503, 504] -> true
      %{__exception__: true} -> true
      _ -> false
    end
  end

  defp retry_delay(n) do
    base = Integer.pow(2, n) * 1_000
    jitter = :rand.uniform(1_000)
    base + jitter
  end
end