defmodule Mix.Tasks.Weather.RebatchAsos do
  @shortdoc "Collapse pending single-station ASOS jobs into batched jobs"

  @moduledoc """
  Converts every pending single-station `WeatherFetchWorker` ASOS job
  into the batched (`asos_batch`) shape — one job per unique
  `(start_dt, end_dt)` window covering every station that had a
  pending job for that window.

  Why: the worker still accepts the old `"asos"` fetch_type for
  backward compat, but every such job pays the full `IemRateLimiter`
  gap + IEM 429-retry tail for a single station's rows. Collapsing
  them into batched jobs gets the queue through the backlog at the
  new per-request efficiency.

  Usage (idempotent — running it twice is a no-op):

      mix weather.rebatch_asos               # rewrites everything
      mix weather.rebatch_asos --dry-run     # report only

  """

  use Mix.Task

  import Ecto.Query

  alias Microwaveprop.Repo
  alias Microwaveprop.Workers.WeatherFetchWorker

  require Logger

  @impl Mix.Task
  def run(args) do
    {opts, _, _} = OptionParser.parse(args, strict: [dry_run: :boolean])
    Mix.Task.run("app.start")
    dry_run? = Keyword.get(opts, :dry_run, false)

    jobs = load_pending_asos_jobs()

    if jobs == [] do
      Mix.shell().info("No pending `asos` jobs. Nothing to do.")
    else
      do_rebatch(jobs, dry_run?)
    end
  end

  @doc false
  def load_pending_asos_jobs do
    query =
      from(j in Oban.Job,
        where:
          j.worker == "Microwaveprop.Workers.WeatherFetchWorker" and
            j.state in ["available", "scheduled", "retryable"] and
            fragment("?->>'fetch_type' = ?", j.args, "asos"),
        select: %{id: j.id, args: j.args}
      )

    Repo.all(query)
  end

  defp do_rebatch(jobs, dry_run?) do
    groups =
      jobs
      |> Enum.group_by(fn %{args: a} ->
        {Map.get(a, "start_dt"), Map.get(a, "end_dt")}
      end)
      |> Enum.reject(fn {{s, e}, _} -> is_nil(s) or is_nil(e) end)

    Mix.shell().info("Found #{length(jobs)} pending `asos` jobs across #{length(groups)} time windows.")

    batch_args =
      Enum.map(groups, fn {{start_dt, end_dt}, members} ->
        # Deterministic sort (matches the enqueuer's ordering) so a
        # re-run with identical membership hits Oban's unique dedup.
        sorted = Enum.sort_by(members, fn %{args: a} -> a["station_code"] end)

        %{
          "fetch_type" => "asos_batch",
          "station_ids" => Enum.map(sorted, fn %{args: a} -> a["station_id"] end),
          "station_codes" => Enum.map(sorted, fn %{args: a} -> a["station_code"] end),
          "start_dt" => start_dt,
          "end_dt" => end_dt,
          "_source_ids" => Enum.map(sorted, & &1.id)
        }
      end)

    if dry_run? do
      Enum.each(batch_args, fn a ->
        Mix.shell().info("[dry-run] #{length(a["station_codes"])} stations for #{a["start_dt"]}..#{a["end_dt"]}")
      end)
    else
      apply_rebatch(batch_args)
    end
  end

  defp apply_rebatch(batch_args) do
    {inserted, cancelled} =
      Enum.reduce(batch_args, {0, 0}, fn a, {ins, can} ->
        source_ids = a["_source_ids"]
        job_args = Map.delete(a, "_source_ids")

        Repo.transaction(fn ->
          _ = job_args |> WeatherFetchWorker.new() |> Oban.insert!()
          {:ok, _n} = Oban.cancel_all_jobs(from(j in Oban.Job, where: j.id in ^source_ids))
        end)

        {ins + 1, can + length(source_ids)}
      end)

    Mix.shell().info("Rebatched: inserted #{inserted} batch jobs, cancelled #{cancelled} single-station jobs.")
  end
end