defmodule Microwaveprop.Weather.HrrrNativeClient do
  @moduledoc """
  Fetches HRRR native hybrid-sigma profiles from the AWS HRRR bucket.

  This is the companion to `HrrrClient`, which works against the
  surface and 25 hPa pressure-level products. The native file
  (`wrfnatf00.grib2`) carries all variables on the 50 hybrid-sigma
  levels native to the HRRR model grid. Vertical spacing near the
  surface is ~10-50 m instead of the ~250 m the pressure-level
  product gives us — crucial for resolving the ducts and
  boundary-layer inversions discussed in
  `docs/plans/2026-04-09-propagation-modeling-improvements.md`.

  ## Design: batch, not per-point

  Each native-level HRRR file is ~566 MB. Essential variables (TMP,
  SPFH, HGT, UGRD, VGRD, TKE, PRES on all 50 hybrid levels) span
  ~530 MB of that file. Per-point on-demand fetching is not viable.
  Instead, the worker fetches the file once per `(date, hour)`,
  extracts native profiles for every point of interest in one pass,
  and bulk-inserts them.

  See `docs/research/hrrr_native_levels.md` for the full analysis.
  """

  alias Microwaveprop.Weather.HrrrClient

  @native_levels 1..50

  @native_variables ~w(TMP SPFH HGT UGRD VGRD TKE PRES)

  @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)

  @doc "Number of native hybrid-sigma levels in HRRR (currently 50)."
  @spec native_level_count() :: pos_integer()
  def native_level_count, do: Enum.count(@native_levels)

  @doc "The seven essential variables we extract on every native level."
  @spec native_variables() :: [String.t()]
  def native_variables, do: @native_variables

  @doc """
  The list of `%{var:, level:}` messages we extract from every native
  file. 7 vars × 50 levels = 350 messages, matching the spike in
  Task 1.1.
  """
  @spec native_messages() :: [%{var: String.t(), level: String.t()}]
  def native_messages do
    for level <- @native_levels, var <- @native_variables do
      %{var: var, level: "#{level} hybrid level"}
    end
  end

  @doc """
  Builds the AWS S3 URL for a native-level HRRR grib2 file.

  ## Examples

      iex> Microwaveprop.Weather.HrrrNativeClient.hrrr_native_url(~D[2026-04-09], 12)
      "https://noaa-hrrr-bdp-pds.s3.amazonaws.com/hrrr.20260409/conus/hrrr.t12z.wrfnatf00.grib2"
  """
  @spec hrrr_native_url(Date.t(), non_neg_integer(), non_neg_integer()) :: String.t()
  def hrrr_native_url(date, hour, 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")
    "#{hrrr_base()}/hrrr.#{date_str}/conus/hrrr.t#{hour_str}z.wrfnatf#{fh_str}.grib2"
  end

  @doc """
  Converts a parsed `%{"VAR:level" => value}` map into a
  `HrrrNativeProfile`-shaped map with level arrays sorted by ascending
  hybrid level (level 1 = surface).

  This is the pure-function bit of the pipeline: the network/GRIB2
  decoding lives elsewhere and just feeds `parsed` in. Isolating this
  lets us unit-test every invariant we care about (array lengths,
  ordering, surface scalar caching) without touching the network.
  """
  @spec build_native_profile(map()) :: map()
  def build_native_profile(parsed) when is_map(parsed) do
    levels =
      @native_levels
      |> Enum.map(fn level ->
        level_str = "#{level} hybrid level"

        %{
          level: level,
          hgt: parsed["HGT:#{level_str}"],
          tmp: parsed["TMP:#{level_str}"],
          spfh: parsed["SPFH:#{level_str}"],
          pres: parsed["PRES:#{level_str}"],
          ugrd: parsed["UGRD:#{level_str}"],
          vgrd: parsed["VGRD:#{level_str}"],
          tke: parsed["TKE:#{level_str}"]
        }
      end)
      |> Enum.reject(fn %{hgt: hgt, tmp: tmp} -> is_nil(hgt) or is_nil(tmp) end)
      |> Enum.sort_by(& &1.hgt)

    level_count = length(levels)

    %{
      level_count: level_count,
      heights_m: Enum.map(levels, & &1.hgt),
      temp_k: Enum.map(levels, & &1.tmp),
      spfh: Enum.map(levels, & &1.spfh),
      pressure_pa: Enum.map(levels, & &1.pres),
      u_wind_ms: Enum.map(levels, & &1.ugrd),
      v_wind_ms: Enum.map(levels, & &1.vgrd),
      tke_m2s2: Enum.map(levels, & &1.tke),
      surface_temp_k: parsed["TMP:surface"] || levels |> List.first() |> safe_get(:tmp),
      surface_spfh: parsed["SPFH:2 m above ground"] || levels |> List.first() |> safe_get(:spfh),
      surface_pressure_pa: parsed["PRES:surface"] || levels |> List.first() |> safe_get(:pres)
    }
  end

  defp safe_get(nil, _key), do: nil
  defp safe_get(map, key), do: Map.get(map, key)

  @doc """
  Returns the list of byte ranges to download for the essentials in
  one native HRRR file. Used by the (still-to-be-built) grid worker.
  Wraps `HrrrClient.byte_ranges_for_messages/2` with our native
  message list so callers don't have to know both.
  """
  @spec essential_byte_ranges([map()]) :: [{non_neg_integer(), non_neg_integer()}]
  def essential_byte_ranges(idx_entries) do
    HrrrClient.byte_ranges_for_messages(idx_entries, native_messages())
  end

  # The 4 variables needed for duct detection (N-profile + height).
  # Skips UGRD, VGRD, TKE (shear/Richardson were dead features).
  @duct_variables ~w(TMP SPFH HGT PRES)

  @doc """
  Messages needed for duct detection only (4 vars × 50 levels = 200 messages).
  ~300 MB download vs ~530 MB for all 7 variables.
  """
  @spec duct_messages() :: [%{var: String.t(), level: String.t()}]
  def duct_messages do
    for level <- @native_levels, var <- @duct_variables do
      %{var: var, level: "#{level} hybrid level"}
    end
  end

  @doc """
  Byte ranges for duct-detection variables only.
  """
  @spec duct_byte_ranges([map()]) :: [{non_neg_integer(), non_neg_integer()}]
  def duct_byte_ranges(idx_entries) do
    HrrrClient.byte_ranges_for_messages(idx_entries, duct_messages())
  end

  @doc """
  Fetch native duct metrics for the CONUS grid for one forecast hour.

  Downloads TMP, SPFH, HGT, PRES on all 50 hybrid levels, extracts to
  the CONUS grid via wgrib2, and computes duct metrics per cell using a
  cell-by-cell reducer (peak memory ~86 MB instead of ~1.8 GB).

  Returns `{:ok, %{{lat, lon} => duct_metrics}}` or `{:error, reason}`.
  """
  @spec fetch_native_duct_grid(Date.t(), non_neg_integer(), map(), non_neg_integer()) ::
          {:ok, %{{float(), float()} => map()}} | {:error, term()}
  def fetch_native_duct_grid(date, hour, grid_spec, forecast_hour \\ 0) do
    alias Microwaveprop.Weather.Grib2.Wgrib2

    url = hrrr_native_url(date, hour, forecast_hour)
    idx_url = url <> ".idx"

    tmp_grib = Path.join(System.tmp_dir!(), "hrrr_native_duct_#{System.unique_integer([:positive])}.grib2")

    try do
      with {:ok, idx_text} <- fetch_idx(idx_url),
           idx_entries = HrrrClient.parse_idx(idx_text),
           ranges = duct_byte_ranges(idx_entries),
           :ok <- HrrrClient.download_grib_ranges_to_file(url, ranges, tmp_grib) do
        match_pattern = ":(#{Enum.join(@duct_variables, "|")}):.*hybrid level:"

        Wgrib2.extract_grid_from_file_mapped(tmp_grib, match_pattern, grid_spec, &compute_duct_metrics/1)
      end
    after
      File.rm(tmp_grib)
    end
  end

  defp fetch_idx(url) do
    case http_get(url, receive_timeout: 120_000) do
      {:ok, %{status: 200, body: body}} -> {:ok, body}
      {:ok, %{status: status}} -> {:error, "HRRR native idx HTTP #{status}"}
      {:error, reason} -> {:error, reason}
    end
  end

  defp http_get(url, opts) do
    runner = Application.get_env(:microwaveprop, :hrrr_native_http_get, &Req.get/2)
    runner.(url, opts)
  end

  # Compute duct metrics from a single cell's native profile data.
  # Input: %{"TMP:1 hybrid level" => 297.5, "SPFH:1 hybrid level" => 0.012, ...}
  # Output: %{native_min_gradient: float, best_duct_freq_ghz: float, max_duct_thickness_m: float}
  defp compute_duct_metrics(parsed) do
    alias Microwaveprop.Propagation.Duct

    profile = build_native_profile(parsed)

    if profile.level_count >= 3 do
      duct_result = Duct.analyze(profile)

      ducts_summary =
        Enum.map(duct_result.ducts, fn d ->
          %{
            base_m: round(d.base_m),
            top_m: round(d.top_m),
            thickness_m: round(d.thickness_m),
            min_freq_ghz: Duct.min_trapped_frequency_ghz(d)
          }
        end)

      %{
        native_min_gradient: min_m_gradient(profile),
        best_duct_freq_ghz: duct_result.best_duct_band_ghz,
        max_duct_thickness_m: max_duct_thickness(duct_result.ducts),
        duct_count: length(duct_result.ducts),
        ducts: ducts_summary
      }
    else
      %{native_min_gradient: nil, best_duct_freq_ghz: nil, max_duct_thickness_m: nil, duct_count: 0, ducts: []}
    end
  end

  defp min_m_gradient(%{heights_m: heights, temp_k: temps, spfh: spfhs, pressure_pa: pressures}) do
    # Compute modified refractivity M at each level, find minimum dM/dh
    ms =
      [heights, temps, spfhs, pressures]
      |> Enum.zip()
      |> Enum.map(fn {h, t, q, p} ->
        # N = 77.6 * P/T + 3.73e5 * e/T^2, where e = q*P/(0.622 + 0.378*q)
        q = max(q || 0.0, 1.0e-8)
        e = q * p / (0.622 + 0.378 * q) / 100.0
        t_safe = max(t, 1.0)
        n = 77.6 * (p / 100.0) / t_safe + 3.73e5 * e / (t_safe * t_safe)
        m = n + 157.0 * h / 1000.0
        {h, m}
      end)

    ms
    |> Enum.chunk_every(2, 1, :discard)
    |> Enum.map(fn [{h1, m1}, {h2, m2}] ->
      dh = h2 - h1
      if dh > 0, do: (m2 - m1) / dh * 1000.0, else: 0.0
    end)
    |> Enum.min(fn -> 0.0 end)
  end

  defp max_duct_thickness([]), do: nil
  defp max_duct_thickness(ducts), do: ducts |> Enum.map(& &1.thickness_m) |> Enum.max()

  @doc """
  Extract native profiles for a list of `{lat, lon}` points from a
  GRIB2 file on disk, using wgrib2. Avoids loading the entire file
  into memory (~530 MB for native-level files).

  Returns `{:ok, %{{lat, lon} => native_profile_map}}` or `{:error, reason}`.
  """
  @spec extract_native_profiles_from_file(Path.t(), [{float(), float()}]) ::
          {:ok, %{{float(), float()} => map()}} | {:error, term()}
  def extract_native_profiles_from_file(grib_path, points) when is_list(points) do
    alias Microwaveprop.Weather.Grib2.Wgrib2

    match_pattern = ":(#{Enum.join(@native_variables, "|")}):.*hybrid level:"

    # Use direct point extraction (-lon) instead of grid extraction (-lola).
    # With geographically dispersed points, -lola creates a coast-to-coast
    # grid (~476k cells × 350 messages ≈ 665 MB), causing OOM.
    # -lon extracts only at the requested points with text output.
    case Wgrib2.extract_points_from_file(grib_path, match_pattern, points) do
      {:ok, point_data} ->
        {:ok, build_profiles_from_point_data(points, point_data)}

      error ->
        error
    end
  end

  @doc """
  Extract native profiles for a list of `{lat, lon}` points from a
  GRIB2 binary, using wgrib2 for speed.

  Under the hood this uses `-lola` on a bounding-box subgrid that
  covers all the requested points, then does nearest-neighbor lookup
  per point. Falls back to the pure-Elixir decoder if wgrib2 is not
  available (expect ~70s per point in that case).

  Returns `%{{lat, lon} => native_profile_map}` where each profile
  map has the shape expected by `HrrrNativeProfile.changeset/2`.
  """
  @spec extract_native_profiles(binary(), [{float(), float()}]) ::
          {:ok, %{{float(), float()} => map()}} | {:error, term()}
  def extract_native_profiles(grib_binary, points) when is_list(points) do
    alias Microwaveprop.Weather.Grib2.Wgrib2

    if Wgrib2.available?() do
      extract_native_profiles_wgrib2(grib_binary, points)
    else
      extract_native_profiles_elixir(grib_binary, points)
    end
  end

  defp build_profiles_from_point_data(points, point_data) do
    Map.new(points, fn point ->
      parsed = Map.get(point_data, point, %{})
      profile = if map_size(parsed) > 0, do: build_native_profile(parsed), else: %{level_count: 0}
      {point, profile}
    end)
  end

  defp extract_native_profiles_wgrib2(grib_binary, points) do
    # Write the binary to a temp file and use the point-extraction path.
    # Using -lola on the binary for geographically dispersed points creates
    # a coast-to-coast grid (~476k cells × 350 messages ≈ 665 MB), causing OOM.
    # -lon extracts only at the requested points with text output.
    tmp_path = Path.join(System.tmp_dir!(), "hrrr_#{System.unique_integer([:positive])}.grib2")

    try do
      File.write!(tmp_path, grib_binary)
      extract_native_profiles_from_file(tmp_path, points)
    after
      File.rm(tmp_path)
    end
  end

  defp extract_native_profiles_elixir(grib_binary, points) do
    alias Microwaveprop.Weather.Grib2.Extractor

    case Extractor.extract_grid(grib_binary, points) do
      {:ok, grid_data} ->
        result = Map.new(grid_data, fn {pt, parsed} -> {pt, build_native_profile(parsed)} end)
        {:ok, result}

      error ->
        error
    end
  end
end