bb

# SPDX-FileCopyrightText: 2025 James Harton
#
# SPDX-License-Identifier: Apache-2.0

defmodule BB.Safety.Controller do
  @moduledoc """
  Global safety controller that owns arm/disarm state for all robots.

  Part of BB's application supervision tree (not per-robot), so it survives
  robot crashes and maintains safety state. Runs at high scheduler priority.

  Uses two ETS tables:

  1. **Robots table** (protected set) - safety state per robot, writes only via GenServer
  2. **Handlers table** (public bag) - direct writes for registration

  Monitors robot supervisors and cleans up state when they terminate.

  ## Safety States

  - `:disarmed` - Robot is safely disarmed, all disarm callbacks succeeded
  - `:armed` - Robot is armed and ready to operate
  - `:disarming` - Disarm in progress, callbacks running concurrently
  - `:error` - Disarm attempted but one or more callbacks failed; hardware may not be safe

  When in `:error` state, the robot cannot be armed until `force_disarm/1` is called
  to acknowledge the error and reset to `:disarmed`.

  Disarm callbacks run concurrently with a timeout. If any callback fails or times out,
  the robot transitions to `:error` state.

  Note: The executing/idle distinction is handled by Runtime as it's not safety-critical.
  """
  use GenServer
  require Logger

  alias BB.{Message, PubSub}
  alias BB.Safety.HardwareError
  alias BB.StateMachine.Transition

  @robots_table Module.concat(__MODULE__, Robots)
  @handlers_table Module.concat(__MODULE__, Handlers)
  @default_disarm_timeout_ms 5_000

  @type safety_state :: :disarmed | :armed | :disarming | :error

  @doc false
  def start_link(opts) do
    GenServer.start_link(__MODULE__, opts, name: __MODULE__)
  end

  @doc """
  Check if a robot is armed.

  Fast ETS read - does not go through GenServer.
  """
  @spec armed?(module()) :: boolean()
  def armed?(robot_module) do
    case :ets.lookup(@robots_table, robot_module) do
      [{^robot_module, :armed, _ref}] -> true
      _ -> false
    end
  end

  @doc """
  Get current safety state for a robot.

  Fast ETS read - does not go through GenServer.
  Returns `:armed`, `:disarmed`, or `:error`.
  """
  @spec state(module()) :: safety_state()
  def state(robot_module) do
    case :ets.lookup(@robots_table, robot_module) do
      [{^robot_module, state, _ref}] -> state
      [] -> :disarmed
    end
  end

  @doc """
  Check if a robot is in error state.

  Returns `true` if a disarm operation failed and the robot requires
  manual intervention via `force_disarm/1`.

  Fast ETS read - does not go through GenServer.
  """
  @spec in_error?(module()) :: boolean()
  def in_error?(robot_module) do
    case :ets.lookup(@robots_table, robot_module) do
      [{^robot_module, :error, _ref}] -> true
      _ -> false
    end
  end

  @doc """
  Check if a robot is currently disarming.

  Returns `true` while disarm callbacks are running.

  Fast ETS read - does not go through GenServer.
  """
  @spec disarming?(module()) :: boolean()
  def disarming?(robot_module) do
    case :ets.lookup(@robots_table, robot_module) do
      [{^robot_module, :disarming, _ref}] -> true
      _ -> false
    end
  end

  @doc """
  Register a robot when it starts.

  Called by `BB.Supervisor` during robot startup. Sets up monitoring of the
  robot's supervision tree for automatic cleanup on crash.
  """
  @spec register_robot(module()) :: :ok | {:error, term()}
  def register_robot(robot_module) do
    GenServer.call(__MODULE__, {:register_robot, robot_module})
  end

  @doc """
  Register a robot's topology supervisor.

  Called by `BB.TopologySupervisor` during its own startup. When the topology
  supervisor terminates (because its restart budget is exhausted) the safety
  controller force-disarms the robot and transitions it to `:error` state.
  """
  @spec register_topology_supervisor(module()) :: :ok | {:error, :not_registered}
  def register_topology_supervisor(robot_module) do
    GenServer.call(__MODULE__, {:register_topology_supervisor, robot_module, self()})
  end

  @doc """
  Register a safety handler (actuator/sensor/controller).

  Called by processes in their `init/1`. The opts should contain all
  hardware-specific parameters needed to call `disarm/1` without GenServer state.

  Writes directly to ETS to avoid blocking on the Controller's mailbox.
  Uses a cast to set up process monitoring for cleanup on handler restart.
  """
  @spec register(module(), keyword()) :: :ok
  def register(module, opts) do
    robot = Keyword.fetch!(opts, :robot)
    path = Keyword.fetch!(opts, :path)
    disarm_opts = Keyword.get(opts, :opts, [])
    pid = self()

    # Direct ETS write - no GenServer call needed
    :ets.insert(@handlers_table, {robot, module, path, disarm_opts, pid})

    # Async monitoring setup - fire and forget
    GenServer.cast(__MODULE__, {:monitor_handler, pid, robot})

    :ok
  end

  @doc """
  Get list of registered handler modules for a robot.

  Used by Runtime to verify all safety handlers have registered on startup.
  """
  @spec registered_handlers(module()) :: [module()]
  def registered_handlers(robot_module) do
    @handlers_table
    |> :ets.lookup(robot_module)
    |> Enum.map(fn {_robot, module, _path, _opts, _pid} -> module end)
  end

  @doc """
  Arm the robot.

  Goes through GenServer to ensure proper state transitions and event publishing.
  Cannot arm if robot is in `:error` state - must call `force_disarm/1` first.
  """
  @spec arm(module()) :: :ok | {:error, :already_armed | :in_error | :not_registered}
  def arm(robot_module), do: GenServer.call(__MODULE__, {:arm, robot_module})

  @doc """
  Disarm the robot.

  Goes through GenServer. Calls all registered `disarm/1` callbacks before
  updating state. If any callback fails, the robot transitions to `:error`
  state instead of `:disarmed`, and this function returns an error with
  details of the failures.

  When in `:error` state, the robot cannot be armed until `force_disarm/1`
  is called to acknowledge the failure and reset to `:disarmed`.

  ## Options

    * `:timeout` - timeout in milliseconds for each disarm callback.
      Defaults to #{@default_disarm_timeout_ms}ms. The GenServer call timeout
      is set to `timeout + 5000` to allow for processing overhead.
  """
  @spec disarm(module(), keyword()) ::
          :ok | {:error, :already_disarmed | :not_registered | {:disarm_failed, list()}}
  def disarm(robot_module, opts \\ []) do
    timeout = Keyword.get(opts, :timeout, @default_disarm_timeout_ms)
    call_timeout = timeout + 5_000
    GenServer.call(__MODULE__, {:disarm, robot_module, timeout}, call_timeout)
  end

  @doc """
  Force disarm from error state.

  Use this function to acknowledge a failed disarm operation and reset the
  robot to `:disarmed` state. This should only be called after manually
  verifying that hardware is in a safe state.

  **WARNING**: This bypasses safety checks. Only use when you have manually
  verified that all actuators are disabled and the robot is safe.
  """
  @spec force_disarm(module()) :: :ok | {:error, :not_in_error | :not_registered}
  def force_disarm(robot_module), do: GenServer.call(__MODULE__, {:force_disarm, robot_module})

  @doc """
  Transition the robot to `:error` state.

  Used by `BB.Safety.disarm/2`'s command-routing layer when a user-defined
  disarm command fails before it can complete the actual hardware-safe
  disarm sequence. By the issue's failure semantics, an incomplete disarm
  means hardware may not be in a safe state, so the robot is parked in
  `:error` and requires `force_disarm/1` to recover.

  Returns `:ok` even if the robot is already in `:error` (idempotent).
  Returns `{:error, :not_registered}` if the robot is unknown.
  Returns `{:error, :already_disarmed}` if the robot is already disarmed —
  no escalation is needed because the disarm command did, in fact, succeed
  at flipping state before returning its error.
  """
  @spec transition_to_error(module()) ::
          :ok | {:error, :already_disarmed | :not_registered}
  def transition_to_error(robot_module) do
    GenServer.call(__MODULE__, {:transition_to_error, robot_module})
  end

  @doc """
  Report a hardware error from a component.

  Publishes a `BB.Safety.HardwareError` message to `[:safety, :error]`. Does
  not affect safety state - components that need to escalate should crash so
  the supervision tree can decide whether the topology supervisor's restart
  budget is exhausted.
  """
  @spec report_error(module(), [atom()], term()) :: :ok
  def report_error(robot_module, path, error) do
    GenServer.cast(__MODULE__, {:report_error, robot_module, path, error})
  end

  # --- GenServer Callbacks ---

  @impl GenServer
  def init(_opts) do
    Process.flag(:priority, :high)

    # Robots table (protected): safety state, writes only via GenServer
    # {robot_module, :armed | :disarmed, supervisor_monitor_ref}
    :ets.new(@robots_table, [:named_table, :protected, :set, read_concurrency: true])

    # Handlers table (public bag): direct writes for registration
    # {robot_module, module, path, opts, handler_pid}
    :ets.new(@handlers_table, [:named_table, :public, :bag, read_concurrency: true])

    {:ok, %{handler_monitors: %{}, topology_monitors: %{}}}
  end

  @impl GenServer
  def handle_call({:register_robot, robot_module}, _from, state) do
    # The supervisor is registered directly with the robot_module name
    case Process.whereis(robot_module) do
      pid when is_pid(pid) ->
        ref = Process.monitor(pid)
        :ets.insert(@robots_table, {robot_module, :disarmed, ref})
        {:reply, :ok, state}

      nil ->
        {:reply, {:error, :supervisor_not_found}, state}
    end
  end

  def handle_call({:register_topology_supervisor, robot_module, pid}, _from, state) do
    case :ets.lookup(@robots_table, robot_module) do
      [{^robot_module, _safety_state, _ref}] ->
        topology_ref = Process.monitor(pid)

        new_state =
          update_in(state.topology_monitors, &Map.put(&1, topology_ref, robot_module))

        {:reply, :ok, new_state}

      [] ->
        {:reply, {:error, :not_registered}, state}
    end
  end

  def handle_call({:arm, robot_module}, _from, state) do
    case :ets.lookup(@robots_table, robot_module) do
      [{^robot_module, :disarmed, ref}] ->
        :ets.insert(@robots_table, {robot_module, :armed, ref})
        publish_transition(robot_module, :disarmed, :armed)
        {:reply, :ok, state}

      [{^robot_module, :armed, _}] ->
        {:reply, {:error, :already_armed}, state}

      [{^robot_module, :disarming, _}] ->
        {:reply, {:error, :disarming}, state}

      [{^robot_module, :error, _}] ->
        {:reply, {:error, :in_error}, state}

      [] ->
        {:reply, {:error, :not_registered}, state}
    end
  end

  def handle_call({:disarm, robot_module, timeout}, _from, state) do
    case :ets.lookup(@robots_table, robot_module) do
      [{^robot_module, :disarmed, _}] ->
        {:reply, {:error, :already_disarmed}, state}

      [{^robot_module, :disarming, _}] ->
        {:reply, {:error, :already_disarming}, state}

      [{^robot_module, :error, _}] ->
        {:reply, {:error, :already_in_error}, state}

      [{^robot_module, :armed, ref}] ->
        # Immediately transition to :disarming to prevent new commands
        :ets.insert(@robots_table, {robot_module, :disarming, ref})
        publish_transition(robot_module, :armed, :disarming)

        # Run callbacks concurrently with timeout
        case disarm_all_handlers(robot_module, timeout) do
          :ok ->
            :ets.insert(@robots_table, {robot_module, :disarmed, ref})
            publish_transition(robot_module, :disarming, :disarmed)
            {:reply, :ok, state}

          {:error, failures} ->
            :ets.insert(@robots_table, {robot_module, :error, ref})
            publish_transition(robot_module, :disarming, :error)
            {:reply, {:error, {:disarm_failed, failures}}, state}
        end

      [] ->
        {:reply, {:error, :not_registered}, state}
    end
  end

  def handle_call({:transition_to_error, robot_module}, _from, state) do
    case :ets.lookup(@robots_table, robot_module) do
      [{^robot_module, :disarmed, _}] ->
        {:reply, {:error, :already_disarmed}, state}

      [{^robot_module, :error, _}] ->
        # Idempotent — already in error
        {:reply, :ok, state}

      [{^robot_module, current, ref}] ->
        Logger.critical(
          "Transitioning #{inspect(robot_module)} to :error - " <>
            "disarm command failed to complete safely from #{current}"
        )

        :ets.insert(@robots_table, {robot_module, :error, ref})
        publish_transition(robot_module, current, :error)
        {:reply, :ok, state}

      [] ->
        {:reply, {:error, :not_registered}, state}
    end
  end

  def handle_call({:force_disarm, robot_module}, _from, state) do
    case :ets.lookup(@robots_table, robot_module) do
      [{^robot_module, :error, ref}] ->
        Logger.warning(
          "Force disarm called for #{inspect(robot_module)} - " <>
            "operator has acknowledged hardware may not be in safe state"
        )

        :ets.insert(@robots_table, {robot_module, :disarmed, ref})
        publish_transition(robot_module, :error, :disarmed)
        {:reply, :ok, state}

      [{^robot_module, _, _}] ->
        {:reply, {:error, :not_in_error}, state}

      [] ->
        {:reply, {:error, :not_registered}, state}
    end
  end

  @impl GenServer
  def handle_cast({:monitor_handler, pid, robot}, state) do
    # Set up monitoring for cleanup on handler restart
    ref = Process.monitor(pid)
    new_state = update_in(state.handler_monitors, &Map.put(&1, ref, {robot, pid}))
    {:noreply, new_state}
  end

  def handle_cast({:report_error, robot_module, path, error}, state) do
    Logger.warning(
      "Hardware error reported for #{inspect(robot_module)} at #{inspect(path)}: #{inspect(error)}"
    )

    publish_hardware_error(robot_module, path, error)

    {:noreply, state}
  end

  @impl GenServer
  def handle_info({:DOWN, ref, :process, pid, _reason}, state) do
    cond do
      Map.has_key?(state.topology_monitors, ref) ->
        handle_topology_down(ref, state)

      match?([[_]], :ets.match(@robots_table, {:"$1", :_, ref})) ->
        handle_root_down(ref, state)

      Map.has_key?(state.handler_monitors, ref) ->
        handle_handler_down(ref, pid, state)

      true ->
        {:noreply, state}
    end
  end

  defp handle_topology_down(ref, state) do
    {robot_module, topology_monitors} = Map.pop(state.topology_monitors, ref)

    if robot_shutting_down?(robot_module) do
      Logger.debug(
        "Topology supervisor for #{inspect(robot_module)} stopped during robot shutdown - ignoring"
      )
    else
      force_disarm_for_topology_failure(robot_module)
    end

    {:noreply, %{state | topology_monitors: topology_monitors}}
  end

  defp force_disarm_for_topology_failure(robot_module) do
    Logger.critical(
      "Topology supervisor for #{inspect(robot_module)} has stopped - force-disarming"
    )

    previous_state = state(robot_module)

    case disarm_all_handlers(robot_module, @default_disarm_timeout_ms) do
      :ok ->
        Logger.info(
          "All disarm callbacks succeeded for #{inspect(robot_module)} after topology failure"
        )

      {:error, failures} ->
        Logger.critical(
          "DISARM CALLBACKS FAILED for #{inspect(robot_module)} after topology failure: " <>
            "#{length(failures)} handler(s) failed to disarm - HARDWARE MAY NOT BE SAFE"
        )
    end

    case :ets.lookup(@robots_table, robot_module) do
      [{^robot_module, _state, sup_ref}] ->
        :ets.insert(@robots_table, {robot_module, :error, sup_ref})
        publish_transition(robot_module, previous_state, :error)

      [] ->
        :ok
    end
  end

  defp robot_shutting_down?(robot_module) do
    case Process.whereis(BB.PubSub.registry_name(robot_module)) do
      nil -> true
      _pid -> false
    end
  end

  defp handle_root_down(ref, state) do
    [[robot_module]] = :ets.match(@robots_table, {:"$1", :_, ref})

    Logger.warning("Robot #{inspect(robot_module)} supervisor crashed, disarming all handlers")

    case disarm_all_handlers(robot_module, @default_disarm_timeout_ms) do
      :ok ->
        Logger.info("All disarm callbacks succeeded for crashed robot #{inspect(robot_module)}")

      {:error, failures} ->
        Logger.critical(
          "DISARM CALLBACKS FAILED for crashed robot #{inspect(robot_module)}: " <>
            "#{length(failures)} handler(s) failed to disarm - HARDWARE MAY NOT BE SAFE"
        )
    end

    :ets.delete(@robots_table, robot_module)
    :ets.match_delete(@handlers_table, {robot_module, :_, :_, :_, :_})

    topology_monitors =
      state.topology_monitors
      |> Enum.reject(fn {_ref, robot} -> robot == robot_module end)
      |> Map.new()

    {:noreply, %{state | topology_monitors: topology_monitors}}
  end

  defp handle_handler_down(ref, pid, state) do
    {{robot, ^pid}, handler_monitors} = Map.pop(state.handler_monitors, ref)
    :ets.match_delete(@handlers_table, {robot, :_, :_, :_, pid})
    {:noreply, %{state | handler_monitors: handler_monitors}}
  end

  # --- Private Functions ---

  defp disarm_all_handlers(robot_module, timeout) do
    # Bag table: each row is {robot_module, module, path, opts, pid}
    handlers = :ets.lookup(@handlers_table, robot_module)

    # Run callbacks concurrently with timeout
    failures =
      handlers
      |> Task.async_stream(
        fn {_robot, module, path, opts, _pid} ->
          {path, safe_disarm(module, path, opts)}
        end,
        timeout: timeout,
        on_timeout: :kill_task,
        ordered: false
      )
      |> Enum.reduce([], fn
        {:ok, {_path, :ok}}, acc -> acc
        {:ok, {path, {:error, error}}}, acc -> [{path, error} | acc]
        {:exit, :timeout}, acc -> [{:unknown, {:timeout, timeout}} | acc]
      end)

    case failures do
      [] -> :ok
      _ -> {:error, failures}
    end
  end

  defp safe_disarm(module, path, opts) do
    case module.disarm(opts) do
      :ok ->
        :ok

      {:error, reason} ->
        Logger.error("Disarm failed for #{inspect(path)}: returned {:error, #{inspect(reason)}}")

        {:error, {:returned_error, reason}}
    end
  rescue
    e ->
      Logger.error("Disarm failed for #{inspect(path)}: #{Exception.message(e)}")
      {:error, {:exception, Exception.message(e)}}
  catch
    kind, reason ->
      Logger.error("Disarm failed for #{inspect(path)}: #{inspect({kind, reason})}")
      {:error, {kind, reason}}
  end

  @impl GenServer
  def terminate(reason, _state) do
    armed_robots =
      @robots_table
      |> :ets.tab2list()
      |> Enum.filter(fn {_robot_module, safety_state, _ref} ->
        safety_state in [:armed, :disarming]
      end)

    if armed_robots != [] do
      Logger.warning("Safety controller terminating (#{inspect(reason)}), disarming all robots")
      Enum.each(armed_robots, &emergency_disarm_robot/1)
    end

    :ok
  end

  defp emergency_disarm_robot({robot_module, _safety_state, _ref}) do
    Logger.info("Attempting emergency disarm for #{inspect(robot_module)}")

    case disarm_all_handlers(robot_module, @default_disarm_timeout_ms) do
      :ok ->
        Logger.info("Emergency disarm succeeded for #{inspect(robot_module)}")

      {:error, failures} ->
        Logger.critical(
          "EMERGENCY DISARM FAILED for #{inspect(robot_module)}: " <>
            "#{length(failures)} handler(s) failed - HARDWARE MAY NOT BE SAFE"
        )
    end
  end

  defp publish_transition(robot_module, from, to) do
    message = Message.new!(Transition, :state_machine, from: from, to: to)
    safe_publish(robot_module, [:state_machine], message)
  end

  defp publish_hardware_error(robot_module, path, error) do
    message = Message.new!(HardwareError, :safety, path: path, error: error)
    safe_publish(robot_module, [:safety, :error], message)
  end

  defp safe_publish(robot_module, path, message) do
    PubSub.publish(robot_module, path, message)
  rescue
    ArgumentError -> :ok
  end
end