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

  @doc """
  Build calibration samples for one hour. Returns the upsert count.
  Returns `0` when no PSKR spots were captured for that hour — for
  example, during PSKR client downtime.
  """
  @spec build_for_hour(DateTime.t()) :: non_neg_integer()
  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")
      0
    else
      hrrr_index = build_hrrr_index(hour_utc)
      kp = fetch_kp(hour_utc)
      cells = bucket_by_cell(spots)

      # Producer side of the two-pass HRRR join: any cell whose
      # midpoint has no profile in the lookahead window gets a
      # `hrrr_fetch_tasks` row enqueued for the Rust hrrr-point-worker
      # to drain. The next rolling-window cron pass (every 5 min) will
      # find the landed profile and UPSERT the sample with non-NULL
      # weather fields. Existing cells with HRRR data already in
      # memory are skipped — no churn against the queue table.
      enqueue_missing_hrrr(cells, hour_utc, hrrr_index)

      samples = Enum.map(cells, &build_sample(&1, hour_utc, hrrr_index, kp))
      upsert(samples)
    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.
  defp enqueue_missing_hrrr(cells, hour_utc, hrrr_index) do
    points =
      cells
      |> Enum.flat_map(fn {{_band, lat, lon}, _spots} ->
        if nearest_hrrr(lat, lon, hrrr_index) do
          []
        else
          [Weather.round_to_hrrr_grid(lat, lon)]
        end
      end)
      |> Enum.uniq()

    case points do
      [] ->
        :ok

      pts ->
        {:ok, _} = HrrrPointEnqueuer.enqueue(%{hour_utc => pts})

        Logger.info("Pskr.CalibrationSampler: enqueued #{length(pts)} HRRR points for #{hour_utc}")
    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}, cell_spots}, hour_utc, hrrr_index, kp) do
    hrrr = nearest_hrrr(lat, lon, hrrr_index)
    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