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)a
  @base_url "https://bigquery.googleapis.com/bigquery/v2"
  @max_results 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.

  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,
        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,
        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)

    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])

      %{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, useLegacySql: false}
  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
         } = initial_response,
         request_options
       ) do
    %Result{
      job_id: job_id,
      num_rows: String.to_integer(num_rows),
      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
         },
         _request_options
       ) do
    %Result{
      job_id: job_id,
      num_rows: String.to_integer(num_rows),
      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" => %{"jobId" => job_id, "projectId" => project_id}
      } ->
        resp = page_request(request_options, project_id, job_id, 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, page_token) do
    uri =
      URI.parse(
        "#{@base_url}/projects/#{project_id}/queries/#{job_id}?maxResults=#{options[:max_results]}&pageToken=#{page_token}"
      )

    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