defmodule Microwaveprop.Weather.RtmaClient do
  @moduledoc """
  Client for NOAA RTMA (Real-Time Mesoscale Analysis) data.
  2.5 km resolution, 15-minute analysis cycles, CONUS coverage.

  Available on AWS S3 at s3://noaa-rtma-pds/. GRIB2 format with
  byte-range requests, same pattern as HRRR.

  RTMA provides surface fields only (no pressure levels):
  - 2m temperature, 2m dewpoint
  - 10m wind U/V
  - Surface pressure
  - Visibility
  - Precipitation analysis
  """

  alias Microwaveprop.Weather.Grib2.Extractor

  require Logger

  @s3_base "https://noaa-rtma-pds.s3.amazonaws.com"

  @wanted_fields [
    "TMP:2 m above ground",
    "DPT:2 m above ground",
    "PRES:surface",
    "UGRD:10 m above ground",
    "VGRD:10 m above ground",
    "VIS:surface"
  ]

  @doc """
  Fetch RTMA observation for a point at a specific time.
  Returns {:ok, attrs} or {:error, reason}.
  """
  @spec fetch_observation(float(), float(), DateTime.t()) :: {:ok, map()} | {:error, term()}
  def fetch_observation(lat, lon, timestamp) do
    Microwaveprop.Instrument.span(
      [:rtma, :fetch_observation],
      %{},
      fn -> do_fetch_observation(lat, lon, timestamp) end
    )
  end

  defp do_fetch_observation(lat, lon, timestamp) do
    # RTMA runs every hour with 15-min updates; round to nearest hour
    valid_time = %{DateTime.truncate(timestamp, :second) | minute: 0, second: 0}
    rlat = Float.round(lat * 40) / 40
    rlon = Float.round(lon * 40) / 40

    url = rtma_url(valid_time)
    idx_url = "#{url}.idx"

    with {:ok, idx_body} <- fetch_idx(idx_url),
         ranges = byte_ranges_for_fields(idx_body, @wanted_fields),
         {:ok, grib_data} <- download_ranges(url, ranges),
         {:ok, fields} <- extract_point(grib_data, rlat, rlon) do
      attrs = %{
        valid_time: valid_time,
        lat: rlat,
        lon: rlon,
        temp_c: kelvin_to_celsius(fields["TMP:2 m above ground"]),
        dewpoint_c: kelvin_to_celsius(fields["DPT:2 m above ground"]),
        pressure_mb: pa_to_mb(fields["PRES:surface"]),
        wind_u_ms: fields["UGRD:10 m above ground"],
        wind_v_ms: fields["VGRD:10 m above ground"],
        visibility_m: fields["VIS:surface"]
      }

      {:ok, attrs}
    end
  end

  @doc "Build S3 URL for an RTMA analysis file."
  @spec rtma_url(DateTime.t()) :: String.t()
  def rtma_url(valid_time) do
    date_str = Calendar.strftime(valid_time, "%Y%m%d")
    hour_str = valid_time.hour |> Integer.to_string() |> String.pad_leading(2, "0")
    "#{@s3_base}/rtma2p5.#{date_str}/rtma2p5.t#{hour_str}z.2dvaranl_ndfd.grb2_wexp"
  end

  defp fetch_idx(url) do
    case Req.get(url, [receive_timeout: 15_000] ++ req_options()) do
      {:ok, %{status: 200, body: body}} -> {:ok, body}
      {:ok, %{status: status}} -> {:error, "RTMA idx HTTP #{status}"}
      {:error, reason} -> {:error, "RTMA idx failed: #{inspect(reason)}"}
    end
  end

  defp byte_ranges_for_fields(idx_body, wanted_fields) do
    lines =
      idx_body
      |> String.split("\n", trim: true)
      |> Enum.map(fn line ->
        parts = String.split(line, ":")
        %{offset: String.to_integer(Enum.at(parts, 1, "0")), field: Enum.at(parts, 3, "") <> ":" <> Enum.at(parts, 4, "")}
      end)

    offsets = Enum.map(lines, & &1.offset)

    lines
    |> Enum.with_index()
    |> Enum.filter(fn {line, _i} ->
      Enum.any?(wanted_fields, &String.contains?(line.field, &1))
    end)
    |> Enum.map(fn {line, i} ->
      next_offset = Enum.at(offsets, i + 1, line.offset + 5_000_000)
      {line.offset, next_offset - 1}
    end)
    |> merge_ranges()
  end

  defp merge_ranges(ranges) do
    ranges
    |> Enum.sort()
    |> Enum.reduce([], fn
      range, [] -> [range]
      {s2, e2}, [{s1, e1} | rest] when s2 <= e1 + 1 -> [{s1, max(e1, e2)} | rest]
      range, acc -> [range | acc]
    end)
    |> Enum.reverse()
  end

  defp download_ranges(url, ranges) do
    data =
      ranges
      |> Task.async_stream(
        fn {range_start, range_end} ->
          Req.get(
            url,
            [
              headers: [{"Range", "bytes=#{range_start}-#{range_end}"}],
              receive_timeout: 30_000
            ] ++ req_options()
          )
        end,
        max_concurrency: 4,
        timeout: 60_000
      )
      |> Enum.reduce({:ok, []}, fn
        {:ok, {:ok, %{status: status, body: chunk}}}, {:ok, acc} when status in [200, 206] ->
          {:ok, [chunk | acc]}

        {:ok, {:ok, %{status: status}}}, _acc ->
          {:error, "RTMA download HTTP #{status}"}

        {:ok, {:error, reason}}, _acc ->
          {:error, "RTMA download failed: #{inspect(reason)}"}

        {:exit, reason}, _acc ->
          Logger.error("RtmaClient.download_ranges task crashed: url=#{url} reason=#{inspect(reason)}")
          {:error, "RTMA download crashed: #{inspect(reason)}"}

        other, acc ->
          Logger.warning("RtmaClient.download_ranges unexpected stream entry: #{inspect(other)}")
          acc
      end)

    case data do
      {:ok, chunks} -> {:ok, chunks |> Enum.reverse() |> IO.iodata_to_binary()}
      error -> error
    end
  end

  defp extract_point(grib_data, lat, lon) do
    case Extractor.extract_points(grib_data, lat, lon) do
      {:ok, fields} when map_size(fields) > 0 -> {:ok, fields}
      {:ok, _} -> {:error, "RTMA: no data at #{lat},#{lon}"}
      {:error, reason} -> {:error, reason}
    end
  end

  defp kelvin_to_celsius(nil), do: nil
  defp kelvin_to_celsius(k), do: k - 273.15

  defp pa_to_mb(nil), do: nil
  defp pa_to_mb(pa), do: pa / 100.0

  defp req_options, do: Application.get_env(:microwaveprop, :rtma_req_options, [])
end