defmodule Microwaveprop.Weather.GefsClient do
  @moduledoc """
  Client for the NCEP Global Ensemble Forecast System (GEFS) served
  anonymously from NOMADS at
  `https://nomads.ncep.noaa.gov/pub/data/nccf/com/gens/prod/`.

  GEFS is the 0.5° / 31-member physics-based global ensemble. It runs
  four times per day (00/06/12/18Z) with 3-hourly output to +240 h and
  6-hourly output to +384 h. The `pgrb2ap5/` product family contains
  the subset of surface and pressure-level variables the propagation
  scorer needs for extended-horizon (Day 2-7) outlooks, once HRRR's
  18 h horizon runs out.

  This module exposes the URL, time, and message-inventory helpers.
  Byte-range idx parsing and HTTP range fetch logic are shared with
  `Microwaveprop.Weather.HrrrClient` since the GRIB2/idx format is
  identical; the GEFS-specific pieces are the NOMADS path shape, the
  3-then-6-hourly forecast-hour list, and the Magnus dewpoint
  derivation (GEFS pgrb2a publishes RH rather than Td).
  """

  require Logger

  @base_url "https://nomads.ncep.noaa.gov/pub/data/nccf/com/gens/prod"

  @run_hours [0, 6, 12, 18]

  # Variables the propagation scorer reads. All present in pgrb2a f000-f384
  # per an April 2026 .idx probe. DPT is absent from GEFS output — derive
  # it from TMP + RH via Magnus (see `dewpoint_from_rh/2`). HPBL is also
  # absent; the scorer tolerates a nil HPBL without penalising the score.
  @surface_messages [
    %{var: "TMP", level: "2 m above ground"},
    %{var: "RH", level: "2 m above ground"},
    %{var: "PRES", 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"}
  ]

  # Pressure levels where GEFS pgrb2a publishes all of TMP, RH, and HGT
  # (the subset needed to build a temperature / dewpoint / height trace
  # the scorer can consume). RH is missing at 300/400/600 mb even though
  # HGT and UGRD/VGRD publish there — those levels are excluded.
  @profile_pressure_levels [1000, 925, 850, 700, 500, 250, 200, 100, 50, 10]

  @doc """
  Rounds a wall-clock `DateTime` down to the nearest GEFS run cycle
  (00/06/12/18Z) on the same UTC day.
  """
  @spec nearest_run(DateTime.t()) :: DateTime.t()
  def nearest_run(%DateTime{} = dt) do
    run_hour = dt.hour |> div(6) |> Kernel.*(6)

    dt
    |> DateTime.truncate(:second)
    |> Map.put(:hour, run_hour)
    |> Map.put(:minute, 0)
    |> Map.put(:second, 0)
  end

  @doc """
  Builds the NOMADS URL for the GEFS ensemble-mean (`geavg`) pgrb2a 0.5°
  GRIB2 file at the given `run_date`, `run_hour`, and `forecast_hour`.
  """
  @spec ensmean_url(Date.t(), non_neg_integer(), non_neg_integer()) :: String.t()
  def ensmean_url(%Date{} = run_date, run_hour, forecast_hour) when run_hour in @run_hours and forecast_hour >= 0 do
    date_str = Calendar.strftime(run_date, "%Y%m%d")
    hh = pad2(run_hour)
    fff = pad3(forecast_hour)

    "#{@base_url}/gefs.#{date_str}/#{hh}/atmos/pgrb2ap5/geavg.t#{hh}z.pgrb2a.0p50.f#{fff}"
  end

  @doc "Appends `.idx` to a GRIB2 URL to locate its wgrib2 byte-range index."
  @spec idx_url(String.t()) :: String.t()
  def idx_url(grib_url) when is_binary(grib_url), do: grib_url <> ".idx"

  @doc """
  Returns the list of GEFS forecast hours published per run: 3-hourly
  from f000 through f240, then 6-hourly from f246 through f384.
  """
  @spec forecast_hours() :: [non_neg_integer()]
  def forecast_hours do
    short = Enum.to_list(0..240//3)
    long = Enum.to_list(246..384//6)
    short ++ long
  end

  @doc """
  GRIB2 message descriptors the propagation scorer needs from the
  GEFS pgrb2a product. Values match the `var` and `level` strings that
  appear in the file's `.idx` sidecar exactly.
  """
  @spec surface_messages() :: [%{var: String.t(), level: String.t()}]
  def surface_messages, do: @surface_messages

  @doc """
  Derives dewpoint in °C from air temperature (°C) and relative humidity
  (%) using the Magnus formula with the Bolton (1980) coefficients.

  GEFS pgrb2a publishes 2 m RH but not 2 m DPT, so downstream code that
  consumed HRRR's `surface_dewpoint_c` must call this helper to fill in
  the equivalent value. RH is clamped to a small positive floor to keep
  the logarithm finite at 0 %.
  """
  @spec dewpoint_from_rh(float(), float()) :: float()
  def dewpoint_from_rh(temp_c, rh_pct) when is_number(temp_c) and is_number(rh_pct) do
    a = 17.625
    b = 243.04
    rh = max(rh_pct, 0.01) / 100.0
    gamma = :math.log(rh) + a * temp_c / (b + temp_c)
    b * gamma / (a - gamma)
  end

  @doc """
  Converts a raw `%{"VAR:LEVEL" => float}` cell map (as returned by the
  GRIB2 extractor) into a scorer-shaped profile map.

  Returns Celsius temperatures, millibar pressure, and a pressure-level
  profile list with dewpoint derived from relative humidity. Missing
  fields fall through as `nil` where the scorer can tolerate it.
  """
  @spec build_profile(%{String.t() => number()}) :: map()
  def build_profile(parsed) when is_map(parsed) do
    profile = build_pressure_profile(parsed)

    sfc_temp_c =
      case parsed["TMP:2 m above ground"] do
        k when is_number(k) -> k - 273.15
        _ -> lowest_profile_value(profile, "tmpc")
      end

    sfc_dewpoint_c =
      case {sfc_temp_c, parsed["RH:2 m above ground"]} do
        {t, rh} when is_number(t) and is_number(rh) -> dewpoint_from_rh(t, rh)
        _ -> lowest_profile_value(profile, "dwpc")
      end

    sfc_pressure_mb =
      case parsed["PRES:surface"] do
        pa when is_number(pa) -> pa / 100.0
        _ -> nil
      end

    %{
      surface_temp_c: sfc_temp_c,
      surface_dewpoint_c: sfc_dewpoint_c,
      surface_pressure_mb: sfc_pressure_mb,
      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

  defp build_pressure_profile(parsed) do
    Enum.flat_map(@profile_pressure_levels, fn level ->
      key = "#{level} mb"
      tmp_k = parsed["TMP:#{key}"]
      rh = parsed["RH:#{key}"]
      hgt = parsed["HGT:#{key}"]

      if is_number(tmp_k) and is_number(rh) and is_number(hgt) do
        tmpc = tmp_k - 273.15

        [
          %{
            "pres" => level * 1.0,
            "tmpc" => tmpc,
            "dwpc" => dewpoint_from_rh(tmpc, rh),
            "hght" => hgt
          }
        ]
      else
        []
      end
    end)
  end

  defp lowest_profile_value([], _key), do: nil
  defp lowest_profile_value([first | _], key), do: first[key]

  @doc """
  Fetches GEFS ensemble-mean (`geavg`) grid data for one forecast hour
  and returns a list of scorer-ready profile attribute maps, one per
  CONUS grid cell (0.125° matching `Microwaveprop.Propagation.Grid`).

  Uses NOMADS byte-range requests against the `.idx` sidecar to pull
  only the surface messages the scorer needs, then hands the slim
  GRIB2 binary to `wgrib2 -lola` for bilinear regridding from GEFS's
  native 0.5° lat/lon grid onto the CONUS 0.125° grid.

  Returns `{:error, :wgrib2_not_available}` if the `wgrib2` binary
  isn't on the host — all callers should degrade gracefully.
  """
  @spec fetch_grid_profiles(Date.t(), non_neg_integer(), non_neg_integer()) ::
          {:ok, [map()]} | {:error, term()}
  def fetch_grid_profiles(%Date{} = run_date, run_hour, forecast_hour) do
    Microwaveprop.Instrument.span(
      [:gefs, :fetch_grid_profiles],
      %{forecast_hour: forecast_hour},
      fn -> do_fetch_grid_profiles(run_date, run_hour, forecast_hour) end
    )
  end

  defp do_fetch_grid_profiles(run_date, run_hour, forecast_hour) do
    alias Microwaveprop.Propagation.Grid
    alias Microwaveprop.Weather.Grib2.Wgrib2
    alias Microwaveprop.Weather.HrrrClient

    url = ensmean_url(run_date, run_hour, forecast_hour)

    with {:ok, idx_text} <- fetch_idx(idx_url(url)),
         idx_entries = HrrrClient.parse_idx(idx_text),
         ranges = HrrrClient.byte_ranges_for_messages(idx_entries, @surface_messages),
         {:ok, grib_binary} <- download_grib_ranges(url, ranges),
         {:ok, grid} <- Wgrib2.extract_grid(grib_binary, wgrib2_match_pattern(), Grid.wgrib2_grid_spec()) do
      profiles =
        Enum.map(grid, fn {{lat, lon}, cell_values} ->
          cell_values
          |> build_profile()
          |> Map.merge(%{lat: lat, lon: lon})
        end)

      {:ok, profiles}
    end
  end

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

  defp fetch_idx(url) do
    case Req.get(url, req_options()) do
      {:ok, %{status: 200, body: body}} -> {:ok, body}
      {:ok, %{status: status}} -> {:error, "GEFS idx HTTP #{status}"}
      {:error, reason} -> {:error, reason}
    end
  end

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

  defp download_grib_ranges(url, ranges) do
    merged =
      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()

    supervisor = {:via, PartitionSupervisor, {Microwaveprop.TaskSupervisor, self()}}

    results =
      supervisor
      |> Task.Supervisor.async_stream_nolink(
        merged,
        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, "GEFS grib HTTP #{status}"}
            {:error, reason} -> {:error, reason}
          end
        end,
        max_concurrency: 8,
        timeout: 120_000
      )
      |> Enum.map(fn
        {:ok, result} ->
          result

        {:exit, reason} ->
          Logger.error("GEFS grib range download crashed: url=#{url} reason=#{inspect(reason)}")

          {:error, reason}
      end)

    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 req_options do
    defaults = [receive_timeout: 120_000, max_retries: 3]
    overrides = Application.get_env(:microwaveprop, :gefs_req_options, [])
    Keyword.merge(defaults, overrides)
  end

  defp pad2(n), do: n |> Integer.to_string() |> String.pad_leading(2, "0")
  defp pad3(n), do: n |> Integer.to_string() |> String.pad_leading(3, "0")
end