jido

defimpl Jido.AgentServer.DirectiveExec, for: Jido.Agent.Directive.Emit do
  @moduledoc false

  require Logger

  alias Jido.Agent.Directive
  alias Jido.AgentServer.ErrorPolicy
  alias Jido.Tracing.Context, as: TraceContext

  def exec(%Directive.Emit{signal: signal, dispatch: dispatch} = directive, input_signal, state) do
    cfg = dispatch || state.default_dispatch
    runtime_context = Map.get(state, :current_runtime_context, %{})

    traced_signal =
      case TraceContext.propagate_to(signal, input_signal.id) do
        {:ok, s} -> s
        {:error, _} -> signal
      end

    case Jido.AgentServer.prepare_emit_signal(
           traced_signal,
           input_signal,
           state,
           directive,
           cfg,
           runtime_context
         ) do
      {:ok, prepared_signal, prepared_dispatch} ->
        dispatch_signal(prepared_signal, prepared_dispatch, state)
        {:async, nil, state}

      {:error, error} ->
        error_directive = %Directive.Error{error: error, context: :plugin_prepare_emit}
        ErrorPolicy.handle(error_directive, state)
    end
  end

  defp dispatch_signal(traced_signal, nil, _state) do
    send(self(), {:signal, traced_signal})
  end

  defp dispatch_signal(traced_signal, cfg, state) do
    if Code.ensure_loaded?(Jido.Signal.Dispatch) do
      task_sup =
        if state.jido, do: Jido.task_supervisor_name(state.jido), else: Jido.TaskSupervisor

      Task.Supervisor.start_child(task_sup, fn ->
        Jido.Signal.Dispatch.dispatch(traced_signal, cfg)
      end)
    else
      Logger.warning(fn -> "Jido.Signal.Dispatch not available, skipping emit" end)
    end
  end
end

defimpl Jido.AgentServer.DirectiveExec, for: Jido.Agent.Directive.Error do
  @moduledoc false

  alias Jido.AgentServer.ErrorPolicy

  def exec(error_directive, _input_signal, state) do
    ErrorPolicy.handle(error_directive, state)
  end
end

defimpl Jido.AgentServer.DirectiveExec, for: Jido.Agent.Directive.RunInstruction do
  @moduledoc false

  require Logger

  alias Jido.AgentServer.State
  alias Jido.Observe.Config, as: ObserveConfig

  def exec(
        %{instruction: instruction, result_action: result_action, meta: meta},
        input_signal,
        state
      ) do
    runtime_context = Map.get(state, :current_runtime_context, %{})

    enriched_instruction = %{
      instruction
      | context: Map.put(instruction.context || %{}, :state, state.agent.state)
    }

    execution_payload =
      enriched_instruction
      |> then(fn instruction ->
        exec_opts = ObserveConfig.action_exec_opts(state.jido, instruction.opts)
        Jido.Exec.run(%{instruction | opts: exec_opts})
      end)
      |> normalize_result_payload()
      |> Map.put(:instruction, instruction)
      |> Map.put(:meta, meta || %{})

    {agent, directives} =
      state.agent_module.cmd(
        state.agent,
        {result_action, execution_payload},
        __jido_instance__: state.jido,
        __partition__: state.partition,
        __jido_signal__: input_signal,
        __jido_action_context__: runtime_context
      )

    state = State.update_agent(state, agent)

    case State.enqueue_all(state, input_signal, runtime_context, List.wrap(directives)) do
      {:ok, state} ->
        {:ok, state}

      {:error, :queue_overflow} ->
        Logger.warning(fn -> "AgentServer #{state.id} queue overflow, dropping directives" end)
        {:ok, state}
    end
  end

  defp normalize_result_payload({:ok, result}) do
    %{
      status: :ok,
      result: result,
      effects: []
    }
  end

  defp normalize_result_payload({:ok, result, effects}) do
    %{
      status: :ok,
      result: result,
      effects: List.wrap(effects)
    }
  end

  defp normalize_result_payload({:error, reason}) do
    %{
      status: :error,
      reason: reason,
      effects: []
    }
  end

  defp normalize_result_payload({:error, reason, effects}) do
    %{
      status: :error,
      reason: reason,
      effects: List.wrap(effects)
    }
  end
end

