altworx_runbox

defmodule Runbox.Runtime.Stage.Emulator.Template do
  @moduledoc group: :internal
  @moduledoc """
  StageBased Template emulator.

  This is similar to `Runbox.Runtime.Stage.TemplateImpl` but designed specifically for the
  in-process emulation via `Runbox.Runtime.Stage.Emulator`.

  If emulating StageBased scenario, do not use this module directly, rather use
  `Runbox.Runtime.Stage.Emulator` to emulate the whole scenario.
  """
  alias Runbox.Message
  alias Runbox.Runtime.OutputAction, as: RuntimeOutputAction
  alias Runbox.Runtime.RuntimeInstruction
  alias Runbox.Runtime.RuntimeInstruction.Timeout
  alias Runbox.Runtime.Stage.Emulator.EmulationError
  alias Runbox.Runtime.Stage.Unit
  alias Runbox.Runtime.Stage.UnitRegistry
  alias Runbox.Scenario.OutputAction
  alias Runbox.Scenario.Template.StageBased, as: Template
  alias Runbox.Utils.TopologySort
  require Logger

  defstruct [:module, :unit_registry, :subscriptions, :deduplicate_inputs?]

  @log_prefix "[StageBased Emulator]"

  @opaque t :: %__MODULE__{
            module: module(),
            unit_registry: UnitRegistry.t(),
            subscriptions: [
              {entity :: (topic :: String.t()) | (template :: module()),
               message_type: Message.type()}
            ],
            deduplicate_inputs?: boolean()
          }

  @type scenario_output :: OutputAction.oa_params() | RuntimeInstruction.t()
  @type template_output :: scenario_output | Message.t()

  @timeout_type :_stage_emulated_timeout

  @doc """
  Returns whether this template requires deduplication on the inputs.

  In a StageBased scenario a deduplication is performed by a Timezip component. So we need to
  perform deduplication only in places where there would be a Timezip. Other places should happily
  receive duplicated messages.
  """
  @spec deduplicate_inputs?(t()) :: boolean()
  def deduplicate_inputs?(%__MODULE__{} = template) do
    template.deduplicate_inputs?
  end

  @doc "Returns input topics the template module is subscribed to."
  @spec input_topics(module()) :: [{type :: :input_topic | :load_topic, topic :: String.t()}]
  def input_topics(module) do
    module
    |> Template.subscriptions()
    # explicitly match all options, so other incorrect subscriptions will fail this case
    |> Enum.filter(fn
      {type, _, _} when type in [:input_topic, :load_topic] -> true
      {:template, _, _} -> false
    end)
    |> Enum.map(fn {type, topic, _} -> {type, topic} end)
  end

  @doc """
  Create a struct representing an emulated template.

  The struct holds all necessary information about the template, including its current state. This
  struct is used during the emulation, so you generally start by creating this struct via this
  function.
  """
  @spec from_module(module()) :: t()
  def from_module(module) do
    template_subscriptions =
      module
      |> Template.subscriptions()
      |> ensure_subscriptions!(module)

    %__MODULE__{
      module: module,
      unit_registry: new_unit_registry(template_subscriptions),
      subscriptions: subscriptions(template_subscriptions),
      deduplicate_inputs?: length(template_subscriptions) > 1
    }
  end

  @doc """
  Sorts a list of templates in a topological order by their subscriptions.

  Templates can depend on each other but these dependencies can never form a loop. Therefore, we can
  sort templates using these dependencies into a topological order. This gives us the templates in
  the order in which they can process messages. Then a latter template might consume outputs of a
  former template, but never the other way round.

  This also ensures the order is the same as in StageBased sorted component network. That is used to
  make Timezip emit outputs in a deterministic manner. Since there is no Timezip in the emulation we
  rely solely on the order of the templates to achieve the same order of outputs.
  """
  @spec topological_sort([t()]) :: [t()]
  def topological_sort(templates) do
    nodes =
      templates
      |> Enum.sort_by(& &1.module)
      |> Enum.map(&{&1, template_subscriptions(&1.subscriptions)})

    ensure_all_templates_exist!(nodes)

    case TopologySort.sort(nodes, & &1.module) do
      {:ok, sorted} ->
        sorted

      {:error, :loop} ->
        raise EmulationError,
          message: "Cannot topologically sort templates, there's a loop in the dependencies"
    end
  end

  @doc """
  Initialize static defined units.

  Static units are defined in `c:Runbox.Scenario.Template.StageBased.instances/0` callback. The
  units are created and initialized via the `c:Runbox.Scenario.Template.StageBased.init/2` callback.
  """
  @spec initialize_static_units(t(), start_from_ts :: non_neg_integer()) ::
          {t(), [scenario_output()]}
  def initialize_static_units(%__MODULE__{} = template, start_from) do
    {unit_registry, outputs} =
      template.module
      |> Template.instances()
      |> remove_invalid_units(template.module)
      |> Enum.map(fn {:unit, id, attributes} -> Unit.new(id, %{}, attributes) end)
      |> Enum.reduce({template.unit_registry, []}, fn unit, {registry, prev_outputs} ->
        {unit, outputs} = initialize_unit(unit, start_from, template.module)

        {
          UnitRegistry.register(registry, unit),
          [prev_outputs | process_outputs(outputs, template.module, unit.id, "init")]
        }
      end)

    {
      %__MODULE__{template | unit_registry: unit_registry},
      List.flatten(outputs)
    }
  end

  @doc """
  Determines whether a template should process a message.

  Template should process a message either when the message matches one of its subscriptions or when
  the message is a planned timeout for a unit of this template. A part of this information is where
  is the message coming from - a specific logical topic or template.
  """
  @spec should_handle_message?(
          t(),
          Message.t(),
          origin :: (logical_topic :: String.t()) | (template :: module()) | nil
        ) :: boolean()
  def should_handle_message?(
        %__MODULE__{} = template,
        %Message{type: @timeout_type, from: @timeout_type} = message,
        _
      ) do
    message.body.template_module == template.module
  end

  def should_handle_message?(%__MODULE__{} = template, %Message{} = message, from) do
    {from, message.type} in template.subscriptions
  end

  @doc """
  Handle a message in the template.

  This involves the following:
  - use the routing part of the subscriptions to find all units that should handle the message
  - handle the message in all these units via the
    `c:Runbox.Scenario.Template.StageBased.handle_message/2` callback
  - if no unit was found trigger `c:Runbox.Scenario.Template.StageBased.handle_asset_discovery/1` to
    potentially create a new unit

  Note that timeouts are routed directly to the unit that created them, they are not subjected to
  the same routing as other messages which come from input topics or other templates.
  """
  @spec handle_message(t(), Message.t()) :: {t(), [template_output()]}
  # special handling for timeouts, as they are routed directly to a unit
  def handle_message(
        %__MODULE__{} = template,
        %Message{type: @timeout_type, from: @timeout_type, body: body, timestamp: timestamp}
      ) do
    case UnitRegistry.get_unit(template.unit_registry, body.unit_id) do
      {:ok, unit} ->
        {unit_registry, outputs} =
          handle_message_in_unit(body.message, unit, template.unit_registry, template.module)

        {%__MODULE__{template | unit_registry: unit_registry}, outputs}
        |> evaluate_timeouts_in_the_same_time(timestamp)

      _not_found ->
        # Unit for the timeout doesn't exist anymore. Note this can occur naturally, since when a
        # unit is stopped we don't clear its timeouts.
        {template, []}
    end
  end

  def handle_message(%__MODULE__{} = template, %Message{} = message) do
    case UnitRegistry.lookup(template.unit_registry, message) do
      {:ok, []} ->
        {unit_registry, outputs} =
          unit_discovery(message, template.unit_registry, template.module)

        {%__MODULE__{template | unit_registry: unit_registry}, outputs}
        |> evaluate_timeouts_in_the_same_time(message.timestamp)

      {:ok, units} ->
        {unit_registry, outputs} =
          Enum.reduce(units, {template.unit_registry, []}, fn unit, {unit_registry, prev_outputs} ->
            {unit_registry, outputs} =
              handle_message_in_unit(message, unit, unit_registry, template.module)

            {unit_registry, [prev_outputs | outputs]}
          end)

        {%__MODULE__{template | unit_registry: unit_registry}, List.flatten(outputs)}
        |> evaluate_timeouts_in_the_same_time(message.timestamp)

      _error ->
        # I guess this shouldn't happen, but StageBased impl also silently ignores this. ¯\_(ツ)_/¯
        {template, []}
    end
  end

  @doc """
  Re-initialize a template.

  This step is necessary when a template is reconstructed from a savepoint or after a similar
  transformation which might corrupt internal template state.

  This is required, because this module depends on `Runbox.Runtime.Stage.UnitRegistry` which keeps
  anonymous functions in its state. Such function references might not survive a persistence (e.g.
  into ETF) and need to be reconstructed.
  """
  @spec reinitialize_template(t()) :: t()
  def reinitialize_template(%__MODULE__{} = template) do
    unit_registry =
      template.unit_registry
      |> UnitRegistry.upgrade()
      |> UnitRegistry.init()

    %__MODULE__{template | unit_registry: unit_registry}
  end

  defp initialize_unit(unit, timestamp, template_module) do
    case Template.init(template_module, timestamp, unit) do
      {:ok, outputs, %Unit{} = unit} ->
        {unit, outputs}

      other ->
        Logger.warning("""
        #{@log_prefix} Unit #{inspect(template_module)}/#{unit.id} returned bad response during \
        init: #{inspect(other)}\
        """)

        {unit, []}
    end
  rescue
    e ->
      log_init_error(template_module, unit, e, __STACKTRACE__)
      {unit, []}
  catch
    e ->
      log_init_error(template_module, unit, e, __STACKTRACE__)
      {unit, []}
  end

  defp evaluate_timeouts_in_the_same_time({template, outputs} = result, timestamp) do
    {timeouts_to_evaluate, outputs} =
      Enum.split_with(
        outputs,
        &match?(
          %RuntimeInstruction{body: %Timeout{timeout_message: %Message{timestamp: ^timestamp}}},
          &1
        )
      )

    case timeouts_to_evaluate do
      [] ->
        result

      _ ->
        {template, outputs} =
          Enum.reduce(timeouts_to_evaluate, {template, outputs}, fn timeout, {template, outputs} ->
            {template, new_outputs} = handle_message(template, timeout.body.timeout_message)
            {template, [outputs, new_outputs]}
          end)

        evaluate_timeouts_in_the_same_time({template, List.flatten(outputs)}, timestamp)
    end
  end

  defp log_init_error(template, unit, error, trace) do
    template = inspect(template)
    error = String.trim(Exception.format(:error, error, trace))

    Logger.warning("""
    #{@log_prefix} TemplateCarrier has ignored initialization of unit #{template}/#{unit.id}
    Error: #{error}\
    """)
  end

  defp unit_discovery(message, unit_registry, template_module) do
    case Template.handle_asset_discovery(template_module, message) do
      {:reply, handle_outputs, %Unit{} = new_unit} when is_list(handle_outputs) ->
        {new_unit, init_outputs} = initialize_unit(new_unit, message.timestamp, template_module)

        {
          UnitRegistry.register(unit_registry, new_unit),
          process_outputs(
            handle_outputs,
            template_module,
            new_unit.id,
            "asset discovery"
          ) ++
            process_outputs(
              init_outputs,
              template_module,
              new_unit.id,
              "init"
            )
        }

      {:reply, outputs} when is_list(outputs) ->
        {unit_registry, process_outputs(outputs, template_module, nil, "asset discovery")}

      other ->
        Logger.warning("""
        #{@log_prefix} Unit #{inspect(template_module)} returned bad response during asset \
        discovery: #{inspect(other)}\
        """)

        {unit_registry, []}
    end
  rescue
    e ->
      log_discovery_error(template_module, message, e, __STACKTRACE__)
      {unit_registry, []}
  catch
    e ->
      log_discovery_error(template_module, message, e, __STACKTRACE__)
      {unit_registry, []}
  end

  defp log_discovery_error(template, msg, error, trace) do
    Logger.warning("""
    #{@log_prefix} TemplateCarrier has ignored discovery message of template #{inspect(template)}
    Error: #{String.trim(Exception.format(:error, error, trace))}
    Message: #{inspect(msg, pretty: true)}
    """)
  end

  defp handle_message_in_unit(message, unit, unit_registry, template_module) do
    case Template.handle_message(template_module, message, unit) do
      {:reply, outputs, %Unit{} = new_unit} when is_list(outputs) ->
        {:ok, unit_registry} = UnitRegistry.update(unit_registry, new_unit)
        {unit_registry, process_outputs(outputs, template_module, new_unit.id, "handle message")}

      {:stop, outputs, _unit} when is_list(outputs) ->
        {UnitRegistry.unregister(unit_registry, unit),
         process_outputs(outputs, template_module, nil, "handle message")}

      other ->
        Logger.warning("""
        #{@log_prefix} Unit #{inspect(template_module)}/#{unit.id} returned bad response during \
        handle message: #{inspect(other)}\
        """)

        {unit_registry, []}
    end
  rescue
    e ->
      log_handle_error(template_module, unit, message, e, __STACKTRACE__)
      {unit_registry, []}
  catch
    e ->
      log_handle_error(template_module, unit, message, e, __STACKTRACE__)
      {unit_registry, []}
  end

  defp log_handle_error(template, unit, message, error, trace) do
    Logger.warning("""
    #{@log_prefix} TemplateCarrier has ignored regular message for unit \
    #{inspect(template)}/#{unit.id}
    Error: #{String.trim(Exception.format(:error, error, trace))}
    Message: #{inspect(message, pretty: true)}
    Unit: #{inspect(unit)}
    """)
  end

  defp template_subscriptions(subscriptions) do
    subscriptions
    |> Enum.flat_map(fn
      {entity, _} when is_atom(entity) -> [entity]
      _ -> []
    end)
    |> Enum.uniq()
  end

  defp new_unit_registry(template_subscriptions) do
    template_subscriptions
    |> Enum.flat_map(fn {_type, _name, unit_routing} -> unit_routing end)
    |> UnitRegistry.new()
  end

  defp subscriptions(template_subscriptions) do
    # explicitly match all options, so other incorrect subscriptions will fail this case
    Enum.flat_map(template_subscriptions, fn
      {type, entity, unit_routings} when type in [:input_topic, :load_topic, :template] ->
        Enum.map(unit_routings, fn {msg_type, _routing_def} -> {entity, msg_type} end)
    end)
  end

  defp process_outputs(outputs, template_module, unit_id, phase) do
    outputs
    |> reject_invalid_outputs(template_module, unit_id, phase)
    |> process_timeout_registrations(template_module, unit_id)
  end

  defp reject_invalid_outputs(outputs, template_module, unit_id, phase) do
    Enum.filter(outputs, fn
      %RuntimeInstruction{body: %Timeout{}} ->
        true

      %Message{} ->
        true

      other ->
        if RuntimeOutputAction.oa_body?(other) do
          true
        else
          log_bad_output(other, template_module, unit_id, phase)
          false
        end
    end)
  end

  defp log_bad_output(output, template_module, unit_id, phase) do
    unit_id_segment = if unit_id, do: "/#{unit_id}"

    Logger.warning("""
    #{@log_prefix} Unit #{inspect(template_module)}#{unit_id_segment} \
    returned bad OA during #{phase}: #{inspect(output)}\
    """)
  end

  defp process_timeout_registrations(outputs, _template_module, nil) do
    # there is no unit so we cannot register a timeout, just discard all timeouts
    Enum.reject(outputs, &match?(%RuntimeInstruction{body: %Timeout{}}, &1))
  end

  defp process_timeout_registrations(outputs, template_module, unit_id) do
    Enum.map(outputs, fn
      %RuntimeInstruction{body: %Timeout{timeout_message: %Message{} = msg}} ->
        %RuntimeInstruction{
          body: %Timeout{
            timeout_message: %Message{
              type: @timeout_type,
              from: @timeout_type,
              body: %{template_module: template_module, unit_id: unit_id, message: msg},
              timestamp: msg.timestamp
            }
          }
        }

      other ->
        other
    end)
  end

  defp ensure_all_templates_exist!(nodes) do
    dependencies =
      nodes
      |> Enum.flat_map(fn {_, deps} -> deps end)
      |> Enum.uniq()

    existing_templates = Enum.map(nodes, fn {%__MODULE__{module: mod}, _} -> mod end)
    missing = dependencies -- existing_templates

    if not Enum.empty?(missing) do
      modules = Enum.map_join(missing, ", ", &inspect(&1))
      raise EmulationError, message: "Missing templates specified in subscriptions: #{modules}"
    end
  end

  defp ensure_subscriptions!([], module) do
    raise EmulationError, message: "Template #{inspect(module)} didn't specify any subscriptions."
  end

  defp ensure_subscriptions!(subs, _module), do: subs

  defp remove_invalid_units(units, module) do
    Enum.filter(units, fn
      {:unit, _id, _attributes} ->
        true

      other ->
        Logger.warning("""
        #{@log_prefix} Ignored unknown instance definition of \
        #{inspect(module)} - #{inspect(other)}
        """)

        false
    end)
  end
end