defmodule Mix.Tasks.HrrrBackfill do
  @moduledoc """
  Re-fetches HRRR profiles for all QSO-linked contacts with the current
  (finer) pressure level set. Replaces existing profiles with updated data.

  Usage:
      mix hrrr_backfill              # backfill all contacts missing HRRR or with < 13 levels
      mix hrrr_backfill --all        # re-fetch all contacts regardless of level count
      mix hrrr_backfill --limit 100  # process at most 100 hours
  """

  use Mix.Task

  import Ecto.Query

  alias Microwaveprop.Radio
  alias Microwaveprop.Radio.Contact
  alias Microwaveprop.Repo
  alias Microwaveprop.Weather
  alias Microwaveprop.Weather.HrrrClient
  alias Microwaveprop.Weather.HrrrProfile
  alias Microwaveprop.Weather.SoundingParams

  require Logger

  @min_levels 13

  @impl Mix.Task
  def run(args) do
    Mix.Task.run("app.start")
    _ = Oban.pause_all_queues(Oban)

    {opts, _, _} = OptionParser.parse(args, switches: [all: :boolean, limit: :integer])
    refetch_all? = Keyword.get(opts, :all, false)
    limit = Keyword.get(opts, :limit, nil)

    # Get all distinct (rounded_hour, path_points) combinations for contacts
    contacts =
      Repo.all(
        from(c in Contact,
          where: not is_nil(c.pos1) and c.qso_timestamp >= ^~U[2014-01-01 00:00:00Z],
          select: %{id: c.id, pos1: c.pos1, pos2: c.pos2, qso_timestamp: c.qso_timestamp},
          order_by: [desc: c.qso_timestamp]
        )
      )

    # Group by rounded HRRR hour to batch requests
    by_hour =
      contacts
      |> Enum.flat_map(fn c ->
        hour = HrrrClient.nearest_hrrr_hour(c.qso_timestamp)

        c
        |> Radio.contact_path_points()
        |> Enum.map(fn {lat, lon} ->
          {rlat, rlon} = Weather.round_to_hrrr_grid(lat, lon)
          {hour, {rlat, rlon}}
        end)
      end)
      |> Enum.group_by(fn {hour, _} -> hour end, fn {_, point} -> point end)
      |> Enum.map(fn {hour, points} -> {hour, Enum.uniq(points)} end)
      |> Enum.sort_by(fn {hour, _} -> hour end, {:desc, DateTime})

    # Filter out points that already have enough levels
    by_hour =
      if refetch_all? do
        by_hour
      else
        by_hour
        |> Enum.map(fn {hour, points} ->
          needed = filter_points_needing_backfill(hour, points)
          {hour, needed}
        end)
        |> Enum.reject(fn {_hour, points} -> points == [] end)
      end

    by_hour = if limit, do: Enum.take(by_hour, limit), else: by_hour

    total = length(by_hour)
    total_points = Enum.reduce(by_hour, 0, fn {_, pts}, acc -> acc + length(pts) end)
    skipped = length(contacts) - total_points
    Logger.info("HRRR backfill: #{total} hours, #{total_points} points to fetch (#{skipped} already up to date)")

    by_hour
    |> Enum.with_index(1)
    |> Enum.each(fn {{hour, points}, idx} ->
      fetch_and_store_hour(hour, points, idx, total)

      # Rate limit to avoid overwhelming NOAA
      if idx < total, do: Process.sleep(500)
    end)

    Logger.info("HRRR backfill complete")
  end

  defp fetch_and_store_hour(hour, points, idx, total) do
    Logger.info("HRRR backfill [#{idx}/#{total}]: #{hour} (#{length(points)} points)")

    case HrrrClient.fetch_grid(points, hour) do
      {:ok, grid_data} ->
        Enum.each(grid_data, fn {{lat, lon}, data} ->
          store_profile(lat, lon, hour, data)
        end)

        Logger.info("HRRR backfill [#{idx}/#{total}]: saved #{map_size(grid_data)} profiles")

      {:error, reason} ->
        Logger.warning("HRRR backfill [#{idx}/#{total}]: failed - #{inspect(reason)}")
    end
  end

  defp filter_points_needing_backfill(hour, points) do
    # Query level counts for all points at this hour in one query
    point_tuples = Enum.map(points, fn {lat, lon} -> {lat, lon} end)
    lats = Enum.map(point_tuples, &elem(&1, 0))
    lons = Enum.map(point_tuples, &elem(&1, 1))

    existing =
      from(h in HrrrProfile,
        where: h.valid_time == ^hour and h.lat in ^lats and h.lon in ^lons,
        select: {h.lat, h.lon, fragment("array_length(profile, 1)")}
      )
      |> Repo.all()
      |> Map.new(fn {lat, lon, count} -> {{lat, lon}, count || 0} end)

    Enum.filter(points, fn {lat, lon} ->
      case Map.get(existing, {lat, lon}) do
        nil -> true
        count when count < @min_levels -> true
        _ -> false
      end
    end)
  end

  defp store_profile(lat, lon, valid_time, data) do
    params = SoundingParams.derive(data.profile)

    attrs =
      maybe_add_derived(
        %{
          valid_time: valid_time,
          lat: lat,
          lon: lon,
          run_time: data.run_time,
          profile: data.profile,
          hpbl_m: data.hpbl_m,
          pwat_mm: data.pwat_mm,
          surface_temp_c: data.surface_temp_c,
          surface_dewpoint_c: data.surface_dewpoint_c,
          surface_pressure_mb: data.surface_pressure_mb
        },
        params
      )

    Weather.upsert_hrrr_profile(attrs)
  end

  defp maybe_add_derived(attrs, nil), do: attrs

  defp maybe_add_derived(attrs, params) do
    Map.merge(attrs, %{
      surface_refractivity: params.surface_refractivity,
      min_refractivity_gradient: params.min_refractivity_gradient,
      ducting_detected: params.ducting_detected,
      duct_characteristics: params.duct_characteristics
    })
  end
end