defimpl Jido.AgentServer.DirectiveExec, for: Jido.Agent.Directive.Spawn do
  @moduledoc false

  require Logger

  def exec(%{child_spec: child_spec, tag: tag}, _input_signal, state) do
    result =
      if is_function(state.spawn_fun, 1) do
        state.spawn_fun.(child_spec)
      else
        agent_sup =
          if state.jido, do: Jido.agent_supervisor_name(state.jido), else: Jido.AgentSupervisor

        DynamicSupervisor.start_child(agent_sup, child_spec)
      end

    case result do
      {:ok, pid} ->
        Logger.debug(fn ->
          "Spawned child process #{inspect(pid)} with tag #{inspect(tag)}"
        end)

        {:ok, state}

      {:ok, pid, _info} ->
        Logger.debug(fn ->
          "Spawned child process #{inspect(pid)} with tag #{inspect(tag)}"
        end)

        {:ok, state}

      {:error, reason} ->
        Logger.error(fn -> "Failed to spawn child: #{inspect(reason)}" end)
        {:ok, state}

      :ignored ->
        {:ok, state}
    end
  end
end

defimpl Jido.AgentServer.DirectiveExec, for: Jido.Agent.Directive.Schedule do
  @moduledoc false

  alias Jido.AgentServer.Signal.Scheduled
  alias Jido.Tracing.Context, as: TraceContext

  def exec(%{delay_ms: delay, message: message}, input_signal, state) do
    signal =
      case message do
        %Jido.Signal{} = s ->
          s

        other ->
          Scheduled.new!(
            %{message: other},
            source: "/agent/#{state.id}"
          )
      end

    traced_signal =
      case TraceContext.propagate_to(signal, input_signal.id) do
        {:ok, s} -> s
        {:error, _} -> signal
      end

    Process.send_after(self(), {:scheduled_signal, traced_signal}, delay)
    {:ok, state}
  end
end

defimpl Jido.AgentServer.DirectiveExec, for: Jido.Agent.Directive.SpawnAgent do
  @moduledoc false

  require Logger

  alias Jido.Agent.Directive
  alias Jido.AgentServer
  alias Jido.AgentServer.{ChildInfo, State}
  alias Jido.RuntimeStore

  @relationship_hive :relationships
  @reserved_child_opts [:agent, :id, :jido, :parent, :partition]

  def exec(
        %{agent: agent, tag: tag, opts: opts, meta: meta, restart: restart},
        _input_signal,
        state
      ) do
    with :ok <- Directive.validate_restart_policy(restart),
         :ok <- Directive.validate_spawn_agent_opts(opts) do
      spawn_child(state, agent, tag, opts, meta, restart)
    else
      {:error, reason} ->
        Logger.error(fn ->
          "AgentServer #{state.id} failed to spawn child: #{inspect(reason)}"
        end)

        {:ok, state}
    end
  end

  defp resolve_agent_module(agent) when is_atom(agent), do: agent
  defp resolve_agent_module(%{__struct__: module}), do: module
  defp resolve_agent_module(_), do: nil

  defp spawn_child(state, agent, tag, opts, meta, restart) do
    child_id = opts[:id] || "#{state.id}/#{tag}"
    child_partition = Map.get(opts, :partition, state.partition)

    parent_ref = %{
      pid: self(),
      id: state.id,
      partition: state.partition,
      tag: tag,
      meta: meta
    }

    child_opts =
      opts
      |> Map.drop(@reserved_child_opts)
      |> Map.put(:agent, agent)
      |> Map.put(:id, child_id)
      |> Map.put(:partition, child_partition)
      |> Map.put(:parent, parent_ref)
      |> maybe_put_jido(state.jido)
      |> Map.to_list()

    child_spec = Supervisor.child_spec({AgentServer, child_opts}, restart: restart)

    supervisor =
      if state.jido, do: Jido.agent_supervisor_name(state.jido), else: Jido.AgentSupervisor

    case DynamicSupervisor.start_child(supervisor, child_spec) do
      {:ok, pid} ->
        case persist_relationship(state, child_id, child_partition, tag, meta) do
          :ok ->
            ref = Process.monitor(pid)

            child_info =
              ChildInfo.new!(%{
                pid: pid,
                ref: ref,
                module: resolve_agent_module(agent),
                id: child_id,
                partition: child_partition,
                tag: tag,
                meta: meta
              })

            new_state = State.add_child(state, tag, child_info)

            Logger.debug(fn ->
              "AgentServer #{state.id} spawned child #{child_id} with tag #{inspect(tag)}"
            end)

            {:ok, new_state}

          {:error, reason} ->
            _ = DynamicSupervisor.terminate_child(supervisor, pid)

            Logger.error(fn ->
              "AgentServer #{state.id} failed to persist relationship for child #{child_id}: #{inspect(reason)}"
            end)

            {:ok, state}
        end

      {:error, reason} ->
        Logger.error(fn ->
          "AgentServer #{state.id} failed to spawn child with restart #{inspect(restart)}: #{inspect(reason)}"
        end)

        {:ok, state}
    end
  end

  defp persist_relationship(state, child_id, child_partition, tag, meta) do
    RuntimeStore.put(
      state.jido,
      @relationship_hive,
      Jido.partition_key(child_id, child_partition),
      %{
        parent_id: state.id,
        parent_partition: state.partition,
        tag: tag,
        meta: normalize_meta(meta)
      }
    )
  end

  defp normalize_meta(meta) when is_map(meta), do: meta
  defp normalize_meta(_meta), do: %{}

  defp maybe_put_jido(opts, nil), do: opts
  defp maybe_put_jido(opts, jido), do: Map.put(opts, :jido, jido)
