defmodule Microwaveprop.Weather.RebatchAsos do
  @moduledoc """
  Converts pending single-station `WeatherFetchWorker` ASOS jobs 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.

  Callable from both a Mix task (`mix weather.rebatch_asos`) and a
  release `eval` shell (`bin/microwaveprop eval
  'Microwaveprop.Weather.RebatchAsos.run()'`). In prod the Mix API
  isn't available, so this module is Mix-free.

  Idempotent: a second run with no pending "asos" jobs is a no-op.
  """

  import Ecto.Query

  alias Microwaveprop.Repo
  alias Microwaveprop.Workers.WeatherFetchWorker

  @doc """
  Run the rebatch. Options:
    * `:dry_run` (boolean) — log the plan but don't mutate the queue.

  Returns a summary map: `%{found: n, windows: n, inserted: n, cancelled: n}`.
  """
  @spec run(keyword()) :: %{
          found: non_neg_integer(),
          windows: non_neg_integer(),
          inserted: non_neg_integer(),
          cancelled: non_neg_integer()
        }
  def run(opts \\ []) do
    dry_run? = Keyword.get(opts, :dry_run, false)
    jobs = load_pending_asos_jobs()

    if jobs == [] do
      IO.puts("No pending `asos` jobs. Nothing to do.")
      %{found: 0, windows: 0, inserted: 0, cancelled: 0}
    else
      do_rebatch(jobs, dry_run?)
    end
  end

  @doc """
  Converts pending `asos_batch` and `asos` jobs into the newer
  `asos_day` shape: one job per unique `(station_id, station_code, date)`
  tuple, where `date` enumerates every UTC date the source window
  touched. Cross-contact dedup is automatic via Oban's unique-args.

  Expected to shrink the pending queue substantially because many
  QSOs reference the same (station, date) and currently have their
  own 4h-window job. Idempotent; dry-run supported.
  """
  @spec to_day_jobs(keyword()) :: %{
          source_jobs: non_neg_integer(),
          unique_day_keys: non_neg_integer(),
          inserted: non_neg_integer(),
          cancelled: non_neg_integer()
        }
  def to_day_jobs(opts \\ []) do
    dry_run? = Keyword.get(opts, :dry_run, false)
    jobs = load_pending_window_jobs()

    if jobs == [] do
      IO.puts("No pending `asos`/`asos_batch` jobs. Nothing to do.")
      %{source_jobs: 0, unique_day_keys: 0, inserted: 0, cancelled: 0}
    else
      do_to_day(jobs, dry_run?)
    end
  end

  defp load_pending_window_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' IN ('asos', 'asos_batch')", j.args),
        select: %{id: j.id, args: j.args}
      )

    Repo.all(query)
  end

  defp do_to_day(jobs, dry_run?) do
    # Expand each source job into a list of (station_id, station_code, date)
    # tuples. asos_batch → one tuple per (station_code, date-in-window);
    # asos → one tuple per (single station_code, date-in-window).
    expanded =
      Enum.flat_map(jobs, fn %{id: id, args: a} ->
        a
        |> expand_tuples()
        |> Enum.map(fn {sid, code, date} -> {sid, code, date, id} end)
      end)

    # Group by the unique dedup key (station_id, date) and keep one
    # station_code per group (always the same since station_id is the
    # key). Track source_ids to cancel afterwards.
    by_key =
      Enum.group_by(expanded, fn {sid, _code, date, _id} -> {sid, date} end)

    IO.puts("Expanded #{length(jobs)} source jobs into #{map_size(by_key)} unique (station, date) keys.")

    inserts =
      Enum.map(by_key, fn {{sid, date}, members} ->
        {_, code, _, _} = hd(members)
        source_ids = members |> Enum.map(fn {_, _, _, id} -> id end) |> Enum.uniq()

        %{
          "fetch_type" => "asos_day",
          "station_id" => sid,
          "station_code" => code,
          "date" => Date.to_iso8601(date),
          "_source_ids" => source_ids
        }
      end)

    if dry_run? do
      Enum.each(Enum.take(inserts, 20), fn a ->
        IO.puts("[dry-run] #{a["station_code"]} #{a["date"]} (collapses #{length(a["_source_ids"])} sources)")
      end)

      %{source_jobs: length(jobs), unique_day_keys: length(inserts), inserted: 0, cancelled: 0}
    else
      apply_day_conversion(inserts, length(jobs))
    end
  end

  defp expand_tuples(%{"fetch_type" => "asos_batch"} = a) do
    case window_dates(a["start_dt"], a["end_dt"]) do
      [] ->
        []

      dates ->
        ids = a["station_ids"] || []
        codes = a["station_codes"] || []

        for {sid, code} <- Enum.zip(ids, codes), date <- dates do
          {sid, code, date}
        end
    end
  end

  defp expand_tuples(%{"fetch_type" => "asos"} = a) do
    case window_dates(a["start_dt"], a["end_dt"]) do
      [] -> []
      dates -> Enum.map(dates, fn date -> {a["station_id"], a["station_code"], date} end)
    end
  end

  defp window_dates(nil, _), do: []
  defp window_dates(_, nil), do: []

  defp window_dates(start_str, end_str) do
    with {:ok, start_dt, _} <- DateTime.from_iso8601(start_str),
         {:ok, end_dt, _} <- DateTime.from_iso8601(end_str) do
      start_date = DateTime.to_date(start_dt)
      end_date = DateTime.to_date(end_dt)

      if Date.compare(start_date, end_date) == :eq do
        [start_date]
      else
        [start_date, end_date]
      end
    else
      _ -> []
    end
  end

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

        Repo.transaction(fn ->
          # Unique-args dedup means duplicate inserts are silently no-ops.
          _ = job_args |> WeatherFetchWorker.new() |> Oban.insert!()
        end)

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

    # Cancel the source jobs in bulk after inserts succeed.
    {:ok, cancelled_n} =
      Oban.cancel_all_jobs(
        from(j in Oban.Job,
          where:
            j.worker == "Microwaveprop.Workers.WeatherFetchWorker" and
              j.state in ["available", "scheduled", "retryable"] and
              fragment("?->>'fetch_type' IN ('asos', 'asos_batch')", j.args)
        )
      )

    IO.puts("Converted: inserted #{inserted} asos_day jobs, cancelled #{cancelled_n} source jobs.")

    %{source_jobs: source_jobs, unique_day_keys: length(inserts), inserted: inserted, cancelled: cancelled}
  end

  defp 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)

    IO.puts("Found #{length(jobs)} pending `asos` jobs across #{length(groups)} time windows.")

    batch_args =
      Enum.map(groups, fn {{start_dt, end_dt}, members} ->
        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 ->
        IO.puts("[dry-run] #{length(a["station_codes"])} stations for #{a["start_dt"]}..#{a["end_dt"]}")
      end)

      %{found: length(jobs), windows: length(groups), inserted: 0, cancelled: 0}
    else
      apply_rebatch(batch_args, length(jobs))
    end
  end

  defp apply_rebatch(batch_args, found) 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)

    IO.puts("Rebatched: inserted #{inserted} batch jobs, cancelled #{cancelled} single-station jobs.")

    %{found: found, windows: length(batch_args), inserted: inserted, cancelled: cancelled}
  end
end