defmodule Microwaveprop.Weather.MrmsClient do
  @moduledoc """
  Fetch NOAA MRMS `PrecipRate` surface rain-rate mosaics and interpolate
  them onto the same 0.125° propagation grid we score against. MRMS is
  ~1 km resolution and updates every 2 minutes, so between hourly HRRR runs
  it's the freshest rain signal we can get over CONUS.

  The module is pure transport + decode: fetch, gunzip, run wgrib2 to
  regrid, return `%{{lat, lon} => rain_rate_mmhr}`. Caching and cron
  scheduling live in `Microwaveprop.Weather.MrmsCache` and
  `Microwaveprop.Workers.MrmsFetchWorker`.
  """

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

  require Logger

  @base_url "https://mrms.ncep.noaa.gov/2D/PrecipRate/"
  @match_pattern ":PrecipRate:"
  @missing_sentinel -3.0

  @type rain_grid :: %{{float(), float()} => float()}
  @type listing_entry :: %{filename: String.t(), valid_time: DateTime.t()}

  @doc """
  List the most recent MRMS PrecipRate files published on the NCEP index.
  Returns entries sorted newest-first so callers can just take the head.
  """
  @spec list_latest(non_neg_integer()) :: {:ok, [listing_entry()]} | {:error, term()}
  def list_latest(limit \\ 10) do
    case Req.get(@base_url, req_options()) do
      {:ok, %{status: 200, body: body}} ->
        entries =
          body
          |> parse_listing()
          |> Enum.sort_by(& &1.valid_time, {:desc, DateTime})
          |> Enum.take(limit)

        {:ok, entries}

      {:ok, %{status: status}} ->
        {:error, {:http_status, status}}

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

  @doc """
  Fetch a single MRMS file by filename, decompress, regrid to the 0.125°
  propagation grid, and return `{valid_time, %{{lat, lon} => rain_rate_mmhr}}`.

  Missing-data sentinels (-3) are dropped from the returned grid so callers
  can test membership directly.
  """
  @spec fetch_and_decode(String.t()) :: {:ok, DateTime.t(), rain_grid()} | {:error, term()}
  def fetch_and_decode(filename) do
    with {:ok, valid_time} <- parse_filename(filename),
         {:ok, grib_bytes} <- download(filename),
         {:ok, tmp_path} <- write_temp(grib_bytes),
         {:ok, grid} <- regrid(tmp_path) do
      File.rm(tmp_path)
      {:ok, valid_time, grid}
    end
  end

  @doc """
  Convenience: list the latest files, grab the newest that isn't already
  cached, and return the decoded grid. `known_valid_time` is whatever the
  caller already has — if the latest file matches it we short-circuit.
  """
  @spec fetch_latest(DateTime.t() | nil) ::
          {:ok, DateTime.t(), rain_grid()} | {:up_to_date, DateTime.t()} | {:error, term()}
  def fetch_latest(known_valid_time \\ nil) do
    case list_latest(1) do
      {:ok, [latest | _]} ->
        if known_valid_time && DateTime.compare(latest.valid_time, known_valid_time) != :gt do
          {:up_to_date, known_valid_time}
        else
          fetch_and_decode(latest.filename)
        end

      {:ok, []} ->
        {:error, :no_files_listed}

      other ->
        other
    end
  end

  # -- internals (public-but-undocumented for tests) --------------------

  @doc false
  def parse_listing(html) do
    ~r/MRMS_PrecipRate_00\.00_(\d{8}-\d{6})\.grib2\.gz/
    |> Regex.scan(html)
    |> Enum.uniq()
    |> Enum.flat_map(fn [full, stamp] ->
      case parse_stamp(stamp) do
        {:ok, dt} -> [%{filename: full, valid_time: dt}]
        _ -> []
      end
    end)
  end

  defp parse_filename(filename) do
    case Regex.run(~r/MRMS_PrecipRate_00\.00_(\d{8}-\d{6})\.grib2\.gz/, filename) do
      [_, stamp] -> parse_stamp(stamp)
      _ -> {:error, :bad_filename}
    end
  end

  defp parse_stamp(stamp) do
    with <<y::binary-4, m::binary-2, d::binary-2, "-", h::binary-2, mi::binary-2, s::binary-2>> <-
           stamp,
         {:ok, naive} <-
           NaiveDateTime.new(si(y), si(m), si(d), si(h), si(mi), si(s)) do
      {:ok, DateTime.from_naive!(naive, "Etc/UTC")}
    else
      _ -> {:error, :bad_stamp}
    end
  end

  defp si(s), do: String.to_integer(s)

  defp download(filename) do
    # Req auto-decompresses gzip responses, so the body we receive is
    # already the raw GRIB2 binary (starts with "GRIB" magic) even though
    # the URL ends in .gz. No gunzip step needed.
    url = @base_url <> filename

    case Req.get(url, req_options()) do
      {:ok, %{status: 200, body: body}} when is_binary(body) ->
        {:ok, body}

      {:ok, %{status: status}} ->
        {:error, {:http_status, status}}

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

  defp write_temp(grib_bytes) do
    path = Path.join(System.tmp_dir!(), "mrms_preciprate_#{System.unique_integer([:positive])}.grib2")
    File.write!(path, grib_bytes)
    {:ok, path}
  end

  defp regrid(tmp_path) do
    case Wgrib2.extract_grid_from_file(tmp_path, @match_pattern, Grid.wgrib2_grid_spec()) do
      {:ok, cells} -> {:ok, flatten_rain_grid(cells)}
      error -> error
    end
  end

  @doc false
  def flatten_rain_grid(cells) do
    Enum.reduce(cells, %{}, fn {{lat, lon}, fields}, acc ->
      case extract_preciprate(fields) do
        nil -> acc
        rate -> Map.put(acc, {lat, lon}, rate)
      end
    end)
  end

  defp extract_preciprate(fields) do
    Enum.find_value(fields, fn
      {"PrecipRate" <> _, value} when is_number(value) and value > @missing_sentinel and value >= 0 ->
        value

      _ ->
        nil
    end)
  end

  defp req_options do
    [receive_timeout: 120_000, retry: :safe_transient, max_retries: 3]
  end
end