end

defimpl Jido.AgentServer.DirectiveExec, for: Jido.Agent.Directive.AdoptChild do
  @moduledoc false

  require Logger

  alias Jido.AgentServer
  alias Jido.AgentServer.{ChildInfo, ParentRef, State}

  def exec(%{child: child, tag: tag, meta: meta}, _input_signal, state) do
    with :ok <- ensure_tag_available(state, tag),
         {:ok, child_pid} <- resolve_child(child, state),
         :ok <- ensure_not_self(child_pid),
         {:ok, child_runtime} <- adopt_child(child_pid, tag, meta, state) do
      child_info =
        ChildInfo.new!(%{
          pid: child_pid,
          ref: Process.monitor(child_pid),
          module: child_runtime.agent_module,
          id: child_runtime.id,
          partition: child_runtime.partition,
          tag: tag,
          meta: meta
        })

      Logger.debug(fn ->
        "AgentServer #{state.id} adopted child #{child_runtime.id} with tag #{inspect(tag)}"
      end)

      {:ok, State.add_child(state, tag, child_info)}
    else
      {:error, reason} ->
        Logger.warning(fn ->
          "AgentServer #{state.id} failed to adopt child #{inspect(child)} " <>
            "with tag #{inspect(tag)}: #{inspect(reason)}"
        end)

        {:ok, state}
    end
  end

  defp ensure_tag_available(state, tag) do
    case State.get_child(state, tag) do
      nil -> :ok
      _child -> {:error, {:tag_in_use, tag}}
    end
  end

  defp resolve_child(pid, _state) when is_pid(pid) do
    if Process.alive?(pid), do: {:ok, pid}, else: {:error, :child_not_alive}
  end

  defp resolve_child(id, state) when is_binary(id) do
    case Jido.whereis(state.jido, id, partition: state.partition) do
      pid when is_pid(pid) -> {:ok, pid}
      nil -> {:error, :child_not_found}
    end
  end

  defp resolve_child(child, _state), do: {:error, {:invalid_child, child}}

  defp ensure_not_self(pid) when pid == self(), do: {:error, :cannot_adopt_self}
  defp ensure_not_self(_pid), do: :ok

  defp adopt_child(child_pid, tag, meta, state) do
    parent_ref =
      ParentRef.new!(%{
        pid: self(),
        id: state.id,
        partition: state.partition,
        tag: tag,
        meta: meta
      })

    AgentServer.adopt_parent(child_pid, parent_ref)
  end
end

defimpl Jido.AgentServer.DirectiveExec, for: Jido.Agent.Directive.StopChild do
  @moduledoc false

  alias Jido.AgentServer.StopChildRuntime

  def exec(%{tag: tag, reason: reason}, input_signal, state) do
    StopChildRuntime.exec(tag, reason, input_signal, state)
  end
end

defimpl Jido.AgentServer.DirectiveExec, for: Jido.Agent.Directive.StartSensor do
  @moduledoc false

  alias Jido.AgentServer.SensorLifecycle

  def exec(
        %{tag: tag, sensor: sensor, config: config, meta: meta, replace?: replace?} = directive,
        _input_signal,
        state
      ) do
    SensorLifecycle.start(state, tag, sensor, config, meta,
      origin: :directive,
      replace?: replace?,
      link?: Map.get(directive, :link?, false)
    )
  end
end

defimpl Jido.AgentServer.DirectiveExec, for: Jido.Agent.Directive.StopSensor do
  @moduledoc false

  alias Jido.AgentServer.SensorLifecycle

  def exec(%{tag: tag, reason: reason}, _input_signal, state) do
    SensorLifecycle.stop(state, tag, reason)
  end
end

defimpl Jido.AgentServer.DirectiveExec, for: Jido.Agent.Directive.Stop do
  @moduledoc false

  def exec(%{reason: reason}, _input_signal, state) do
    {:stop, reason, state}
  end
end

defimpl Jido.AgentServer.DirectiveExec, for: Any do
  @moduledoc false

  require Logger

  def exec(directive, _input_signal, state) do
    Logger.debug(fn -> "Ignoring unknown directive: #{inspect(directive.__struct__)}" end)
    {:ok, state}
  end
end