quackdb

defmodule QuackDB.Server do
  @moduledoc """
  Supervises a local DuckDB Quack server process with MuonTrap.

  This is a convenience for local development, tests, demos, and notebooks. It
  starts the external `duckdb` executable and serves DuckDB's Quack HTTP
  protocol. It is not an embedded DuckDB driver and is not required when your
  Quack server runs elsewhere.

      children =
        QuackDB.Server.child_specs(
          server: [name: MyApp.DuckDB, duckdb: :managed, endpoint: "quack:localhost:9494"],
          client: [name: MyApp.QuackDB, pool_size: 5]
        )

  Use `duckdb: :managed` to download and cache DuckDB's official CLI binary via
  `QuackDB.Binary`. Pass `duckdb: "/path/to/duckdb"` or set
  `QUACKDB_BINARY_PATH` when you want to provide the executable yourself.

  `child_specs/1` generates one shared random token when neither side provides
  `:token`, then injects the same token and URI into the server and client specs.

  By default the server runs DuckDB directly under MuonTrap with `-interactive`
  so the process stays alive after `quack_serve/2` starts:

      duckdb :memory: -interactive -init /dev/null -cmd "LOAD quack; ..."

  """

  use GenServer

  alias QuackDB.Protocol.Codec
  alias QuackDB.Protocol.Message.ConnectionRequest
  alias QuackDB.Protocol.Message.ConnectionResponse
  alias QuackDB.Protocol.Message.Disconnect
  alias QuackDB.Protocol.Message.ErrorResponse

  @ready_check_timeout 1_000

  defstruct [
    :daemon,
    :duckdb,
    :database,
    :endpoint,
    :uri,
    :token,
    :boot_sql,
    :daemon_command,
    :daemon_args,
    :daemon_options
  ]

  @type option ::
          {:name, GenServer.name()}
          | {:duckdb, String.t() | :managed}
          | {:duckdb_options, keyword()}
          | {:database, String.t()}
          | {:endpoint, String.t()}
          | {:uri, String.t()}
          | {:token, String.t()}
          | {:load_quack?, boolean()}
          | {:boot_sql, String.t()}
          | {:settings, keyword(QuackDB.SQL.parameter())}
          | {:global_settings, keyword(QuackDB.SQL.parameter())}
          | {:wait, boolean()}
          | {:wait_timeout, timeout()}
          | {:poll_interval, pos_integer()}
          | {:daemon_options, Keyword.t()}
          | {:daemon_command, {String.t(), [String.t()]}}

  @spec child_specs(keyword()) :: [Supervisor.child_spec()]
  def child_specs(options \\ []) do
    server_options = Keyword.get(options, :server, [])
    client_options = Keyword.get(options, :client, [])
    endpoint = Keyword.get(server_options, :endpoint, "quack:localhost")

    uri =
      Keyword.get(server_options, :uri) || Keyword.get(client_options, :uri) ||
        default_uri(endpoint)

    token =
      Keyword.get(server_options, :token) || Keyword.get(client_options, :token) || random_token()

    server_options =
      server_options
      |> Keyword.put_new(:endpoint, endpoint)
      |> Keyword.put(:uri, uri)
      |> Keyword.put(:token, token)

    client_options = client_options |> Keyword.put(:uri, uri) |> Keyword.put(:token, token)

    [child_spec(server_options), QuackDB.child_spec(client_options)]
  end

  @spec child_spec([option()]) :: Supervisor.child_spec()
  def child_spec(options) do
    %{
      id: Keyword.get(options, :name, __MODULE__),
      start: {__MODULE__, :start_link, [options]},
      type: :worker,
      restart: :permanent,
      shutdown: 5_000
    }
  end

  @spec start_link([option()]) :: GenServer.on_start()
  def start_link(options \\ []) do
    {genserver_options, options} = Keyword.split(options, [:name])
    GenServer.start_link(__MODULE__, options, genserver_options)
  end

  @spec uri(GenServer.server()) :: String.t()
  def uri(server), do: GenServer.call(server, :uri)

  @spec token(GenServer.server()) :: String.t()
  def token(server), do: GenServer.call(server, :token)

  @spec info(GenServer.server()) :: map()
  def info(server), do: GenServer.call(server, :info)

  @spec os_pid(GenServer.server()) :: non_neg_integer() | :error
  def os_pid(server), do: GenServer.call(server, :os_pid)

  @spec statistics(GenServer.server()) :: map()
  def statistics(server), do: GenServer.call(server, :statistics)

  @impl true
  def init(options) do
    state = build_state(options)

    with {:ok, daemon} <- start_daemon(state) do
      state = %{state | daemon: daemon}

      if Keyword.get(options, :wait, true) do
        wait_ready!(
          state,
          Keyword.get(options, :wait_timeout, 5_000),
          Keyword.get(options, :poll_interval, 100)
        )
      end

      {:ok, state}
    end
  end

  @impl true
  def handle_call(:uri, _from, state), do: {:reply, state.uri, state}
  def handle_call(:token, _from, state), do: {:reply, state.token, state}

  def handle_call(:info, _from, state) do
    info = %{
      duckdb: state.duckdb,
      database: state.database,
      endpoint: state.endpoint,
      uri: state.uri,
      token: state.token,
      boot_sql: state.boot_sql,
      os_pid: daemon_os_pid(state.daemon),
      statistics: daemon_statistics(state.daemon)
    }

    {:reply, info, state}
  end

  def handle_call(:os_pid, _from, state), do: {:reply, daemon_os_pid(state.daemon), state}
  def handle_call(:statistics, _from, state), do: {:reply, daemon_statistics(state.daemon), state}

  @impl true
  def terminate(_reason, %{daemon: daemon}) when is_pid(daemon) do
    Process.exit(daemon, :shutdown)
    :ok
  end

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

  defp build_state(options) do
    duckdb = duckdb_path(options)
    database = Keyword.get(options, :database, ":memory:")
    endpoint = Keyword.get(options, :endpoint, "quack:localhost")
    uri = Keyword.get(options, :uri, default_uri(endpoint))
    token = Keyword.get_lazy(options, :token, &random_token/0)
    boot_sql = Keyword.get_lazy(options, :boot_sql, fn -> boot_sql(endpoint, token, options) end)
    daemon_options = daemon_options(options)

    {command, args} =
      Keyword.get_lazy(options, :daemon_command, fn ->
        daemon_command(duckdb, database, boot_sql)
      end)

    %__MODULE__{
      duckdb: duckdb,
      database: database,
      endpoint: endpoint,
      uri: uri,
      token: token,
      boot_sql: boot_sql,
      daemon_command: command,
      daemon_args: args,
      daemon_options: daemon_options
    }
  end

  defp duckdb_path(options) do
    case Keyword.get(options, :duckdb, "duckdb") do
      :managed -> QuackDB.Binary.path!(Keyword.get(options, :duckdb_options, []))
      path -> path
    end
  end

  defp daemon_options(options) do
    options
    |> Keyword.get(:daemon_options, [])
    |> Keyword.put_new(:stderr_to_stdout, true)
    |> Keyword.put_new(:log_prefix, "[quackdb-server] ")
  end

  defp daemon_command(duckdb, database, boot_sql) do
    {duckdb, [database, "-interactive", "-init", "/dev/null", "-cmd", boot_sql]}
  end

  defp start_daemon(state) do
    MuonTrap.Daemon.start_link(state.daemon_command, state.daemon_args, state.daemon_options)
  end

  defp wait_ready!(state, timeout, poll_interval) do
    deadline = System.monotonic_time(:millisecond) + timeout
    do_wait_ready!(state, deadline, poll_interval, nil)
  end

  defp do_wait_ready!(state, deadline, poll_interval, last_error) do
    case check_ready(state) do
      :ok ->
        :ok

      {:error, error} ->
        if System.monotonic_time(:millisecond) >= deadline do
          raise QuackDB.Error.new(
                  :server_start_timeout,
                  "DuckDB Quack server did not become ready",
                  source: :client,
                  metadata: %{last_error: error, uri: state.uri}
                )
        else
          Process.sleep(poll_interval)
          do_wait_ready!(state, deadline, poll_interval, error || last_error)
        end
    end
  end

  defp check_ready(state) do
    with {:ok, uri} <- QuackDB.URI.normalize(state.uri),
         request <- connection_request(state),
         {:ok, response} <- QuackDB.Transport.post(uri, request, timeout: @ready_check_timeout),
         {:ok, {header, body}} <- Codec.decode(response),
         :ok <- ready_response(header, body) do
      disconnect(uri, header.connection_id)
      :ok
    end
  end

  defp connection_request(state) do
    %ConnectionRequest{
      auth_string: state.token,
      client_duckdb_version: "quackdb/server-check",
      client_platform: client_platform()
    }
    |> Codec.encode()
  end

  defp ready_response(_header, %ConnectionResponse{}), do: :ok

  defp ready_response(_header, %ErrorResponse{message: message}) do
    {:error, QuackDB.Error.new(:server_error, message, source: :server)}
  end

  defp ready_response(header, _body) do
    {:error,
     QuackDB.Error.new(:unexpected_message, "expected connection response, got #{header.type}",
       source: :protocol
     )}
  end

  defp disconnect(_uri, nil), do: :ok
  defp disconnect(_uri, ""), do: :ok

  defp disconnect(uri, connection_id) do
    request = Codec.encode(%Disconnect{}, connection_id: connection_id)
    _ignored = QuackDB.Transport.post(uri, request, timeout: @ready_check_timeout)
    :ok
  end

  defp daemon_os_pid(nil), do: :error
  defp daemon_os_pid(pid), do: MuonTrap.Daemon.os_pid(pid)

  defp daemon_statistics(nil), do: %{}
  defp daemon_statistics(pid), do: MuonTrap.Daemon.statistics(pid)

  defp boot_sql(endpoint, token, options) do
    [
      if(Keyword.get(options, :load_quack?, true), do: [QuackDB.SQL.load(:quack), " "], else: []),
      server_settings(options),
      server_global_settings(options),
      QuackDB.SQL.call(:quack_serve, [endpoint], token: token)
    ]
    |> IO.iodata_to_binary()
  end

  defp server_settings(options) do
    options
    |> Keyword.get(:settings, default_settings())
    |> Enum.map(fn {name, value} -> [QuackDB.SQL.set(name, value), " "] end)
  end

  defp server_global_settings(options) do
    options
    |> Keyword.get(:global_settings, default_global_settings())
    |> Enum.map(fn {name, value} -> [QuackDB.SQL.set_global(name, value), " "] end)
  end

  defp default_settings do
    [threads: System.schedulers_online()]
  end

  defp default_global_settings do
    [quack_fetch_batch_chunks: 4]
  end

  defp default_uri(endpoint) do
    case parse_endpoint(endpoint) do
      {:ok, "localhost", port} -> "http://[::1]:#{port}"
      {:ok, host, port} -> "http://#{host}:#{port}"
      :error -> "http://[::1]:9494"
    end
  end

  defp parse_endpoint("quack:" <> rest) do
    case String.split(rest, ":", parts: 2) do
      [host] when host != "" -> {:ok, host, 9494}
      [host, port] when host != "" -> parse_port(host, port)
      _other -> :error
    end
  end

  defp parse_endpoint(_endpoint), do: :error

  defp parse_port(host, port) do
    case Integer.parse(port) do
      {port, ""} when port > 0 -> {:ok, host, port}
      _other -> :error
    end
  end

  defp random_token do
    24 |> :crypto.strong_rand_bytes() |> Base.url_encode64(padding: false)
  end

  defp client_platform do
    :system_architecture
    |> :erlang.system_info()
    |> List.to_string()
  end
end