exsql

defmodule ExSQL.Ecto.Connection do
  @moduledoc false

  use DBConnection

  alias ExSQL.{Database, Executor, FileFormat, Registry}
  alias ExSQL.Ecto.{Error, Query, Result}

  defstruct db: nil,
            database: nil,
            dirty?: false,
            transaction_status: :idle,
            status: :idle,
            journal_mode: :memory,
            synced_version: 0

  @type t :: %__MODULE__{
          db: Database.t(),
          database: String.t() | :memory | nil,
          dirty?: boolean(),
          transaction_status: :idle | :transaction,
          status: :idle | :busy,
          journal_mode: atom() | String.t(),
          synced_version: non_neg_integer() | :no_registry
        }

  @impl true
  def connect(opts) do
    database = Keyword.get(opts, :database, :memory)
    journal_mode = Keyword.get(opts, :journal_mode, :memory)

    # Capture the version *before* reading the file: if a write lands in
    # between, we record the older version and reload on our next read rather
    # than silently trusting a snapshot that is already one commit stale.
    synced_version = Registry.current_version(database)

    with {:ok, db} <- open_database(database) do
      {:ok,
       %__MODULE__{
         db: db,
         database: database,
         journal_mode: journal_mode,
         synced_version: synced_version
       }}
    end
  end

  @impl true
  def disconnect(_err, state) do
    _ = persist(state)
    :ok
  end

  @impl true
  def checkout(%__MODULE__{status: :idle} = state) do
    {:ok, %{state | status: :busy}}
  end

  def checkout(%__MODULE__{status: :busy} = state) do
    {:disconnect, %Error{message: "Database is busy"}, state}
  end

  @impl true
  def ping(state), do: {:ok, state}

  @impl true
  def handle_prepare(%Query{} = query, _opts, state) do
    {:ok, query, state}
  end

  @impl true
  def handle_execute(%Query{} = query, params, _opts, state) do
    execute_query(query, params, state)
  end

  @impl true
  def handle_begin(opts, state) do
    state = maybe_reload(state)
    mode = Keyword.get(opts, :mode, :deferred)

    sql =
      if state.transaction_status == :idle do
        begin_sql(mode)
      else
        "SAVEPOINT exsql_savepoint"
      end

    run_transaction_sql(sql, state, :transaction, :without_query)
  end

  @impl true
  def handle_commit(opts, state) do
    mode = Keyword.get(opts, :mode, :deferred)

    sql =
      if mode == :savepoint do
        "RELEASE SAVEPOINT exsql_savepoint"
      else
        "COMMIT"
      end

    status =
      if outer_transaction_end?(state, mode) do
        :idle
      else
        :transaction
      end

    run_transaction_sql(sql, state, status, :without_query)
  end

  @impl true
  def handle_rollback(opts, state) do
    mode = Keyword.get(opts, :mode, :deferred)

    sql =
      if mode == :savepoint do
        "ROLLBACK TO SAVEPOINT exsql_savepoint; RELEASE SAVEPOINT exsql_savepoint"
      else
        "ROLLBACK"
      end

    status =
      if outer_transaction_end?(state, mode) do
        :idle
      else
        :transaction
      end

    case run_transaction_sql(sql, state, status, :without_query) do
      {:ok, result, state} ->
        if mode == :savepoint do
          {:ok, result, state}
        else
          {:ok, result, %{state | dirty?: false}}
        end

      other ->
        other
    end
  end

  @impl true
  def handle_status(_opts, state), do: {state.transaction_status, state}

  @impl true
  def handle_close(_query, _opts, state), do: {:ok, nil, state}

  @impl true
  def handle_declare(%Query{} = query, params, _opts, state) do
    state = maybe_reload(state)
    statement = IO.iodata_to_binary(query.statement)

    case Executor.run(state.db, statement, params) do
      {:ok, [result], db} ->
        cursor = %{result: Result.from_exsql(result), offset: 0}
        state = %{state | db: db, dirty?: state.dirty? or mutating?(result.command, false)}
        {:ok, query, cursor, state}

      {:ok, results, db} ->
        {:error,
         %Error{
           message: "expected one statement, got #{length(results)}",
           statement: statement
         }, %{state | db: db, dirty?: true}}

      {:error, error, db} ->
        {:error, Error.from_exsql(error, statement), %{state | db: db, dirty?: true}}
    end
  end

  @impl true
  def handle_fetch(_query, %{result: result, offset: offset} = cursor, opts, state) do
    max_rows = Keyword.get(opts, :max_rows, 500)
    rows = Enum.slice(result.rows, offset, max_rows)
    next_offset = offset + length(rows)
    result = %{result | rows: rows, num_rows: length(rows)}

    if next_offset >= length(cursor.result.rows) do
      {:halt, result, state}
    else
      {:cont, result, state}
    end
  end

  @impl true
  def handle_deallocate(_query, _cursor, _opts, state), do: {:ok, nil, state}

  defp execute_query(%Query{} = query, params, state) do
    state = maybe_reload(state)
    statement = IO.iodata_to_binary(query.statement)
    mutating_statement = mutating_statement?(statement)

    case Executor.run(state.db, statement, params) do
      {:ok, [result], db} ->
        state =
          %{state | db: db, dirty?: state.dirty? or mutating?(result.command, mutating_statement)}
          |> maybe_persist_after(result.command, mutating_statement)

        {:ok, query, Result.from_exsql(result), state}

      {:ok, results, db} ->
        {:error,
         %Error{
           message: "expected one statement, got #{length(results)}",
           statement: statement
         }, %{state | db: db, dirty?: true}}

      {:error, error, db} ->
        {:error, Error.from_exsql(error, statement), %{state | db: db, dirty?: true}}
    end
  end

  defp run_transaction_sql(sql, state, status, return_shape) do
    query = Query.build(statement: sql)
    mutating_statement = mutating_statement?(sql)

    case Executor.run(state.db, sql, []) do
      {:ok, [result], db} ->
        command = result.command
        result = Result.from_exsql(result)

        state =
          %{state | db: db, transaction_status: status, dirty?: true}
          |> maybe_persist_after(command, mutating_statement)

        transaction_result(query, result, state, return_shape)

      {:ok, results, db} ->
        result = results |> List.last() |> Result.from_exsql()
        state = %{state | db: db, transaction_status: status, dirty?: true}
        transaction_result(query, result, state, return_shape)

      {:error, error, db} ->
        {:error, Error.from_exsql(error, sql), %{state | db: db, dirty?: true}}
    end
  end

  defp transaction_result(_query, result, state, :without_query), do: {:ok, result, state}

  defp open_database(database) when database in [:memory, ":memory:"] do
    {:ok, Database.new()}
  end

  defp open_database(path) when is_binary(path) do
    if File.exists?(path) do
      case FileFormat.read(path) do
        {:ok, db} -> {:ok, db}
        {:error, message} -> {:error, %Error{message: message}}
      end
    else
      path |> Path.dirname() |> File.mkdir_p!()
      {:ok, Database.new()}
    end
  end

  defp open_database(_database), do: {:ok, Database.new()}

  defp maybe_reload(
         %__MODULE__{database: database, transaction_status: :idle, dirty?: false} = state
       )
       when is_binary(database) and database != ":memory:" do
    current = Registry.current_version(database)

    cond do
      current != :no_registry and current == state.synced_version ->
        # Our in-memory database already reflects the latest committed write
        # (no other connection has committed since), so skip re-parsing the
        # whole file. This is the common single-writer hot path.
        state

      File.exists?(database) ->
        case FileFormat.read(database) do
          {:ok, db} -> %{state | db: db, synced_version: current}
          {:error, _message} -> state
        end

      true ->
        state
    end
  end

  defp maybe_reload(state), do: state

  defp persist(%__MODULE__{database: database, dirty?: true} = state)
       when is_binary(database) and database != ":memory:" do
    case FileFormat.write(state.db, database, journal_mode: state.journal_mode) do
      {:ok, _path} -> {:ok, Registry.bump(database)}
      {:error, message} -> {:error, %Error{message: message}}
    end
  end

  defp persist(_state), do: :ok

  defp maybe_persist_after(state, command, mutating_statement) do
    if state.transaction_status == :idle and mutating?(command, mutating_statement) do
      case persist(state) do
        {:ok, version} -> %{state | dirty?: false, synced_version: version}
        :ok -> %{state | dirty?: false}
        {:error, _error} -> state
      end
    else
      state
    end
  end

  defp mutating?(_command, true), do: true

  defp mutating?(command, false) do
    command not in [
      nil,
      :select,
      :pragma,
      :explain,
      :begin,
      :savepoint,
      :rollback
    ]
  end

  defp mutating_statement?(statement) when is_binary(statement) do
    statement =
      statement
      |> String.trim()
      |> String.upcase()

    Regex.match?(
      ~r/\A(?:WITH\b[\s\S]*?\b(INSERT|UPDATE|DELETE|REPLACE)\b|\b(INSERT|UPDATE|DELETE|REPLACE)\b)/,
      statement
    )
  end

  defp mutating_statement?(_statement), do: false

  defp begin_sql(:immediate), do: "BEGIN IMMEDIATE TRANSACTION"
  defp begin_sql(:exclusive), do: "BEGIN EXCLUSIVE TRANSACTION"
  defp begin_sql(_mode), do: "BEGIN TRANSACTION"

  defp outer_transaction_end?(state, mode) do
    state.transaction_status == :transaction and mode != :savepoint
  end
end