firebird

defmodule Firebird.Lazy do
  @moduledoc """
  Lazy-loading WASM instances that initialize on first use.

  Perfect for modules that may not always be needed, or for
  keeping startup fast while still having WASM available.

  ## Usage

      # Start a lazy instance (no WASM loaded yet)
      {:ok, lazy} = Firebird.Lazy.start_link(wasm: "math.wasm")

      # First call triggers loading, then executes
      {:ok, [8]} = Firebird.Lazy.call(lazy, :add, [5, 3])

      # Subsequent calls use the already-loaded instance
      {:ok, [55]} = Firebird.Lazy.call(lazy, :fibonacci, [10])

      Firebird.Lazy.stop(lazy)

  ## Supervision Tree

      children = [
        {Firebird.Lazy, wasm: "priv/wasm/math.wasm", name: :lazy_math}
      ]

      # Later:
      {:ok, [8]} = Firebird.Lazy.call(:lazy_math, :add, [5, 3])
  """

  use GenServer

  defstruct [:wasm_source, :opts, :instance, :loaded?]

  @typedoc "Lazy loader state."
  @type t :: %__MODULE__{
          wasm_source: String.t() | binary(),
          opts: keyword(),
          instance: pid() | nil,
          loaded?: boolean()
        }

  @typedoc "Options for starting a lazy loader."
  @type option ::
          {:wasm, String.t() | binary()}
          | {:name, GenServer.name()}
          | {:wasi, boolean()}
          | {:preload, boolean()}

  @doc """
  Start a lazy WASM loader.

  ## Options

  - `:wasm` - Path to WASM file or binary bytes (required)
  - `:name` - Optional GenServer name
  - `:wasi` - Enable WASI support
  - `:preload` - Load immediately instead of lazily (default: false)
  """
  @spec start_link([option()]) :: GenServer.on_start()
  def start_link(opts) do
    {name, opts} = Keyword.pop(opts, :name)
    gen_opts = if name, do: [name: name], else: []

    # Trap exits so that {:stop, reason} from init doesn't kill the caller
    old_trap = Process.flag(:trap_exit, true)

    try do
      GenServer.start_link(__MODULE__, opts, gen_opts)
    after
      Process.flag(:trap_exit, old_trap)

      receive do
        {:EXIT, _, _} -> :ok
      after
        0 -> :ok
      end
    end
  end

  @spec child_spec([option()]) :: Supervisor.child_spec()
  def child_spec(opts) do
    name = Keyword.get(opts, :name, __MODULE__)

    %{
      id: name,
      start: {__MODULE__, :start_link, [opts]},
      type: :worker,
      restart: :permanent
    }
  end

  @doc "Call a WASM function, loading the instance if needed."
  @spec call(GenServer.server(), atom() | String.t(), list()) :: {:ok, list()} | {:error, term()}
  def call(server, function, args) do
    GenServer.call(server, {:call, function, args})
  end

  @doc "Call a WASM function, raising on error."
  @spec call!(GenServer.server(), atom() | String.t(), list()) :: list()
  def call!(server, function, args) do
    case call(server, function, args) do
      {:ok, result} -> result
      {:error, reason} -> raise "Lazy WASM call failed: #{inspect(reason)}"
    end
  end

  @doc "Call a WASM function and return an unwrapped single value."
  @spec call_one(GenServer.server(), atom() | String.t(), list()) ::
          {:ok, term()} | {:error, term()}
  def call_one(server, function, args) do
    case call(server, function, args) do
      {:ok, [single]} -> {:ok, single}
      {:ok, multi} -> {:ok, multi}
      error -> error
    end
  end

  @doc "Call a WASM function and return an unwrapped single value, raising on error."
  @spec call_one!(GenServer.server(), atom() | String.t(), list()) :: term()
  def call_one!(server, function, args) do
    case call_one(server, function, args) do
      {:ok, result} -> result
      {:error, reason} -> raise "Lazy WASM call failed: #{inspect(reason)}"
    end
  end

  @doc "Check if the WASM instance has been loaded."
  @spec loaded?(GenServer.server()) :: boolean()
  def loaded?(server) do
    GenServer.call(server, :loaded?)
  end

  @doc "Force-load the WASM instance (useful for warming up)."
  @spec preload(GenServer.server()) :: :ok | {:error, term()}
  def preload(server) do
    GenServer.call(server, :preload)
  end

  @doc "Stop the lazy loader and underlying WASM instance."
  @spec stop(GenServer.server()) :: :ok
  def stop(server) do
    GenServer.stop(server, :normal)
  end

  # GenServer callbacks

  @impl true
  def init(opts) do
    wasm_source = Keyword.fetch!(opts, :wasm)
    wasm_opts = Keyword.drop(opts, [:wasm, :preload])
    preload = Keyword.get(opts, :preload, false)

    state = %__MODULE__{
      wasm_source: wasm_source,
      opts: wasm_opts,
      instance: nil,
      loaded?: false
    }

    if preload do
      case do_load(state) do
        {:ok, state} -> {:ok, state}
        {:error, reason} -> {:stop, reason}
      end
    else
      {:ok, state}
    end
  end

  @impl true
  def handle_call({:call, function, args}, _from, state) do
    case ensure_loaded(state) do
      {:ok, state} ->
        result = Firebird.call(state.instance, function, args)
        {:reply, result, state}

      {:error, reason} ->
        {:reply, {:error, reason}, state}
    end
  end

  def handle_call(:loaded?, _from, state) do
    {:reply, state.loaded?, state}
  end

  def handle_call(:preload, _from, state) do
    case ensure_loaded(state) do
      {:ok, state} -> {:reply, :ok, state}
      {:error, reason} -> {:reply, {:error, reason}, state}
    end
  end

  @impl true
  def terminate(_reason, %{instance: pid}) when is_pid(pid) do
    if Process.alive?(pid), do: Firebird.stop(pid)
    :ok
  end

  def terminate(_reason, _state), do: :ok

  defp ensure_loaded(%{loaded?: true} = state), do: {:ok, state}

  defp ensure_loaded(state) do
    do_load(state)
  end

  defp do_load(state) do
    case Firebird.load(state.wasm_source, state.opts) do
      {:ok, pid} ->
        {:ok, %{state | instance: pid, loaded?: true}}

      {:error, reason} ->
        {:error, reason}
    end
  end
end