defmodule Microwaveprop.Weather.HrdpsClient do
  @moduledoc """
  Client for ECCC's HRDPS (High Resolution Deterministic Prediction System) —
  the Canadian analog to HRRR. Mirrors `Microwaveprop.Weather.HrrrClient`'s
  shape (`fetch_grid/3` returning `{:ok, %{{lat, lon} => profile}}`) so the
  propagation chain can call either polymorphically.

  ## Architectural differences from HRRR

  * **URL structure** — date-prefixed datamart paths
    (`https://dd.weather.gc.ca/{YYYYMMDD}/WXO-DD/model_hrdps/continental/2.5km/{HH}/{FFF}/...`).
    No flat `/model_hrdps/` root; the old structure 404s as of the 2026-04-29
    re-verification.
  * **One variable per file** — HRDPS publishes each variable as its own
    GRIB2 file (~3.5 MB each) instead of HRRR's bundled `wrfsfcf`/`wrfprsf`.
    No idx files, no byte-range partials. We fetch each needed variable's
    file concurrently and **byte-concatenate** them into a single multi-record
    binary — wgrib2 reads the result as if it were a normal multi-message
    GRIB2 file. No `wgrib2 -merge` step needed.
  * **Rotated lat/lon projection** — handled transparently by `wgrib2 -lon`
    inside `Microwaveprop.Weather.Grib2.Wgrib2.extract_points_from_file/3`.
  * **Cycle cadence** — 4×/day at 00/06/12/18Z (vs HRRR's hourly), forecasts
    out to 48h. Publish latency is ~3-4h after each cycle.
  """

  alias Microwaveprop.Weather.Grib2.Wgrib2

  require Logger

  @datamart_base_default "https://dd.weather.gc.ca"

  defp datamart_base, do: Application.get_env(:microwaveprop, :hrdps_datamart_base, @datamart_base_default)

  # Surface variables required by the propagation scorer. HRRR's equivalent
  # set is in `HrrrClient.@surface_messages`; HRDPS doesn't publish the
  # bundled APCP / PWAT inventory at f000, so the scorer's PWAT factor falls
  # back to its default for HRDPS-derived points (see plan stage 1.2).
  @surface_vars [
    {:tmp_2m, "TMP", "AGL-2m"},
    {:depr_2m, "DEPR", "AGL-2m"},
    {:dpt_2m, "DPT", "AGL-2m"},
    {:pres_sfc, "PRES", "Sfc"},
    {:hpbl_sfc, "HPBL", "Sfc"},
    {:ugrd_10m, "UGRD", "AGL-10m"},
    {:vgrd_10m, "VGRD", "AGL-10m"},
    {:tcdc_sfc, "TCDC", "Sfc"}
  ]

  # Pressure-level set covering the lower troposphere where the refractivity
  # gradient is computed. Matches HRRR's `@grid_pressure_levels` (1000→700 mb)
  # so SoundingParams.derive sees a comparable profile shape.
  @grid_pressure_levels [1000, 950, 900, 850, 800, 750, 700]

  @pressure_var_kinds [{:tmp, "TMP"}, {:depr, "DEPR"}, {:hgt, "HGT"}]

  @pressure_vars (for level <- @grid_pressure_levels,
                      {kind, msc} <- @pressure_var_kinds do
                    atom = String.to_atom("#{kind}_#{level}mb")
                    msc_level = "ISBL_#{level |> Integer.to_string() |> String.pad_leading(4, "0")}"
                    {atom, msc, msc_level}
                  end)

  @all_vars @surface_vars ++ @pressure_vars

  @var_lookup Map.new(@all_vars, fn {atom, msc_var, msc_level} -> {atom, {msc_var, msc_level}} end)

  # --- Public API ---

  @doc "List of all variable atoms this client knows how to fetch."
  @spec variables() :: [atom()]
  def variables, do: Enum.map(@all_vars, fn {atom, _, _} -> atom end)

  @doc """
  Build the MSC Datamart URL for a single HRDPS variable file.

  ECCC reorganized to a date-prefixed structure as of 2026; the old
  `/model_hrdps/...` flat root 404s. URLs always look like:

      https://dd.weather.gc.ca/{YYYYMMDD}/WXO-DD/model_hrdps/continental/2.5km/{HH}/{FFF}/
        {YYYYMMDD}T{HH}Z_MSC_HRDPS_{VAR}_{LEVEL}_RLatLon0.0225_PT{FFF}H.grib2
  """
  @spec grib_url(atom(), DateTime.t(), non_neg_integer()) :: String.t()
  def grib_url(variable, %DateTime{} = run_time, forecast_hour) do
    {msc_var, msc_level} =
      Map.get(@var_lookup, variable) ||
        raise ArgumentError, "unknown HRDPS variable #{inspect(variable)}"

    cycle = nearest_hrdps_cycle(run_time)
    date_str = Calendar.strftime(cycle, "%Y%m%d")
    hour_str = cycle.hour |> Integer.to_string() |> String.pad_leading(2, "0")
    fff_str = forecast_hour |> Integer.to_string() |> String.pad_leading(3, "0")

    filename =
      "#{date_str}T#{hour_str}Z_MSC_HRDPS_#{msc_var}_#{msc_level}_RLatLon0.0225_PT#{fff_str}H.grib2"

    "#{datamart_base()}/#{date_str}/WXO-DD/model_hrdps/continental/2.5km/#{hour_str}/#{fff_str}/#{filename}"
  end

  @doc """
  Snap a `DateTime` to the start of the most recent HRDPS cycle hour
  (00/06/12/18Z). Always rounds DOWN — never points at a future cycle that
  may not have published yet. Callers needing a later valid_time should
  reach it via `forecast_hour` instead.
  """
  @spec nearest_hrdps_cycle(DateTime.t()) :: DateTime.t()
  def nearest_hrdps_cycle(%DateTime{} = dt) do
    cycle_hour = div(dt.hour, 6) * 6
    seconds_into_cycle = (dt.hour - cycle_hour) * 3600 + dt.minute * 60 + dt.second
    DateTime.add(dt, -seconds_into_cycle, :second)
  end

  @doc """
  Returns true if the HRDPS cycle for `run_time` is published. Probes the
  cycle's f000 directory existence; the directory only appears once the
  first forecast hour's files are uploaded.

  Override the probe in tests via the `:microwaveprop, :hrdps_cycle_available_fn`
  Application env (a 1-arity function from `DateTime` to `boolean`).
  """
  @spec cycle_available?(DateTime.t()) :: boolean()
  def cycle_available?(run_time) do
    rounded = nearest_hrdps_cycle(run_time)

    case Application.get_env(:microwaveprop, :hrdps_cycle_available_fn) do
      fun when is_function(fun, 1) -> fun.(rounded)
      _ -> probe_cycle(rounded)
    end
  end

  defp probe_cycle(%DateTime{} = run_time) do
    date_str = Calendar.strftime(run_time, "%Y%m%d")
    hour_str = run_time.hour |> Integer.to_string() |> String.pad_leading(2, "0")
    url = "#{datamart_base()}/#{date_str}/WXO-DD/model_hrdps/continental/2.5km/#{hour_str}/000/"

    case Req.head(url, req_options()) do
      {:ok, %{status: 200}} -> true
      _ -> false
    end
  end

  @doc """
  Fetch HRDPS profiles for the given list of `{lat, lon}` points at the
  cycle's `forecast_hour`. Returns `{:ok, %{{lat, lon} => profile}}` matching
  the shape `HrrrClient.fetch_grid/3` produces, so downstream code is
  source-agnostic.

  Internally:

  1. Build URLs for every variable in `variables/0`.
  2. Concurrently download each file as a binary.
  3. Byte-concatenate into a single multi-record GRIB2.
  4. Stage to a temp file; call `Wgrib2.extract_points_from_file/3` with all
     points in one pass — wgrib2 handles the rotated-lat/lon reprojection.
  5. Build per-point profiles via `build_profile_from_extracted/1`.
  """
  @spec fetch_grid([{float(), float()}], DateTime.t(), keyword()) ::
          {:ok, %{{float(), float()} => map()}} | {:error, term()}
  def fetch_grid(points, run_time, opts \\ []) do
    Microwaveprop.Instrument.span(
      [:hrdps, :fetch_grid],
      %{point_count: length(points), forecast_hour: Keyword.get(opts, :forecast_hour, 0)},
      fn -> do_fetch_grid(points, run_time, opts) end
    )
  end

  defp do_fetch_grid(points, run_time, opts) do
    cycle = nearest_hrdps_cycle(run_time)
    forecast_hour = Keyword.get(opts, :forecast_hour, 0)

    with {:ok, binaries} <- fetch_all_variables(cycle, forecast_hour),
         combined = concat_grib_binaries(binaries),
         {:ok, path} <- write_temp_grib(combined),
         {:ok, extracted} <- extract_points(path, points) do
      _ = File.rm(path)

      profiles =
        Map.new(extracted, fn {point, raw} ->
          profile =
            raw
            |> build_profile_from_extracted()
            |> Map.put(:run_time, cycle)
            |> Map.put(:forecast_hour, forecast_hour)

          {point, profile}
        end)

      {:ok, profiles}
    end
  end

  @doc """
  Concatenate a list of GRIB2 binaries. wgrib2 reads multi-record GRIB2 files
  natively, so byte-concatenation is a valid stand-in for the `wgrib2 -merge`
  workflow — and skips the temp-file shuffle.
  """
  @spec concat_grib_binaries([binary()]) :: binary()
  def concat_grib_binaries(binaries), do: IO.iodata_to_binary(binaries)

  @doc """
  Translate the `Wgrib2.extract_points_from_file/3` output (`%{"VAR:LEVEL" => float}`)
  into the same profile shape `HrrrClient.build_profile/1` produces, so
  downstream code can consume HRDPS-derived rows interchangeably with
  HRRR-derived rows.

  HRDPS publishes DEPR (T-Td depression in K) as a primary surface variable
  rather than DPT. When DPT is absent this builder derives dewpoint as
  `temp - depression`.
  """
  @spec build_profile_from_extracted(%{String.t() => number()}) :: map()
  def build_profile_from_extracted(extracted) do
    sfc_temp_k = extracted["TMP:2 m above ground"]
    sfc_dpt_k = surface_dewpoint_k(extracted, sfc_temp_k)
    sfc_pres_pa = extracted["PRES:surface"]

    profile_levels = build_pressure_profile(extracted)
    lowest = List.first(profile_levels)

    sfc_temp_c =
      cond do
        sfc_temp_k -> sfc_temp_k - 273.15
        lowest -> lowest["tmpc"]
        true -> nil
      end

    sfc_dpt_c =
      cond do
        sfc_dpt_k -> sfc_dpt_k - 273.15
        lowest -> lowest["dwpc"]
        true -> nil
      end

    %{
      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: extracted["HPBL:surface"],
      pwat_mm: extracted["PWAT:entire atmosphere (considered as a single layer)"],
      wind_u: extracted["UGRD:10 m above ground"],
      wind_v: extracted["VGRD:10 m above ground"],
      cloud_cover_pct: extracted["TCDC:surface"] || extracted["TCDC:entire atmosphere"],
      precip_mm: extracted["APCP:surface"],
      profile: profile_levels
    }
  end

  defp surface_dewpoint_k(extracted, sfc_temp_k) do
    dpt = extracted["DPT:2 m above ground"]
    depr = extracted["DEPR:2 m above ground"]

    cond do
      not is_nil(dpt) -> dpt
      not is_nil(sfc_temp_k) and not is_nil(depr) -> sfc_temp_k - depr
      true -> nil
    end
  end

  defp build_pressure_profile(extracted) do
    Enum.flat_map(@grid_pressure_levels, fn level ->
      level_str = "#{level} mb"
      tmp = extracted["TMP:#{level_str}"]
      dpt = pressure_level_dewpoint(extracted, level_str, tmp)
      hgt = extracted["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)
  end

  defp pressure_level_dewpoint(extracted, level_str, tmp) do
    dpt = extracted["DPT:#{level_str}"]
    depr = extracted["DEPR:#{level_str}"]

    cond do
      not is_nil(dpt) -> dpt
      not is_nil(tmp) and not is_nil(depr) -> tmp - depr
      true -> nil
    end
  end

  # --- Internal: file fetch + concat ---

  # Fetch every variable's file concurrently. Datamart files are small
  # (~3.5 MB) and individually quick, but ~14 files done serially still adds
  # up to several seconds. 4-way concurrency stays well under any sensible
  # rate limit.
  defp fetch_all_variables(%DateTime{} = cycle, forecast_hour) do
    supervisor = {:via, PartitionSupervisor, {Microwaveprop.TaskSupervisor, self()}}

    results =
      supervisor
      |> Task.Supervisor.async_stream_nolink(
        variables(),
        fn variable -> fetch_variable(variable, cycle, forecast_hour) end,
        max_concurrency: 4,
        timeout: 120_000,
        ordered: true
      )
      |> Enum.map(fn
        {:ok, result} ->
          result

        {:exit, reason} ->
          # Per CLAUDE.md: never silently drop async failures. A swallowed
          # download here means a Canadian dead-zone we won't notice for
          # hours.
          Logger.error(
            "HRDPS variable fetch crashed: cycle=#{inspect(cycle)} f=#{forecast_hour} reason=#{inspect(reason)}"
          )

          {:error, reason}
      end)

    case Enum.find(results, &match?({:error, _}, &1)) do
      nil ->
        binaries = Enum.map(results, fn {:ok, _var, body} -> body end)
        {:ok, binaries}

      {:error, _} = error ->
        error
    end
  end

  defp fetch_variable(variable, cycle, forecast_hour) do
    url = grib_url(variable, cycle, forecast_hour)

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

      {:ok, %{status: status}} ->
        Logger.error("HRDPS fetch HTTP #{status} for variable=#{variable} url=#{url}")
        {:error, "HRDPS fetch HTTP #{status}"}

      {:error, reason} ->
        Logger.error("HRDPS fetch crashed: variable=#{variable} url=#{url} reason=#{inspect(reason)}")
        {:error, reason}
    end
  end

  defp write_temp_grib(binary) do
    path = Path.join(System.tmp_dir!(), "hrdps_#{System.unique_integer([:positive])}.grib2")

    case File.write(path, binary) do
      :ok -> {:ok, path}
      {:error, reason} -> {:error, {:tmp_write, reason}}
    end
  end

  defp extract_points(path, points) do
    Wgrib2.extract_points_from_file(path, ":(TMP|DPT|DEPR|PRES|HPBL|UGRD|VGRD|TCDC|HGT|APCP|PWAT):", points)
  end

  defp req_options do
    defaults = [receive_timeout: 120_000, retry: &retry?/2, max_retries: 5, retry_delay: &retry_delay/1]
    overrides = Application.get_env(:microwaveprop, :hrdps_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
