package bigquery

import (
	"context"
	"fmt"

	"cloud.google.com/go/bigquery"
	"go.uber.org/zap"
	"google.golang.org/api/iterator"
)

type BigQueryProvider interface {
	ExecQuery(ctx context.Context, sql string) (BigQueryRows, error)
}
type BigQueryRows []map[string]bigquery.Value
type BigQuery struct {
	client    *bigquery.Client
	projectID string
	log       *zap.Logger
	location  string
}

// NewBigQuery new a bigquery client
func NewBigQuery(projectID string, location string, log *zap.Logger) (*BigQuery, error) {
	ctx := context.Background()
	client, err := bigquery.NewClient(ctx, projectID)
	if err != nil {
		return nil, fmt.Errorf("bigquery.NewClient: %w", err)
	}
	c := BigQuery{
		client:    client,
		projectID: projectID,
		log:       log,
		location:  location,
	}
	return &c, nil
}

func (t BigQuery) query(ctx context.Context, sql string) (BigQueryRows, error) {
	q := t.client.Query(sql)
	q.Location = t.location

	job, err := q.Run(ctx)
	if err != nil {
		return nil, err
	}
	// wait job
	status, err := job.Wait(ctx)
	if err != nil {
		return nil, err
	}
	// check status
	if err := status.Err(); err != nil {
		return nil, err
	}

	var rows BigQueryRows

	it, err := job.Read(ctx)
	if err != nil {
		return nil, err
	}
	for {
		var row map[string]bigquery.Value
		err2 := it.Next(&row)
		if err2 == iterator.Done {
			break
		}
		if err2 != nil {
			return nil, err2
		}
		rows = append(rows, row)
	}
	return rows, nil
}

func (t BigQuery) ExecQuery(ctx context.Context, sql string) (BigQueryRows, error) {
	rows, err := t.query(ctx, sql)
	if err != nil {
		t.log.Error("sql error", zap.String("sql", sql), zap.Error(err))
		return nil, err
	}
	return rows, err
}