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 "https://noaa-hrrr-bdp-pds.s3.amazonaws.com"

  @pressure_levels [1000, 975, 950, 925, 900, 850, 800, 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 ---

  def surface_messages, do: @surface_messages

  def 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

      merged = merge_grid_data(sfc_grid, prs_grid)

      profiles =
        Map.new(merged, fn {point, data} ->
          profile = build_profile(data)
          {point, Map.put(profile, :run_time, hour_dt)}
        end)

      {:ok, profiles}
    end
  end

  def 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

  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

  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

  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

  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} ->
        if entry.var == w.var && entry.level == w.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)
    end)
  end

  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

  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"]

    profile =
      Enum.flat_map(@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()
      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()
      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

  defp pressure_messages do
    for level <- @pressure_levels, var <- ["TMP", "DPT", "HGT"] do
      %{var: var, level: "#{level} mb"}
    end
  end

  defp fetch_idx(url) do
    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

  defp download_grib_ranges(_url, []), do: {:ok, <<>>}

  defp download_grib_ranges(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: 60_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 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: 60_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