defmodule Microwaveprop.Pskr.CalibrationSampler do
  @moduledoc """
  Builds `Pskr.CalibrationSample` rows by joining PSKR spot density
  with HRRR atmospheric state and SWPC geomagnetic state at a given
  hour.

  Workflow:

  1. Pull every `pskr_spots_hourly` row for the target hour.
  2. Snap each row's midpoint to the propagation grid (0.125°) and
     bucket by `(band, snapped_lat, snapped_lon)`.
  3. For each bucket, aggregate the spot stats (sum count, count
     distinct paths, dedup modes, median distance).
  4. For each bucket, look up the nearest `hrrr_profiles` row to the
     midpoint at the target hour and copy the surface/refractivity
     fields onto the sample. HRRR is pulled in one window query and
     joined in-memory to keep the per-bucket cost a hashtable lookup.
  5. Look up SWPC Kp once per run.
  6. Bulk-upsert the result with `(hour_utc, band, midpoint_lat,
     midpoint_lon)` as the conflict target.

  The sampler is intentionally idempotent: re-running for the same
  hour is a no-op except for the `updated_at` bump, so the worker
  can safely overlap with manual reruns and the cron retry on
  failure.
  """

  import Ecto.Query

  alias Microwaveprop.Pskr.CalibrationSample
  alias Microwaveprop.Pskr.SpotHourly
  alias Microwaveprop.Repo
  alias Microwaveprop.SpaceWeather
  alias Microwaveprop.Weather
  alias Microwaveprop.Weather.HrrrPointEnqueuer
  alias Microwaveprop.Weather.HrrrProfile

  require Logger

  @grid_step_deg 0.125
  # HRRR domain is ~3 km native, so a ±0.07° (~7 km) lat/lon window
  # plus ±60 min reliably finds at least one profile for any CONUS
  # midpoint at most hours.
  @hrrr_lookahead_deg 0.07
  @hrrr_lookahead_seconds 3600

  @typep cell_key :: {band :: String.t(), lat :: float(), lon :: float()}

  @typedoc """
  Result of a single `build_for_hour/1` pass:

    * `upserted` — number of `pskr_calibration_samples` rows
      inserted-or-updated.
    * `missing_hrrr_cells` — number of distinct cell midpoints whose
      HRRR profile was unavailable. The caller (typically
      `Microwaveprop.Workers.PskrCalibrationWorker`) uses this to
      decide whether to schedule a follow-up sampler pass after the
      Rust hrrr-point-worker has had time to drain the just-enqueued
      fetch tasks.
  """
  @type result :: %{upserted: non_neg_integer(), missing_hrrr_cells: non_neg_integer()}

  @doc """
  Build calibration samples for one hour. Returns
  `%{upserted: count, missing_hrrr_cells: count}`. Both counts are 0
  when no PSKR spots were captured for that hour — for example,
  during PSKR client downtime.

  The `missing_hrrr_cells` count tells the caller whether the Rust
  hrrr-point-worker has work to drain after this fire. The
  `PskrCalibrationWorker` uses it to decide whether to schedule a
  follow-up sampler pass for the same hour ~10 min later, so the
  newly-landed HRRR data gets UPSERTed onto the samples we just
  wrote with NULL fields.
  """
  @spec build_for_hour(DateTime.t()) :: result()
  def build_for_hour(%DateTime{} = hour_utc) do
    hour_utc = align_to_hour(hour_utc)
    spots = fetch_spots(hour_utc)

    if spots == [] do
      Logger.info("Pskr.CalibrationSampler: no spots for #{hour_utc} — skipping")
      %{upserted: 0, missing_hrrr_cells: 0}
    else
      hrrr_index = build_hrrr_index(hour_utc)
      kp = fetch_kp(hour_utc)
      cells = bucket_by_cell(spots)

      # Compute (band, lat, lon) → nearest_hrrr_or_nil ONCE. Both the
      # producer (enqueue_missing_hrrr) and the per-bucket sample
      # assembly need it, and `nearest_hrrr/3` is an O(profiles) scan;
      # walking the cell list twice was ~10M extra comparisons per
      # fire at typical sizes (~1k cells × ~5k profiles).
      cell_hrrr =
        Map.new(cells, fn {{_band, lat, lon} = key, _spots} ->
          {key, nearest_hrrr(lat, lon, hrrr_index)}
        end)

      missing = enqueue_missing_hrrr(cell_hrrr, hour_utc)

      samples = Enum.map(cells, &build_sample(&1, hour_utc, cell_hrrr, kp))
      %{upserted: upsert(samples), missing_hrrr_cells: missing}
    end
  end

  # ── Spot fetch + cell bucketing ─────────────────────────────────

  defp fetch_spots(hour_utc) do
    Repo.all(
      from(s in SpotHourly,
        where: s.hour_utc == ^hour_utc and not is_nil(s.midpoint_lat) and not is_nil(s.midpoint_lon),
        select: %{
          band: s.band,
          midpoint_lat: s.midpoint_lat,
          midpoint_lon: s.midpoint_lon,
          spot_count: s.spot_count,
          distance_km: s.distance_km,
          modes: s.modes
        }
      )
    )
  end

  defp bucket_by_cell(spots) do
    Enum.group_by(spots, fn s -> {s.band, snap(s.midpoint_lat), snap(s.midpoint_lon)} end)
  end

  defp snap(deg), do: Float.round(deg / @grid_step_deg) * @grid_step_deg

  # ── HRRR feature lookup ─────────────────────────────────────────

  defp build_hrrr_index(hour_utc) do
    time_start = DateTime.add(hour_utc, -@hrrr_lookahead_seconds, :second)
    time_end = DateTime.add(hour_utc, @hrrr_lookahead_seconds, :second)

    Repo.all(
      from(h in HrrrProfile,
        where: h.valid_time >= ^time_start and h.valid_time <= ^time_end,
        select: %{
          lat: h.lat,
          lon: h.lon,
          surface_temp_c: h.surface_temp_c,
          surface_dewpoint_c: h.surface_dewpoint_c,
          pwat_mm: h.pwat_mm,
          surface_pressure_mb: h.surface_pressure_mb,
          min_refractivity_gradient: h.min_refractivity_gradient,
          hpbl_m: h.hpbl_m,
          ducting_detected: h.ducting_detected
        }
      )
    )
  end

  # Pick the HRRR row whose lat/lon is closest to (lat, lon). Linear
  # scan over the in-memory window — at the typical ~5k profile
  # window size and ~1k cells per hour, this is faster than 1k DB
  # round trips. Returns `nil` if no profile is in the window.
  defp nearest_hrrr(_lat, _lon, []), do: nil

  defp nearest_hrrr(lat, lon, profiles) do
    candidates =
      Enum.filter(profiles, fn p ->
        abs(p.lat - lat) <= @hrrr_lookahead_deg and abs(p.lon - lon) <= @hrrr_lookahead_deg
      end)

    case candidates do
      [] -> nil
      list -> Enum.min_by(list, fn p -> abs(p.lat - lat) + abs(p.lon - lon) end)
    end
  end

  # ── HRRR fetch enqueue (producer side) ──────────────────────────

  # For every PSKR cell that has no HRRR profile in the lookahead
  # window, queue a point fetch at the cell midpoint. Dedup on
  # `(rounded_hrrr_lat, rounded_hrrr_lon)` so multiple bands sharing
  # a midpoint enqueue once, and let `HrrrPointEnqueuer.enqueue/1`
  # union the points into the existing `hrrr_fetch_tasks` row for
  # this hour rather than racing to insert duplicates.
  #
  # Returns the number of unique missing-HRRR cells. The caller uses
  # that to decide whether a follow-up sampler pass is needed (see
  # `Microwaveprop.Workers.PskrCalibrationWorker`).
  @spec enqueue_missing_hrrr(%{cell_key() => map() | nil}, DateTime.t()) :: non_neg_integer()
  defp enqueue_missing_hrrr(cell_hrrr, hour_utc) do
    points =
      cell_hrrr
      |> Enum.flat_map(fn
        {_key, %{}} -> []
        {{_band, lat, lon}, nil} -> [Weather.round_to_hrrr_grid(lat, lon)]
      end)
      |> Enum.uniq()

    case points do
      [] ->
        0

      pts ->
        {:ok, _} = HrrrPointEnqueuer.enqueue(%{hour_utc => pts})
        Logger.info("Pskr.CalibrationSampler: enqueued #{length(pts)} HRRR points for #{hour_utc}")
        length(pts)
    end
  end

  # ── Kp lookup ───────────────────────────────────────────────────

  # Single Kp value applies cluster-wide for the hour. We prefer the
  # integer `kp_index` (the official 3-hour Kp) and fall back to the
  # truncated `estimated_kp` when SWPC hasn't published the official
  # value yet — same precedence as `Propagation.current_kp_index/0`.
  defp fetch_kp(_hour_utc) do
    case SpaceWeather.latest_kp() do
      %{kp_index: kp} when is_integer(kp) -> kp
      %{estimated_kp: kp} when is_number(kp) -> trunc(kp)
      _ -> nil
    end
  end

  # ── Per-bucket sample assembly ──────────────────────────────────

  defp build_sample({{band, lat, lon} = key, cell_spots}, hour_utc, cell_hrrr, kp) do
    hrrr = Map.get(cell_hrrr, key)
    now = DateTime.truncate(DateTime.utc_now(), :second)

    %{
      id: Ecto.UUID.generate(),
      hour_utc: hour_utc,
      band: band,
      midpoint_lat: lat,
      midpoint_lon: lon,
      spot_count: Enum.sum(Enum.map(cell_spots, & &1.spot_count)),
      distinct_paths: length(cell_spots),
      median_distance_km: median_distance(cell_spots),
      modes: dedup_modes(cell_spots),
      surface_temp_c: hrrr && hrrr.surface_temp_c,
      surface_dewpoint_c: hrrr && hrrr.surface_dewpoint_c,
      pwat_mm: hrrr && hrrr.pwat_mm,
      surface_pressure_mb: hrrr && hrrr.surface_pressure_mb,
      min_refractivity_gradient: hrrr && hrrr.min_refractivity_gradient,
      hpbl_m: hrrr && hrrr.hpbl_m,
      hrrr_ducting_detected: hrrr && hrrr.ducting_detected,
      kp_index: kp,
      inserted_at: now,
      updated_at: now
    }
  end

  defp median_distance(cell_spots) do
    cell_spots
    |> Enum.map(& &1.distance_km)
    |> Enum.reject(&is_nil/1)
    |> case do
      [] ->
        nil

      list ->
        sorted = Enum.sort(list)
        n = length(sorted)
        mid = div(n, 2)
        if rem(n, 2) == 0, do: (Enum.at(sorted, mid - 1) + Enum.at(sorted, mid)) / 2, else: Enum.at(sorted, mid)
    end
  end

  defp dedup_modes(cell_spots) do
    cell_spots |> Enum.flat_map(& &1.modes) |> Enum.uniq()
  end

  # ── Upsert ──────────────────────────────────────────────────────

  defp upsert([]), do: 0

  defp upsert(samples) do
    {count, _} =
      Repo.insert_all(CalibrationSample, samples,
        on_conflict:
          from(c in CalibrationSample,
            update: [
              set: [
                spot_count: fragment("EXCLUDED.spot_count"),
                distinct_paths: fragment("EXCLUDED.distinct_paths"),
                median_distance_km: fragment("EXCLUDED.median_distance_km"),
                modes: fragment("EXCLUDED.modes"),
                surface_temp_c: fragment("EXCLUDED.surface_temp_c"),
                surface_dewpoint_c: fragment("EXCLUDED.surface_dewpoint_c"),
                pwat_mm: fragment("EXCLUDED.pwat_mm"),
                surface_pressure_mb: fragment("EXCLUDED.surface_pressure_mb"),
                min_refractivity_gradient: fragment("EXCLUDED.min_refractivity_gradient"),
                hpbl_m: fragment("EXCLUDED.hpbl_m"),
                hrrr_ducting_detected: fragment("EXCLUDED.hrrr_ducting_detected"),
                kp_index: fragment("EXCLUDED.kp_index"),
                updated_at: fragment("EXCLUDED.updated_at")
              ]
            ]
          ),
        conflict_target: [:hour_utc, :band, :midpoint_lat, :midpoint_lon]
      )

    count
  end

  defp align_to_hour(%DateTime{} = dt) do
    %{dt | minute: 0, second: 0, microsecond: {0, 0}}
  end
end