defmodule ReqBigQuery do
  @moduledoc """
  `Req` plugin for [Google BigQuery](https://cloud.google.com/bigquery/docs/reference/rest).

  ReqBigQuery makes it easy to make BigQuery queries. It uses `Goth` for authentication.
  Query results are decoded into the `ReqBigQuery.Result` struct.
  The struct implements the `Table.Reader` protocol and thus can be efficiently traversed by rows or columns.
  """

  alias Req.Request
  alias ReqBigQuery.Result

  @allowed_options ~w(goth default_dataset_id project_id bigquery max_results use_legacy_sql timeout_ms dry_run)a
  @base_url "https://bigquery.googleapis.com/bigquery/v2"
  @max_results 10_000
  @use_legacy_sql false
  @timeout_ms 10_000

  @doc """
  Attaches to Req request.

  ## Request Options

    * `:goth` - Required. The goth server name.

    * `:project_id` - Required. The GCP project id.

    * `:bigquery` - Required. The query to execute. It can be a plain sql string or
      a `{query, params}` tuple, where `query` can contain `?` placeholders and `params`
      is a list of corresponding values.

    * `:default_dataset_id` - Optional. If set, the dataset to assume for any unqualified table
      names in the query. If not set, all table names in the query string must be qualified in the
      format 'datasetId.tableId'.

    * `:max_results` - Optional. Number of rows to be returned by BigQuery in each request (paging).
      The rows Stream can make multiple requests if `num_rows` returned is grather than `:max_results`.
      Defaults to 10000.

    * `:use_legacy_sql` - Optional. Specifies whether to use BigQuery's legacy SQL dialect for this query.
      If set to false, the query will use BigQuery's GoogleSQL: https://cloud.google.com/bigquery/sql-reference/
      The default value is false.

    * `:timeout_ms` - Optional. How long to wait for the query to complete, in milliseconds,
      before the request times out and returns. Note: The call is not guaranteed to wait for
      the specified timeout. If the query takes longer to run than the timeout value, the call
      returns without any results and with the 'jobComplete' flag set to false. The default
      value is 10000 milliseconds (10 seconds).

    * `:dry_run` - Optional. Specifies whether to run the given query in dry run mode.

  If you want to set any of these options when attaching the plugin, pass them as the second argument.

  ## Examples

  With plain query string:

      iex> credentials = File.read!("credentials.json") |> Jason.decode!()
      iex> source = {:service_account, credentials, []}
      iex> {:ok, _} = Goth.start_link(name: MyGoth, source: source, http_client: &Req.request/1)
      iex> project_id = System.fetch_env!("PROJECT_ID")
      iex> query = \"""
      ...> SELECT title, SUM(views) AS views
      ...>   FROM `bigquery-public-data.wikipedia.table_bands`
      ...>  WHERE EXTRACT(YEAR FROM datehour) <= 2021
      ...>  GROUP BY title
      ...>  ORDER BY views DESC
      ...>  LIMIT 10
      ...> \"""
      iex> req = Req.new() |> ReqBigQuery.attach(goth: MyGoth, project_id: project_id)
      iex> res = Req.post!(req, bigquery: query).body
      iex> res
      %ReqBigQuery.Result{
        columns: ["title", "views"],
        job_id: "job_JDDZKquJWkY7x0LlDcmZ4nMQqshb",
        num_rows: 10,
        total_bytes_processed: 18161868216,
        rows: %Stream{}
      }
      iex> Enum.to_list(res.rows)
      [
          ["The_Beatles", 13758950],
          ["Queen_(band)", 12019563],
          ["Pink_Floyd", 9522503],
          ["AC/DC", 8972364],
          ["Led_Zeppelin", 8294994],
          ["Linkin_Park", 8242802],
          ["The_Rolling_Stones", 7825952],
          ["Red_Hot_Chili_Peppers", 7302904],
          ["Fleetwood_Mac", 7199563],
          ["Twenty_One_Pilots", 6970692]
      ]

  With parameterized query:

      iex> credentials = File.read!("credentials.json") |> Jason.decode!()
      iex> source = {:service_account, credentials, []}
      iex> {:ok, _} = Goth.start_link(name: MyGoth, source: source, http_client: &Req.request/1)
      iex> project_id = System.fetch_env!("PROJECT_ID")
      iex> query = \"""
      ...> SELECT EXTRACT(YEAR FROM datehour) AS year, SUM(views) AS views
      ...>   FROM `bigquery-public-data.wikipedia.table_bands`
      ...>  WHERE EXTRACT(YEAR FROM datehour) <= 2021
      ...>    AND title = ?
      ...>  GROUP BY 1
      ...>  ORDER BY views DESC
      ...> \"""
      iex> req = Req.new() |> ReqBigQuery.attach(goth: MyGoth, project_id: project_id)
      iex> res = Req.post!(req, bigquery: {query, ["Linkin_Park"]}).body
      %ReqBigQuery.Result{
        columns: ["year", "views"],
        job_id: "job_GXiJvALNsTAoAOJ39Eg3Mw94XMUQ",
        num_rows: 7,
        total_bytes_processed: 15686357820,
        rows: %Stream{}
      }
      iex> Enum.to_list(res.rows)
      [[2017, 2895889], [2016, 1173359], [2018, 1133770], [2020, 906538], [2015, 860899], [2019, 790747], [2021, 481600]]

  """
  @spec attach(Request.t(), keyword()) :: Request.t()
  def attach(%Request{} = request, options \\ []) do
    checked_options =
      options
      |> Keyword.put_new(:base_url, @base_url)
      |> Keyword.put_new(:max_results, @max_results)
      |> Keyword.put_new(:use_legacy_sql, @use_legacy_sql)
      |> Keyword.put_new(:timeout_ms, @timeout_ms)

    request
    |> Request.prepend_request_steps(bigquery_run: &run/1)
    |> Request.register_options(@allowed_options)
    |> Request.merge_options(checked_options)
  end

  defp run(%Request{options: options} = request) do
    if query = options[:bigquery] do
      goth = options[:goth] || raise ":goth is missing"
      project_id = options[:project_id] || raise ":project_id is missing"
      base_url = options[:base_url]
      token = Goth.fetch!(goth).token
      uri = URI.parse("#{base_url}/projects/#{project_id}/queries")

      json =
        query
        |> build_request_body(options[:default_dataset_id])
        |> Map.put(:maxResults, options[:max_results])
        |> Map.put(:useLegacySql, options[:use_legacy_sql])
        |> Map.put(:timeoutMs, options[:timeout_ms])
        |> Map.put(:dryRun, options[:dry_run] || false)

      %{request | url: uri}
      |> Request.merge_options(auth: {:bearer, token}, json: json)
      |> Request.append_response_steps(bigquery_decode: &decode/1)
    else
      request
    end
  end

  defp build_request_body({query, []}, dataset) when is_binary(query) do
    build_request_body(query, dataset)
  end

  defp build_request_body({query, params}, dataset) when is_binary(query) do
    map = build_request_body(query, dataset)

    query_params =
      for value <- params do
        {value, type} = encode_value(value)
        %{parameterType: %{type: type}, parameterValue: %{value: value}}
      end

    Map.merge(map, %{
      queryParameters: query_params,
      useLegacySql: false,
      parameterMode: "POSITIONAL"
    })
  end

  defp build_request_body(query, dataset) when dataset in ["", nil] do
    %{query: query}
  end

  defp build_request_body(query, dataset) when is_binary(query) do
    %{defaultDataset: %{datasetId: dataset}, query: query}
  end

  defp decode({request, %{status: 200} = response}) do
    {request, update_in(response.body, &decode_body(&1, request.options))}
  end

  defp decode(any), do: any

  defp decode_body(
         %{
           "jobReference" => %{"jobId" => job_id},
           "kind" => "bigquery#queryResponse",
           "rows" => _rows,
           "schema" => %{"fields" => fields},
           "totalRows" => num_rows,
           "totalBytesProcessed" => total_bytes
         } = initial_response,
         request_options
       ) do
    %Result{
      job_id: job_id,
      num_rows: String.to_integer(num_rows),
      total_bytes_processed: String.to_integer(total_bytes),
      rows: initial_response |> rows_stream(request_options) |> decode_rows(fields),
      columns: decode_columns(fields)
    }
  end

  defp decode_body(
         %{
           "jobReference" => %{"jobId" => job_id},
           "kind" => "bigquery#queryResponse",
           "schema" => %{"fields" => fields},
           "totalRows" => num_rows,
           "totalBytesProcessed" => total_bytes
         },
         _request_options
       ) do
    %Result{
      job_id: job_id,
      num_rows: String.to_integer(num_rows),
      total_bytes_processed: String.to_integer(total_bytes),
      rows: [],
      columns: decode_columns(fields)
    }
  end

  defp decode_body(
         %{
           "jobReference" => %{},
           "kind" => "bigquery#queryResponse",
           "schema" => %{"fields" => fields},
           "totalBytesProcessed" => total_bytes
         },
         _request_options
       ) do
    %Result{
      job_id: nil,
      num_rows: 0,
      total_bytes_processed: String.to_integer(total_bytes),
      rows: [],
      columns: decode_columns(fields)
    }
  end

  defp rows_stream(initial_response, request_options) do
    Stream.unfold({:initial, initial_response}, fn
      {:initial, %{"rows" => rows} = initial_body} ->
        {rows, initial_body}

      %{
        "pageToken" => page_token,
        "jobReference" => job_reference
      } ->
        %{"jobId" => job_id, "projectId" => project_id} = job_reference

        job_location = Map.get(job_reference, "location")

        resp = page_request(request_options, project_id, job_id, job_location, page_token)
        {resp.body["rows"], resp.body}

      _end ->
        # last iteration didn't have pageToken
        nil
    end)
    |> Stream.flat_map(& &1)
  end

  defp page_request(options, project_id, job_id, job_location, page_token) do
    uri =
      URI.parse(
        "#{@base_url}/projects/#{project_id}/queries/#{job_id}?maxResults=#{options[:max_results]}&pageToken=#{page_token}"
      )

    uri =
      case job_location do
        nil -> uri
        job_location -> URI.append_query(uri, "location=#{job_location}")
      end

    token = Goth.fetch!(options[:goth]).token

    Req.new(url: uri)
    |> Request.merge_options(auth: {:bearer, token})
    |> Req.get!()
  end

  defp decode_rows(rows, fields) do
    Stream.map(rows, fn %{"f" => columns} ->
      Enum.with_index(columns, fn %{"v" => value}, index ->
        field = Enum.at(fields, index)
        decode_value(value, field)
      end)
    end)
  end

  defp decode_columns(fields) do
    Enum.map(fields, & &1["name"])
  end

  @decimal_types ~w(NUMERIC BIGNUMERIC)

  defp decode_value(nil, _), do: nil
  defp decode_value(%{"v" => value}, field), do: decode_value(value, field)

  @invalid_float_values ["-Infinity", "Infinity", "NaN"]

  defp decode_value(value, %{"type" => "FLOAT"}) when value in @invalid_float_values do
    raise "float value #{inspect(value)} is not supported"
  end

  defp decode_value(values, %{"mode" => "REPEATED"} = field) do
    Enum.map(values, &decode_value(&1, Map.delete(field, "mode")))
  end

  defp decode_value(value, %{"type" => "FLOAT"}), do: String.to_float(value)
  defp decode_value(value, %{"type" => "INTEGER"}), do: String.to_integer(value)

  defp decode_value(value, %{"type" => type}) when type in @decimal_types,
    do: Decimal.new(value)

  defp decode_value("true", %{"type" => "BOOLEAN"}), do: true
  defp decode_value("false", %{"type" => "BOOLEAN"}), do: false

  defp decode_value(value, %{"fields" => fields, "type" => "RECORD"}) do
    decode_record(value, fields)
  end

  defp decode_value(value, %{"type" => "DATE"}), do: Date.from_iso8601!(value)

  defp decode_value(value, %{"type" => "DATETIME"}),
    do: NaiveDateTime.from_iso8601!(value)

  defp decode_value(value, %{"type" => "TIME"}), do: Time.from_iso8601!(value)

  defp decode_value(value, %{"type" => "TIMESTAMP"}) do
    float = String.to_float(value)
    DateTime.from_unix!(round(float * 1_000_000), :microsecond)
  end

  defp decode_value(value, _), do: value

  defp decode_record(values, fields) when is_list(values) do
    Enum.map(values, &decode_record(&1, fields))
  end

  defp decode_record(%{"f" => columns}, fields) do
    for {%{"v" => value}, index} <- Enum.with_index(columns), into: %{} do
      field = Enum.at(fields, index)

      {field["name"], decode_value(value, field)}
    end
  end

  defp encode_value(%DateTime{time_zone: "Etc/UTC"} = datetime) do
    naive_datetime = DateTime.to_naive(datetime)
    {to_string(naive_datetime), "TIMESTAMP"}
  end

  defp encode_value(%Date{} = date), do: {to_string(date), "DATE"}
  defp encode_value(%Time{} = time), do: {to_string(time), "TIME"}
  defp encode_value(%NaiveDateTime{} = timestamp), do: {to_string(timestamp), "DATETIME"}
  defp encode_value(%Decimal{} = decimal), do: {to_string(decimal), "BIGNUMERIC"}

  defp encode_value(value) when is_boolean(value), do: {value, "BOOL"}
  defp encode_value(value) when is_float(value), do: {value, "FLOAT"}
  defp encode_value(value) when is_integer(value), do: {value, "INTEGER"}
  defp encode_value(value), do: {value, "STRING"}
end