defmodule Microwaveprop.Propagation do
  @moduledoc false

  import Ecto.Query

  alias Microwaveprop.Propagation.BandConfig
  alias Microwaveprop.Propagation.Grid
  alias Microwaveprop.Propagation.ProfilesFile
  alias Microwaveprop.Propagation.RunTiming
  alias Microwaveprop.Propagation.ScoreCache
  alias Microwaveprop.Propagation.Scorer
  alias Microwaveprop.Propagation.ScoresFile
  alias Microwaveprop.Repo
  alias Microwaveprop.Weather.SoundingParams

  require Logger

  @ml_key :propagation_ml
  @ml_module Microwaveprop.Propagation.Model

  @doc """
  Loads the ML model from disk, compiles the predict function, and caches both
  in persistent_term. No-op if the model file doesn't exist or ML deps unavailable.
  """
  @spec load_ml_model() :: :ok
  def load_ml_model do
    if Code.ensure_loaded?(@ml_module) do
      case @ml_module.load() do
        {:ok, params} ->
          predict_fn = @ml_module.compile_predict()
          :persistent_term.put(@ml_key, {predict_fn, params})
          Logger.info("PropagationML: model loaded and compiled")
          :ok

        :error ->
          Logger.info("PropagationML: no model file found, using algorithm scorer only")
          :ok
      end
    else
      Logger.info("PropagationML: ML dependencies not available")
      :ok
    end
  end

  @doc "Returns cached {predict_fn, params} tuple, or nil if not loaded."
  @spec ml_model() :: {function(), term()} | nil
  def ml_model do
    :persistent_term.get(@ml_key, nil)
  end

  @doc """
  Score a single grid point across all bands using HRRR profile data.
  Uses ML model if loaded, falls back to algorithm scorer.
  Returns a list of %{band_mhz, score, factors} maps.
  """
  @spec score_grid_point(map(), DateTime.t(), float(), float()) ::
          [%{band_mhz: non_neg_integer(), score: non_neg_integer(), factors: map()}]
  def score_grid_point(hrrr_profile, valid_time, latitude, longitude) do
    derived = derive_from_hrrr(hrrr_profile)

    temp_c = hrrr_profile.surface_temp_c
    dewpoint_c = hrrr_profile.surface_dewpoint_c

    # Skip points with missing or physically impossible surface data
    if is_nil(temp_c) or is_nil(dewpoint_c) or temp_c < -80 or temp_c > 60 or
         dewpoint_c < -80 or dewpoint_c > 50 do
      []
    else
      score_grid_point_with_data(hrrr_profile, valid_time, temp_c, dewpoint_c, derived, latitude, longitude)
    end
  end

  defp score_grid_point_with_data(hrrr_profile, valid_time, temp_c, dewpoint_c, derived, latitude, longitude) do
    # Algorithm is the primary scorer — always used for the map score.
    # ML score stored in factors as :ml_score for comparison/analysis.
    score_with_algorithm(hrrr_profile, valid_time, temp_c, dewpoint_c, derived, latitude, longitude)
  end

  defp score_with_algorithm(hrrr_profile, valid_time, temp_c, dewpoint_c, derived, _latitude, longitude) do
    temp_f = Scorer.c_to_f(temp_c)
    dewpoint_f = Scorer.c_to_f(dewpoint_c)

    conditions = %{
      abs_humidity: Scorer.absolute_humidity(temp_c, dewpoint_c),
      temp_f: temp_f,
      dewpoint_f: dewpoint_f,
      wind_speed_kts: Scorer.wind_speed_kts(hrrr_profile[:wind_u], hrrr_profile[:wind_v]),
      sky_cover_pct: hrrr_profile[:cloud_cover_pct],
      utc_hour: valid_time.hour,
      utc_minute: valid_time.minute,
      month: valid_time.month,
      longitude: longitude,
      pressure_mb: hrrr_profile.surface_pressure_mb,
      prev_pressure_mb: nil,
      rain_rate_mmhr: merged_rain_rate(hrrr_profile),
      min_refractivity_gradient: hrrr_profile[:native_min_gradient] || derived[:min_refractivity_gradient],
      bl_depth_m: hrrr_profile[:hpbl_m],
      pwat_mm: hrrr_profile[:pwat_mm],
      best_duct_band_ghz: hrrr_profile[:best_duct_freq_ghz] || hrrr_profile[:best_duct_band_ghz],
      bulk_richardson: hrrr_profile[:bulk_richardson]
    }

    # Hoist the four band-invariant factors out of the 17-band inner
    # loop. time_of_day / sky / wind / pressure depend on conditions
    # alone, not the band — precomputing once per point drops ~30% of
    # the scoring wall time on the hourly chain.
    conditions = Map.merge(conditions, Scorer.precompute_band_invariants(conditions))

    duct_info =
      if hrrr_profile[:duct_count] && hrrr_profile[:duct_count] > 0 do
        %{
          duct_count: hrrr_profile[:duct_count],
          best_duct_freq_ghz: hrrr_profile[:best_duct_freq_ghz],
          max_duct_thickness_m: hrrr_profile[:max_duct_thickness_m],
          ducts: hrrr_profile[:ducts] || []
        }
      end

    link_degradation = hrrr_profile[:commercial_link_degradation]

    Enum.map(BandConfig.all_bands(), fn band_config ->
      result = Scorer.composite_score(conditions, band_config)

      boosted_score =
        if link_degradation do
          Scorer.commercial_link_boost(result.score, link_degradation)
        else
          result.score
        end

      result =
        result
        |> Map.put(:score, boosted_score)
        |> Map.put(:band_mhz, band_config.freq_mhz)

      result =
        if link_degradation do
          put_in(result, [:factors, :commercial_link_degradation], link_degradation)
        else
          result
        end

      if duct_info, do: put_in(result, [:factors, :duct_info], duct_info), else: result
    end)
  end

  # Pick the heavier of HRRR's hourly accumulation-derived rate and NEXRAD's
  # reflectivity-derived rate. NEXRAD catches fast-moving convective cells that
  # fall between HRRR hourly analyses; HRRR catches broad stratiform rain that
  # NEXRAD reports as low dBZ. Taking max lets either source trigger the rain
  # penalty without double-counting.
  defp merged_rain_rate(hrrr_profile) do
    hrrr_rate = Scorer.precip_to_rate_mmhr(hrrr_profile[:precip_mm])
    nexrad_rate = Scorer.dbz_to_rain_rate_mmhr(hrrr_profile[:nexrad_max_reflectivity_dbz])
    max(hrrr_rate, nexrad_rate)
  end

  @doc """
  Replace every propagation score for `valid_time` with `scores`.

  Used by `PropagationGridWorker` on the hot path. Scores are written
  as binary files on disk via `ScoresFile.write!/3`, one file per
  band.

  Consumes `scores` in a single streaming pass that folds each score
  straight into a per-band accumulator. Previously this function ran
  `Enum.to_list/1` followed by `Enum.group_by/2`, which held two full
  copies of the ~460k-entry grid (list + grouped list) in memory at
  once — the hot path's largest transient spike after native-duct
  merge. The single-pass reduce keeps only one copy and buys back
  ~100 MB of headroom per forecast-hour step.
  """
  @spec replace_scores(Enumerable.t(), DateTime.t()) :: {:ok, non_neg_integer()} | {:error, term()}
  def replace_scores(scores, %DateTime{} = valid_time) do
    Microwaveprop.Instrument.span(
      [:db, :replace_scores],
      %{valid_time: valid_time},
      fn -> do_replace_scores(scores, valid_time) end
    )
  end

  defp do_replace_scores(scores, valid_time) do
    {per_band, total} =
      Enum.reduce(scores, {%{}, 0}, fn score, {acc, count} ->
        {Map.update(acc, score.band_mhz, [score], &[score | &1]), count + 1}
      end)

    Enum.each(per_band, fn {band_mhz, band_scores} ->
      try do
        ScoresFile.write!(band_mhz, valid_time, band_scores)
      rescue
        e ->
          Logger.warning("Propagation: ScoresFile write failed for band=#{band_mhz} vt=#{valid_time}: #{inspect(e)}")
      end
    end)

    {:ok, total}
  end

  @doc """
  Remove score files with valid_times older than 2 hours. Called on
  a cron by `Microwaveprop.Workers.PropagationPruneWorker`.
  """
  @spec prune_old_scores() :: :ok
  def prune_old_scores do
    cutoff = DateTime.add(DateTime.utc_now(), -2, :hour)
    file_deleted = ScoresFile.prune_older_than(cutoff)
    profiles_deleted = ProfilesFile.prune_older_than(cutoff)

    if file_deleted + profiles_deleted > 0 do
      Logger.info(
        "PropagationScores: pruned #{file_deleted} old score files + " <>
          "#{profiles_deleted} profile files (before #{cutoff})"
      )
    end

    :ok
  end

  @doc """
  Returns distinct valid_times for a band, ordered ascending. Always
  reads from the on-disk `ScoresFile` store — the `ScoreCache` only
  holds whatever hours have been fetched or broadcast, which can be a
  partial view of what's actually on disk, so using it as the source
  of truth for the timeline makes new forecast hours invisible until
  the cache happens to catch up. Filters out times more than 1 hour in
  the past, but always includes the most recent valid_time so there's
  always data to display.
  """
  # HRRR forecast horizon: f00..f18 covers the next 18 hours from cycle
  # time. Anything beyond that in the score store is a leftover from a
  # stale cycle and clutters the timeline without adding information.
  @hrrr_forecast_horizon_hours 18

  @spec available_valid_times(non_neg_integer()) :: [DateTime.t()]
  def available_valid_times(band_mhz) do
    now = DateTime.utc_now()
    past_cutoff = DateTime.add(now, -3600, :second)
    future_cutoff = DateTime.add(now, @hrrr_forecast_horizon_hours * 3600, :second)

    case ScoresFile.list_valid_times(band_mhz) do
      [] -> []
      times -> filter_or_latest(times, past_cutoff, future_cutoff)
    end
  end

  defp filter_or_latest(times, past_cutoff, future_cutoff) do
    fresh =
      Enum.filter(times, fn t ->
        DateTime.compare(t, past_cutoff) != :lt and DateTime.compare(t, future_cutoff) != :gt
      end)

    if fresh == [] do
      [Enum.max(times, DateTime)]
    else
      fresh
    end
  end

  @doc """
  Get scores for a band at a specific valid_time, optionally within a bounding box.
  If valid_time is nil, uses the earliest available (current analysis hour).
  Excludes factors for performance.
  """
  @spec scores_at(non_neg_integer(), DateTime.t() | nil, map() | nil) ::
          [%{lat: float(), lon: float(), score: non_neg_integer(), valid_time: DateTime.t()}]
  def scores_at(band_mhz, valid_time, bounds \\ nil) do
    Microwaveprop.Instrument.span([:propagation, :scores_at], %{band_mhz: band_mhz}, fn ->
      time = valid_time || earliest_valid_time(band_mhz)

      case time do
        nil -> []
        _ -> scores_at_fetch(band_mhz, time, bounds)
      end
    end)
  end

  defp scores_at_fetch(band_mhz, time, bounds) do
    case ScoreCache.fetch_bounds(band_mhz, time, bounds) do
      {:ok, scores} ->
        :telemetry.execute([:microwaveprop, :propagation, :scores_at, :cache], %{}, %{hit: true})
        Enum.map(scores, &Map.put(&1, :valid_time, time))

      :miss ->
        :telemetry.execute([:microwaveprop, :propagation, :scores_at, :cache], %{}, %{hit: false})
        read_from_disk_and_cache(band_mhz, time, bounds)
    end
  end

  @doc """
  Variant of `scores_at/3` that always reads from the `.ntms` file on
  disk and overwrites the cache entry, rather than returning whatever
  the cache happens to hold. Use from update paths (the map's
  `propagation_updated` handler) where the underlying file has just
  been rewritten but the cache may still contain the previous chain's
  scores because of the race between `propagation:cache` fan-out and
  `propagation:updated` delivery.
  """
  @spec scores_at_fresh(non_neg_integer(), DateTime.t(), map() | nil) ::
          [%{lat: float(), lon: float(), score: non_neg_integer(), valid_time: DateTime.t()}]
  def scores_at_fresh(band_mhz, %DateTime{} = valid_time, bounds \\ nil) do
    read_from_disk_and_cache(band_mhz, valid_time, bounds)
  end

  defp read_from_disk_and_cache(band_mhz, time, bounds) do
    full = ScoresFile.read_bounds(band_mhz, time)
    ScoreCache.put(band_mhz, time, full)

    full
    |> filter_bounds(bounds)
    |> Enum.map(&Map.put(&1, :valid_time, time))
  end

  @doc """
  Load the full CONUS score set for `{band_mhz, valid_time}` from the
  on-disk binary file and broadcast it to every `ScoreCache` in the
  cluster. Called from `PropagationGridWorker` after each forecast
  hour so all pods have a warm cache by the time clients begin
  requesting the new hour.
  """
  @spec warm_cache_and_broadcast(non_neg_integer(), DateTime.t()) :: :ok
  def warm_cache_and_broadcast(band_mhz, valid_time) do
    scores = ScoresFile.read_bounds(band_mhz, valid_time)
    ScoreCache.broadcast_put(band_mhz, valid_time, scores)
    :ok
  end

  defp filter_bounds(scores, nil), do: scores

  defp filter_bounds(scores, %{"south" => s, "north" => n, "west" => w, "east" => e}) do
    Enum.filter(scores, fn %{lat: lat, lon: lon} ->
      lat >= s and lat <= n and lon >= w and lon <= e
    end)
  end

  @doc "Get the latest scores for a band (alias for scores_at with earliest valid_time)."
  @spec latest_scores(non_neg_integer(), map() | nil) ::
          [%{lat: float(), lon: float(), score: non_neg_integer(), valid_time: DateTime.t()}]
  def latest_scores(band_mhz, bounds \\ nil) do
    scores_at(band_mhz, nil, bounds)
  end

  defp earliest_valid_time(band_mhz) do
    case ScoresFile.list_valid_times(band_mhz) do
      [earliest | _] -> earliest
      [] -> nil
    end
  end

  @doc "Get scores across all forecast hours for a single grid point (for sparkline)."
  @spec point_forecast(non_neg_integer(), float(), float()) ::
          [%{valid_time: DateTime.t(), score: non_neg_integer()}]
  def point_forecast(band_mhz, lat, lon) do
    Microwaveprop.Instrument.span([:propagation, :point_forecast], %{band_mhz: band_mhz}, fn ->
      {snapped_lat, snapped_lon} = snap_to_grid(lat, lon)
      now = DateTime.utc_now()

      # Use the on-disk .ntms file list as the authoritative timeline so
      # the chart never falls behind the main-map timeline (which also
      # reads the disk). The cache is still consulted per-hour for a
      # fast score lookup; a miss falls through to the file.
      band_mhz
      |> ScoresFile.list_valid_times()
      |> forecast_window(now)
      |> Enum.map(&point_forecast_entry(band_mhz, &1, snapped_lat, snapped_lon))
      |> Enum.reject(&is_nil/1)
    end)
  end

  defp point_forecast_entry(band_mhz, valid_time, lat, lon) do
    case ScoreCache.fetch_point(band_mhz, valid_time, lat, lon) do
      {:ok, score} ->
        %{valid_time: valid_time, score: score}

      :miss ->
        case ScoresFile.read_point(band_mhz, valid_time, lat, lon) do
          nil -> nil
          score -> %{valid_time: valid_time, score: score}
        end
    end
  end

  # Select the set of valid_times the forecast chart should render.
  # Mirrors `available_valid_times`: keep everything from one hour
  # before now onward so the most recent analysis hour (typically
  # ~30–60 min behind wall clock due to HRRR publishing lag) sits at
  # the left edge of the chart as "now". When every hour on disk is
  # older than that cutoff, fall back to just the newest entry so the
  # chart can still render a single data point.
  defp forecast_window([], _now), do: []

  defp forecast_window(times, now) do
    past_cutoff = DateTime.add(now, -3600, :second)
    future_cutoff = DateTime.add(now, @hrrr_forecast_horizon_hours * 3600, :second)
    filter_or_latest(times, past_cutoff, future_cutoff)
  end

  defp snap_to_grid(lat, lon) do
    step = Grid.step()
    {Float.round(Float.round(lat / step) * step, 3), Float.round(Float.round(lon / step) * step, 3)}
  end

  @doc "Get the full score and factors for a specific grid point, snapped to nearest grid."
  @spec point_detail(non_neg_integer(), float(), float(), DateTime.t() | nil) ::
          %{
            lat: float(),
            lon: float(),
            score: non_neg_integer(),
            factors: map(),
            valid_time: DateTime.t()
          }
          | nil
  def point_detail(band_mhz, lat, lon, valid_time \\ nil) do
    Microwaveprop.Instrument.span([:propagation, :point_detail], %{band_mhz: band_mhz}, fn ->
      do_point_detail(band_mhz, lat, lon, valid_time)
    end)
  end

  defp do_point_detail(band_mhz, lat, lon, valid_time) do
    {snapped_lat, snapped_lon} = snap_to_grid(lat, lon)
    time = valid_time || latest_valid_time(band_mhz)

    case time do
      nil ->
        nil

      _ ->
        case ScoresFile.read_point(band_mhz, time, snapped_lat, snapped_lon) do
          nil ->
            nil

          score ->
            %{
              lat: snapped_lat,
              lon: snapped_lon,
              score: score,
              factors: factors_for(band_mhz, time, snapped_lat, snapped_lon),
              valid_time: time
            }
        end
    end
  end

  # Rebuild the factor breakdown for a clicked grid cell by rescoring
  # the persisted HRRR profile. Only analysis hours (f00) persist
  # profiles, so forecast hours fall back to the most recent analysis
  # profile at or before the requested time — the atmospheric breakdown
  # still explains what's driving the score even if sampled an hour or
  # two earlier.
  defp factors_for(band_mhz, valid_time, lat, lon) do
    case ProfilesFile.read_point(valid_time, lat, lon) do
      nil -> factors_from_fallback_profile(band_mhz, valid_time, lat, lon)
      profile -> factors_from_profile(band_mhz, valid_time, profile, lat, lon)
    end
  end

  defp factors_from_fallback_profile(band_mhz, valid_time, lat, lon) do
    case latest_profile_time_at_or_before(valid_time) do
      nil ->
        %{}

      fallback_time ->
        case ProfilesFile.read_point(fallback_time, lat, lon) do
          nil -> %{}
          profile -> factors_from_profile(band_mhz, fallback_time, profile, lat, lon)
        end
    end
  end

  defp factors_from_profile(band_mhz, valid_time, profile, lat, lon) do
    profile
    |> score_grid_point(valid_time, lat, lon)
    |> Enum.find(fn r -> r.band_mhz == band_mhz end)
    |> case do
      %{factors: factors} when is_map(factors) -> factors
      _ -> %{}
    end
  end

  defp latest_profile_time_at_or_before(%DateTime{} = valid_time) do
    ProfilesFile.list_valid_times()
    |> Enum.filter(&(DateTime.compare(&1, valid_time) != :gt))
    |> case do
      [] -> nil
      past -> Enum.max(past, DateTime)
    end
  end

  @doc "Get the latest valid_time across all bands."
  @spec latest_valid_time() :: DateTime.t() | nil
  def latest_valid_time do
    ScoresFile.latest_valid_time()
  end

  @doc "Get the latest valid_time for a specific band."
  @spec latest_valid_time(non_neg_integer()) :: DateTime.t() | nil
  def latest_valid_time(band_mhz) do
    case ScoresFile.list_valid_times(band_mhz) do
      [] -> nil
      times -> Enum.max(times, DateTime)
    end
  end

  ## Run timings

  @doc """
  Record wall-clock duration for a single forecast-hour chain step.

  Called by `PropagationGridWorker` at the end of every step (success or
  failure) so the timing history survives pod restarts and can be
  inspected later to see which steps are slow or flaky.
  """
  @spec record_run_timing(map()) ::
          {:ok, RunTiming.t()} | {:error, Ecto.Changeset.t()}
  def record_run_timing(attrs) do
    %RunTiming{}
    |> RunTiming.changeset(attrs)
    |> Repo.insert()
  end

  @doc """
  List the most-recently-started run-timing rows, newest first.

  Defaults to 100 rows; pass `:limit` to override.
  """
  @spec list_recent_run_timings(keyword()) :: [RunTiming.t()]
  def list_recent_run_timings(opts \\ []) do
    limit = Keyword.get(opts, :limit, 100)

    RunTiming
    |> order_by(desc: :started_at)
    |> limit(^limit)
    |> Repo.all()
  end

  # Prefer the persisted scalar — `hrrr_profiles` already stored this at
  # ingestion time and AsosAdjustmentWorker loads 92k rows per tick without
  # the JSONB `profile` column to avoid a Jason.decode! storm on the DB pool.
  defp derive_from_hrrr(%{min_refractivity_gradient: grad}) when is_number(grad) do
    %{min_refractivity_gradient: grad * 1.0}
  end

  defp derive_from_hrrr(%{profile: profile}) when is_list(profile) and length(profile) >= 3 do
    case SoundingParams.derive(profile) do
      nil -> %{}
      derived -> %{min_refractivity_gradient: derived.min_refractivity_gradient}
    end
  end

  defp derive_from_hrrr(_), do: %{}
end