ch

defmodule Ch.Connection do
  @moduledoc false
  use DBConnection
  require Logger
  alias Ch.{Error, Query, Result, RowBinary}
  alias Mint.HTTP1, as: HTTP

  @user_agent "ch/" <> Mix.Project.config()[:version]

  @typep conn :: HTTP.t()

  @impl true
  @spec connect([Ch.start_option()]) :: {:ok, conn} | {:error, Error.t() | Mint.Types.error()}
  def connect(opts) do
    scheme = String.to_existing_atom(opts[:scheme] || "http")
    address = opts[:hostname] || "localhost"
    port = opts[:port] || 8123
    mint_opts = [mode: :passive] ++ Keyword.take(opts, [:hostname, :transport_opts])

    with {:ok, conn} <- HTTP.connect(scheme, address, port, mint_opts) do
      conn =
        conn
        |> HTTP.put_private(:timeout, opts[:timeout] || :timer.seconds(15))
        |> maybe_put_private(:database, opts[:database])
        |> maybe_put_private(:username, opts[:username])
        |> maybe_put_private(:password, opts[:password])
        |> maybe_put_private(:settings, opts[:settings])

      handshake = Query.build("select 1, version()")
      params = DBConnection.Query.encode(handshake, _params = [], _opts = [])

      case handle_execute(handshake, params, _opts = [], conn) do
        {:ok, handshake, responses, conn} ->
          case DBConnection.Query.decode(handshake, responses, _opts = []) do
            %Result{rows: [[1, version]]} ->
              conn =
                if parse_version(version) >= parse_version("24.10") do
                  settings =
                    HTTP.get_private(conn, :settings, [])
                    |> Keyword.put_new(:input_format_binary_read_json_as_string, 1)
                    |> Keyword.put_new(:output_format_binary_write_json_as_string, 1)

                  HTTP.put_private(conn, :settings, settings)
                else
                  conn
                end

              {:ok, conn}

            result ->
              {:ok, _conn} = HTTP.close(conn)
              reason = Error.exception("unexpected result for '#{handshake}': #{inspect(result)}")
              {:error, reason}
          end

        {:error, reason, conn} ->
          {:ok, _conn} = HTTP.close(conn)
          {:error, reason}

        {disconnect, reason, conn} when disconnect in [:disconnect, :disconnect_and_retry] ->
          {:ok, _conn} = HTTP.close(conn)
          {:error, reason}
      end
    end
  catch
    _kind, reason -> {:error, reason}
  end

  defp parse_version(version) do
    version
    |> String.split(".")
    |> Enum.flat_map(fn segment ->
      case Integer.parse(segment) do
        {int, _rest} -> [int]
        :error -> []
      end
    end)
  end

  @impl true
  @spec ping(conn) :: {:ok, conn} | {:disconnect, Mint.Types.error() | Error.t(), conn}
  def ping(conn) do
    headers = [{"user-agent", @user_agent}]

    case request(conn, "GET", "/ping", headers, _body = "", _opts = []) do
      {:ok, conn, _response} -> {:ok, conn}
      {:error, error, conn} -> {:disconnect, error, conn}
      {:disconnect, _error, _conn} = disconnect -> disconnect
    end
  end

  @impl true
  @spec checkout(conn) :: {:ok, conn}
  def checkout(conn), do: {:ok, conn}

  # we "support" these four tx callbacks for Repo.checkout
  # even though ClickHouse doesn't support txs

  @impl true
  def handle_begin(_opts, conn), do: {:ok, %{}, conn}
  @impl true
  def handle_commit(_opts, conn), do: {:ok, %{}, conn}
  @impl true
  def handle_rollback(_opts, conn), do: {:ok, %{}, conn}
  @impl true
  def handle_status(_opts, conn), do: {:idle, conn}

  @impl true
  def handle_prepare(_query, _opts, conn) do
    {:error, Error.exception("prepared statements are not supported"), conn}
  end

  @impl true
  def handle_close(_query, _opts, conn) do
    {:error, Error.exception("prepared statements are not supported"), conn}
  end

  @impl true
  def handle_declare(query, params, opts, conn) do
    %Query{command: command, decode: decode} = query
    {query_params, extra_headers, body} = params

    path = path(conn, query_params, opts)
    headers = headers(conn, extra_headers, opts)
    timeout = timeout(conn, opts)

    with {:ok, conn, _ref} <- send_request(conn, "POST", path, headers, body),
         {:ok, conn, columns, headers, reader} <- recv_declare(conn, decode, timeout) do
      result = %Result{
        command: command,
        columns: columns,
        rows: [],
        num_rows: 0,
        headers: headers,
        data: []
      }

      {:ok, query, result, {conn, reader}}
    else
      {:error, _reason, _conn} = client_error -> client_error
      {:disconnect, reason, conn} -> {:disconnect_and_retry, reason, conn}
    end
  end

  defp recv_declare(conn, decode, timeout) do
    acc = %{decode: decode, step: :status, buffer: [], headers: []}
    recv_declare_continue(conn, acc, timeout)
  end

  defp recv_declare_continue(conn, acc, timeout) do
    case HTTP.recv(conn, 0, timeout) do
      {:ok, conn, responses} ->
        case handle_recv_declare(responses, acc) do
          {:ok, columns, headers, reader} ->
            {:ok, conn, columns, headers, reader}

          {:more, acc} ->
            recv_declare_continue(conn, acc, timeout)

          :error ->
            all_responses_result =
              case handle_all_responses(responses, []) do
                {:ok, responses} -> {:ok, conn, responses}
                {:more, acc} -> recv_all(conn, acc, timeout)
              end

            with {:ok, conn, responses} <- all_responses_result do
              [_status, headers | data] = responses
              message = IO.iodata_to_binary(data)

              code =
                if code = get_header(headers, "x-clickhouse-exception-code") do
                  String.to_integer(code)
                end

              {:error, Error.exception(code: code, message: message), conn}
            end
        end

      {:error, conn, error, _responses} ->
        {:disconnect, error, conn}
    end
  end

  defp handle_recv_declare([{:status, _ref, status} | responses], %{step: :status} = acc) do
    case status do
      200 -> handle_recv_declare(responses, %{acc | step: :headers})
      _other -> :error
    end
  end

  defp handle_recv_declare([{:headers, _ref, headers} | responses], %{step: :headers} = acc) do
    with %{decode: true} <- acc,
         "RowBinaryWithNamesAndTypes" <- get_header(headers, "x-clickhouse-format") do
      handle_recv_declare(responses, %{acc | headers: headers, step: :columns})
    else
      _ ->
        reader = %{decode: false, responses: responses}
        {:ok, _columns = nil, headers, reader}
    end
  end

  defp handle_recv_declare([{:data, _ref, data} | responses], %{step: :columns} = acc) do
    buffer = maybe_concat_buffer(acc.buffer, data)

    case RowBinary.decode_header(buffer) do
      {:ok, names, types, buffer} ->
        reader = %{buffer: buffer, types: types, state: nil, responses: responses}
        {:ok, names, acc.headers, reader}

      :more ->
        handle_recv_declare(responses, %{acc | buffer: buffer})
    end
  end

  defp handle_recv_declare([], acc), do: {:more, acc}

  @compile inline: [maybe_concat_buffer: 2]
  defp maybe_concat_buffer("", data), do: data
  defp maybe_concat_buffer(buffer, data) when is_binary(buffer), do: buffer <> data
  defp maybe_concat_buffer([], data), do: data

  @impl true
  def handle_fetch(query, %Result{} = result, opts, {conn, reader}) do
    case reader do
      %{responses: []} ->
        handle_fetch_recv(query, result, opts, conn, reader)

      %{decode: false, responses: responses} ->
        case responses do
          [{:data, _ref, data} | responses] ->
            result = %Result{result | data: data}
            reader = %{reader | responses: responses}
            {:cont, result, {conn, reader}}

          [{:done, _ref}] ->
            reader = %{reader | responses: []}
            {:halt, result, {conn, reader}}
        end

      %{buffer: buffer, types: types, state: state, responses: responses} ->
        case responses do
          [{:data, _ref, data} | responses] ->
            buffer = maybe_concat_buffer(buffer, data)
            {rows, buffer, state} = RowBinary.decode_rows_continue(buffer, types, state)
            result = %Result{result | data: data, rows: rows, num_rows: length(rows)}
            reader = %{reader | buffer: buffer, state: state, responses: responses}
            {:cont, result, {conn, reader}}

          [{:done, _ref}] ->
            reader = %{reader | responses: []}
            {:halt, result, {conn, reader}}
        end
    end
  end

  defp handle_fetch_recv(query, result, opts, conn, reader) do
    timeout = timeout(conn, opts)

    case HTTP.recv(conn, 0, timeout) do
      {:ok, conn, responses} ->
        reader = %{reader | responses: responses}
        handle_fetch(query, result, opts, {conn, reader})

      {:error, conn, reason, _responses} ->
        {:disconnect, reason, conn}
    end
  end

  @impl true
  def handle_deallocate(_query, %Result{} = result, _opts, {conn, _reader}) do
    case HTTP.open_request_count(conn) do
      0 ->
        {:ok, %{result | data: []}, conn}

      1 ->
        error =
          Error.exception("stopping stream before receiving full response by closing connection")

        {:disconnect, error, conn}
    end
  end

  @impl true
  def handle_execute(%Query{} = query, {:stream, params}, opts, conn) do
    {query_params, extra_headers, body} = params

    path = path(conn, query_params, opts)
    headers = headers(conn, extra_headers, opts)

    with {:ok, conn, ref} <- send_request(conn, "POST", path, headers, :stream) do
      case HTTP.stream_request_body(conn, ref, body) do
        {:ok, conn} -> {:ok, query, ref, conn}
        {:error, conn, reason} -> {:disconnect_and_retry, reason, conn}
      end
    end
  end

  def handle_execute(%Query{} = query, {:stream, ref, body}, opts, conn) do
    case HTTP.stream_request_body(conn, ref, body) do
      {:ok, conn} ->
        case body do
          :eof ->
            with {:ok, conn, responses} <- receive_full_response(conn, timeout(conn, opts)) do
              {:ok, query, responses, conn}
            end

          _other ->
            {:ok, query, ref, conn}
        end

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

  def handle_execute(%Query{command: :insert} = query, params, opts, conn) do
    {query_params, extra_headers, body} = params

    path = path(conn, query_params, opts)
    headers = headers(conn, extra_headers, opts)

    result =
      if is_function(body, 2) do
        request_chunked(conn, "POST", path, headers, body, opts)
      else
        request(conn, "POST", path, headers, body, opts)
      end

    case result do
      {:ok, conn, responses} -> {:ok, query, responses, conn}
      {:error, _reason, _conn} = client_error -> client_error
      {:disconnect, reason, conn} -> {:disconnect_and_retry, reason, conn}
    end
  end

  def handle_execute(query, params, opts, conn) do
    {query_params, extra_headers, body} = params

    path = path(conn, query_params, opts)
    headers = headers(conn, extra_headers, opts)

    case request(conn, "POST", path, headers, body, opts) do
      {:ok, conn, responses} -> {:ok, query, responses, conn}
      {:error, _reason, _conn} = client_error -> client_error
      {:disconnect, reason, conn} -> {:disconnect_and_retry, reason, conn}
    end
  end

  @impl true
  def disconnect(error, {conn, _reader}) do
    disconnect(error, conn)
  end

  def disconnect(_error, conn) do
    {:ok = ok, _conn} = HTTP.close(conn)
    ok
  end

  @typep response :: Mint.Types.status() | Mint.Types.headers() | binary

  @spec request(conn, binary, binary, Mint.Types.headers(), iodata, [Ch.query_option()]) ::
          {:ok, conn, [response]}
          | {:error, Error.t(), conn}
          | {:disconnect, Mint.Types.error(), conn}
  defp request(conn, method, path, headers, body, opts) do
    with {:ok, conn, _ref} <- send_request(conn, method, path, headers, body) do
      receive_full_response(conn, timeout(conn, opts))
    end
  end

  @spec request_chunked(conn, binary, binary, Mint.Types.headers(), Enumerable.t(), Keyword.t()) ::
          {:ok, conn, [response]}
          | {:error, Error.t(), conn}
          | {:disconnect, Mint.Types.error(), conn}
  def request_chunked(conn, method, path, headers, stream, opts) do
    with {:ok, conn, ref} <- send_request(conn, method, path, headers, :stream),
         {:ok, conn} <- stream_body(conn, ref, stream),
         do: receive_full_response(conn, timeout(conn, opts))
  end

  @spec stream_body(conn, Mint.Types.request_ref(), Enumerable.t()) ::
          {:ok, conn} | {:disconnect, Mint.Types.error(), conn}
  defp stream_body(conn, ref, stream) do
    result =
      stream
      |> Stream.concat([:eof])
      |> Enum.reduce_while({:ok, conn}, fn
        chunk, {:ok, conn} -> {:cont, HTTP.stream_request_body(conn, ref, chunk)}
        _chunk, {:error, _conn, _reason} = error -> {:halt, error}
      end)

    case result do
      {:ok, _conn} = ok -> ok
      {:error, conn, reason} -> {:disconnect, reason, conn}
    end
  end

  # stacktrace is a bit cleaner with this function inlined
  @compile inline: [send_request: 5]
  defp send_request(conn, method, path, headers, body) do
    case HTTP.request(conn, method, path, headers, body) do
      {:ok, _conn, _ref} = ok -> ok
      {:error, conn, reason} -> {:disconnect, reason, conn}
    end
  end

  @spec receive_full_response(conn, timeout) ::
          {:ok, conn, [response]}
          | {:error, Error.t(), conn}
          | {:disconnect, Mint.Types.error(), conn}
  defp receive_full_response(conn, timeout) do
    with {:ok, conn, responses} <- recv_all(conn, [], timeout) do
      case responses do
        [200, headers | _rest] ->
          conn = ensure_same_server(conn, headers)
          {:ok, conn, responses}

        [_status, headers | data] ->
          message = IO.iodata_to_binary(data)

          code =
            if code = get_header(headers, "x-clickhouse-exception-code") do
              String.to_integer(code)
            end

          {:error, Error.exception(code: code, message: message), conn}
      end
    end
  end

  @spec recv_all(conn, [response], timeout()) ::
          {:ok, conn, [response]} | {:disconnect, Mint.Types.error(), conn}
  defp recv_all(conn, acc, timeout) do
    case HTTP.recv(conn, 0, timeout) do
      {:ok, conn, responses} ->
        case handle_all_responses(responses, acc) do
          {:ok, responses} -> {:ok, conn, responses}
          {:more, acc} -> recv_all(conn, acc, timeout)
        end

      {:error, conn, reason, _responses} ->
        {:disconnect, reason, conn}
    end
  end

  for tag <- [:data, :status, :headers] do
    defp handle_all_responses([{unquote(tag), _ref, data} | rest], acc) do
      handle_all_responses(rest, [data | acc])
    end
  end

  defp handle_all_responses([{:done, _ref}], acc), do: {:ok, :lists.reverse(acc)}
  defp handle_all_responses([], acc), do: {:more, acc}

  defp maybe_put_private(conn, _k, nil), do: conn
  defp maybe_put_private(conn, k, v), do: HTTP.put_private(conn, k, v)

  defp timeout(conn), do: HTTP.get_private(conn, :timeout)
  defp timeout(conn, opts), do: Keyword.get(opts, :timeout) || timeout(conn)

  defp settings(conn, opts) do
    default_settings = HTTP.get_private(conn, :settings, [])
    opts_settings = Keyword.get(opts, :settings, [])
    Keyword.merge(default_settings, opts_settings)
  end

  defp headers(conn, extra_headers, opts) do
    extra_headers
    |> maybe_put_new_header("x-clickhouse-user", get_opts_or_private(conn, opts, :username))
    |> maybe_put_new_header("x-clickhouse-key", get_opts_or_private(conn, opts, :password))
    |> maybe_put_new_header("x-clickhouse-database", get_opts_or_private(conn, opts, :database))
    |> maybe_put_new_header("user-agent", @user_agent)
  end

  defp get_opts_or_private(conn, opts, key) do
    Keyword.get(opts, key) || HTTP.get_private(conn, key)
  end

  defp maybe_put_new_header(headers, _name, _no_value = nil), do: headers

  defp maybe_put_new_header(headers, name, value) do
    if List.keymember?(headers, name, 0) do
      headers
    else
      [{name, value} | headers]
    end
  end

  defp get_header(headers, key) do
    case List.keyfind(headers, key, 0) do
      {_, value} -> value
      nil = not_found -> not_found
    end
  end

  defp path(conn, query_params, opts) do
    settings = settings(conn, opts)
    "/?" <> URI.encode_query(settings ++ query_params)
  end

  @server_display_name_key :server_display_name

  @spec ensure_same_server(conn, Mint.Types.headers()) :: conn
  defp ensure_same_server(conn, headers) do
    expected_name = HTTP.get_private(conn, @server_display_name_key)
    actual_name = get_header(headers, "x-clickhouse-server-display-name")

    cond do
      expected_name && actual_name ->
        unless actual_name == expected_name do
          Logger.warning(
            "Server mismatch detected. Expected #{inspect(expected_name)} but got #{inspect(actual_name)}!" <>
              " Connection pooling might be unstable."
          )
        end

        conn

      actual_name ->
        HTTP.put_private(conn, @server_display_name_key, actual_name)

      true ->
        conn
    end
  end
end