strom

defmodule Strom.GenMix do
  @moduledoc """
  Generic functionality used by other components.
  """

  use GenServer

  @chunk 1
  @buffer 1000

  defstruct pid: nil,
            process_chunk: nil,
            inputs: [],
            outputs: %{},
            accs: %{},
            opts: [],
            chunk: @chunk,
            buffer: @buffer,
            no_wait: false,
            tasks: %{},
            tasks_started: false,
            tasks_run: false,
            asks: %{},
            data: %{},
            data_size: 0,
            waiting_tasks: %{}

  def start(%__MODULE__{process_chunk: process_chunk, opts: opts} = gm) when is_list(opts) do
    gm = %{
      gm
      | process_chunk: if(process_chunk, do: process_chunk, else: &process_chunk/4),
        chunk: Keyword.get(opts, :chunk, @chunk),
        buffer: Keyword.get(opts, :buffer, @buffer),
        no_wait: Keyword.get(opts, :no_wait, false)
    }

    partitions = PartitionSupervisor.partitions(Strom.ComponentSupervisor)
    partition_key = Enum.random(1..partitions)

    {:ok, pid} =
      DynamicSupervisor.start_child(
        {:via, PartitionSupervisor, {Strom.ComponentSupervisor, partition_key}},
        %{id: __MODULE__, start: {__MODULE__, :start_link, [gm]}, restart: :temporary}
      )

    %{gm | pid: pid}
  end

  def start_link(%__MODULE__{} = gm) do
    GenServer.start_link(__MODULE__, gm)
  end

  @impl true
  def init(%__MODULE__{} = gm) do
    {:ok, %{gm | pid: self()}}
  end

  def call(flow, gm) do
    input_streams =
      Enum.reduce(gm.inputs, %{}, fn name, acc ->
        Map.put(acc, name, Map.fetch!(flow, name))
      end)

    gm_pid = GenServer.call(gm.pid, {:start_tasks, input_streams})
    sub_flow = build_sub_flow(gm.outputs, gm_pid)

    flow
    |> Map.drop(gm.inputs)
    |> Map.merge(sub_flow)
  end

  defp build_sub_flow(outputs, gm_pid) do
    Enum.reduce(outputs, %{}, fn {output_name, _fun}, flow ->
      stream =
        Stream.resource(
          fn ->
            GenServer.call(gm_pid, :run_tasks)
            gm_pid
          end,
          fn gm_pid ->
            GenServer.cast(gm_pid, {:ask, output_name, self()})

            receive do
              {^output_name, :done} ->
                {:halt, gm_pid}

              {^output_name, events} ->
                {events, gm_pid}
            end
          end,
          fn gm_pid -> gm_pid end
        )

      Map.put(flow, output_name, stream)
    end)
  end

  def stop(gm) do
    GenServer.call(gm.pid, :stop)
  end

  def process_chunk(_input_stream_name, chunk, outputs, nil) do
    outputs
    |> Enum.reduce({%{}, false, nil}, fn {output_name, output_stream_fun}, {acc, any?, nil} ->
      {data, _} = Enum.split_with(chunk, output_stream_fun)
      {Map.put(acc, output_name, data), any? || Enum.any?(data), nil}
    end)
  end

  defp run_stream_in_task({name, stream}, {gm_pid, outputs, chunk}, process_chunk) do
    Task.Supervisor.start_child(
      {:via, PartitionSupervisor, {Strom.TaskSupervisor, self()}},
      fn ->
        acc =
          case GenServer.call(gm_pid, {:register_task, name, self()}) do
            {true, acc} ->
              acc

            {false, _} ->
              receive do
                {:run_task, acc} ->
                  acc
              end
          end

        stream
        |> Stream.chunk_every(chunk)
        |> Stream.transform(
          fn -> acc end,
          fn chunk, acc ->
            case process_chunk.(name, chunk, outputs, acc) do
              {new_data, true, new_acc} ->
                GenServer.cast(gm_pid, {:new_data, name, {new_data, new_acc}})

                receive do
                  :continue_task ->
                    :ok
                end

                {[], new_acc}

              {_new_data, false, new_acc} ->
                {[], new_acc}
            end
          end,
          fn acc -> acc end
        )
        |> Stream.run()

        GenServer.cast(gm_pid, {:task_done, name})
      end,
      restart: :transient
    )
  end

  @impl true
  def handle_call(
        {:start_tasks, input_streams},
        _from,
        %__MODULE__{tasks_started: false} = gm
      ) do
    tasks =
      Enum.reduce(input_streams, %{}, fn {name, stream}, acc ->
        {:ok, task_pid} =
          run_stream_in_task({name, stream}, {gm.pid, gm.outputs, gm.chunk}, gm.process_chunk)

        Map.put(acc, name, task_pid)
      end)

    {:reply, gm.pid, %{gm | tasks_started: true, tasks: tasks}}
  end

  def handle_call(
        {:start_tasks, _input_streams},
        _from,
        %__MODULE__{tasks_started: true} = gm
      ) do
    {:reply, gm.pid, gm}
  end

  def handle_call(:run_tasks, _from, %__MODULE__{tasks_started: true, tasks_run: false} = gm) do
    Enum.each(gm.tasks, fn {name, task_pid} ->
      send(task_pid, {:run_task, gm.accs[name]})
    end)

    {:reply, gm.pid, %{gm | tasks_run: true}}
  end

  def handle_call(:run_tasks, _from, %__MODULE__{tasks_run: true} = gm) do
    {:reply, gm.pid, gm}
  end

  def handle_call({:register_task, name, task_pid}, _from, %__MODULE__{} = gm) do
    tasks = Map.put(gm.tasks, name, task_pid)

    {:reply, {gm.tasks_run, gm.accs[name]}, %{gm | tasks: tasks}}
  end

  def handle_call(:stop, _from, %__MODULE__{} = gm) do
    Enum.each(
      Map.values(gm.tasks),
      &DynamicSupervisor.terminate_child(Strom.TaskSupervisor, &1)
    )

    {:stop, :normal, :ok, gm}
  end

  defp process_new_data(new_data, data, asks) do
    Enum.reduce(new_data, {data, asks, 0}, fn {output_name, data}, {all_data, asks, count} ->
      data_for_output = Map.get(all_data, output_name, []) ++ data

      {data_for_output, asks} =
        case {data_for_output, asks[output_name]} do
          {[], _} ->
            {data_for_output, asks}

          {data_for_output, nil} ->
            {data_for_output, asks}

          {data_for_output, client_pid} ->
            send(client_pid, {output_name, data_for_output})
            {[], Map.delete(asks, output_name)}
        end

      {Map.put(all_data, output_name, data_for_output), asks, count + length(data_for_output)}
    end)
  end

  defp continue_or_wait(name, {tasks, waiting_tasks}, {total_count, buffer}) do
    case {tasks[name], total_count < buffer} do
      {nil, _} ->
        waiting_tasks

      {task_pid, true} ->
        send(task_pid, :continue_task)
        waiting_tasks

      {task_pid, false} ->
        Map.put(waiting_tasks, name, task_pid)
    end
  end

  @impl true
  def handle_cast(
        {:new_data, input_name, {new_data, new_acc}},
        %__MODULE__{} = gm
      ) do
    {all_data, remaining_asks, total_count} = process_new_data(new_data, gm.data, gm.asks)

    waiting_tasks =
      continue_or_wait(input_name, {gm.tasks, gm.waiting_tasks}, {total_count, gm.buffer})

    {:noreply,
     %{
       gm
       | data: all_data,
         accs: Map.put(gm.accs, input_name, new_acc),
         data_size: total_count,
         asks: remaining_asks,
         waiting_tasks: waiting_tasks
     }}
  end

  def handle_cast({:ask, output_name, client_pid}, %__MODULE__{asks: asks} = gm) do
    {asks, new_data, data_size_for_output} =
      case Map.get(gm.data, output_name, []) do
        [] ->
          if map_size(gm.tasks) == 0 do
            send(client_pid, {output_name, :done})
            {Map.delete(asks, output_name), gm.data, 0}
          else
            {Map.put(asks, output_name, client_pid), gm.data, 0}
          end

        events ->
          send(client_pid, {output_name, events})

          {Map.delete(asks, output_name), Map.put(gm.data, output_name, []), length(events)}
      end

    new_data_size = gm.data_size - data_size_for_output

    waiting_tasks =
      if new_data_size < gm.buffer do
        Enum.each(gm.waiting_tasks, &send(elem(&1, 1), :continue_task))
        %{}
      else
        gm.waiting_tasks
      end

    {:noreply,
     %{
       gm
       | asks: asks,
         data: new_data,
         data_size: new_data_size,
         waiting_tasks: waiting_tasks
     }}
  end

  def handle_cast({:task_done, input_name}, %__MODULE__{} = gm) do
    {tasks, waiting_tasks} =
      if gm.no_wait do
        Enum.each(
          Map.values(gm.tasks),
          &DynamicSupervisor.terminate_child(Strom.TaskSupervisor, &1)
        )

        {%{}, %{}}
      else
        {Map.delete(gm.tasks, input_name), Map.delete(gm.waiting_tasks, input_name)}
      end

    asks =
      case map_size(tasks) do
        0 ->
          Enum.each(gm.asks, fn {output_name, client_pid} ->
            send(client_pid, {output_name, :done})
          end)

          %{}

        _more ->
          gm.asks
      end

    {:noreply, %{gm | tasks: tasks, waiting_tasks: waiting_tasks, asks: asks}}
  end
end