gust

defmodule Gust.DAG.RunRestarter.Worker do
  @moduledoc false
  alias Gust.Flows
  @behaviour Gust.DAG.RunRestarter
  use GenServer

  @impl true
  def init(init_arg) do
    {:ok, init_arg}
  end

  def child_spec(arg) do
    %{
      id: __MODULE__,
      start: {__MODULE__, :start_link, [arg]},
      type: :worker
    }
  end

  def start_link(args) do
    GenServer.start_link(__MODULE__, args, name: __MODULE__)
  end

  @impl true
  def restart(dags) do
    GenServer.call(__MODULE__, {:restart, dags})
  end

  @impl true
  def restart_enqueued(dag_id) do
    GenServer.cast(__MODULE__, {:restart_enqueued, dag_id})
  end

  # TODO: unify dag reload; and do not start dag_def with errors
  @impl true
  def handle_cast({:restart_enqueued, dag_id}, state) do
    {:ok, dag_def} = Gust.DAG.Loader.reload_definition(dag_id)

    _ =
      get_runs([dag_id], :enqueued)
      |> Stream.each(fn run -> start_run(run, dag_def) end)
      |> Enum.to_list()

    {:noreply, state}
  end

  @impl true
  def handle_call({:restart, dags}, _from, state) do
    dag_ids = Map.keys(dags)

    runs =
      get_runs(dag_ids, :running)
      |> Stream.filter(fn run -> map_size(dags[run.dag_id].error) == 0 end)
      |> Stream.map(fn run -> start_run(run, dags[run.dag_id]) end)
      |> Enum.to_list()

    {:reply, runs, state}
  end

  defp get_runs(dag_ids, status) do
    Flows.get_running_runs_by_dag(dag_ids, status)
  end

  defp start_run(run, dag_def) do
    {:ok, _pid} = Gust.DAG.RunnerSupervisor.start_child(run, dag_def)
    run
  end
end