package prom

import (
	"context"
	"crypto/tls"
	"fmt"
	"net/http"
	"net/url"
	"time"

	"github.com/pingcap/dbaas/central/config"
	infraapi "github.com/pingcap/dbaas/central/infra-api"
	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	"github.com/prometheus/client_golang/api"
	promAPIv1 "github.com/prometheus/client_golang/api/prometheus/v1"
	"github.com/prometheus/common/model"
	"github.com/tidbcloud/pkgv2/auth/auth0"
	"go.uber.org/zap"
)

// BasePromClient is the client for the prometheus instances who scrape TiDB clusters directly.
// For dedicated TiDB clusters, those instances are in the monitoring pods.
// For dev-tier TiDB clusters, those instances are shared.
type BasePromClient struct {
	promAPI promAPIv1.API
}

const (
	defaultQueryTimeout = time.Second * 30
)

func (q *BasePromClient) Query(ctx context.Context, query string) (model.Value, error) {
	ctx, cancel := context.WithTimeout(ctx, defaultQueryTimeout)
	defer cancel()
	v, _, err := q.promAPI.Query(ctx, query, time.Time{})
	if err != nil {
		return nil, errors.Annotatef(err, "prometheus query `%s` failed", query)
	}
	return v, nil
}

func (q *BasePromClient) QueryVec(ctx context.Context, query string) (model.Vector, error) {
	val, err := q.Query(ctx, query)
	if err != nil {
		return nil, err
	}
	vec, ok := val.(model.Vector)
	if !ok {
		log.Warn("query metric value type mismatch",
			zap.String("query", query),
			zap.String("expect", "vector"),
			zap.Any("received", val.Type()))
		return nil, errors.New("unexpected metric type")
	}
	return vec, nil
}

func (q *BasePromClient) MultiQueryVecWithTimeout(ctx context.Context, timePoint time.Time, query string) (model.Value, error) {
	ctx, cancel := context.WithTimeout(ctx, defaultQueryTimeout)
	defer cancel()
	v, _, err := q.promAPI.Query(ctx, query, timePoint)
	if err != nil {
		return nil, errors.Annotatef(err, "prometheus query `%s` failed", query)
	}
	return v, nil
}

func (q *BasePromClient) MultiQueryVec(ctx context.Context, queries ...string) ([]model.Vector, error) {
	type result struct {
		value model.Vector
		err   error
		index int
	}
	resultsCh := make(chan result, len(queries))
	for idx, query := range queries {
		go func(idx int, query string) {
			vec, err := q.QueryVec(ctx, query)
			resultsCh <- result{
				value: vec,
				err:   err,
				index: idx,
			}
		}(idx, query)
	}
	retVec := make([]model.Vector, len(queries))
	for range queries {
		result := <-resultsCh
		if result.err != nil {
			return nil, result.err
		}
		retVec[result.index] = result.value
	}
	return retVec, nil
}

func NewBasePromClient(projectID, clusterID uint64, api *infraapi.API, proxyCfg config.ProxyConfig, tlsConfig *tls.Config) (*BasePromClient, error) {
	cluster, err := api.GetCluster(context.TODO(), projectID, clusterID)
	if err != nil {
		return nil, err
	}

	if cluster.Status.Metrics == nil {
		return nil, errors.New("cluster metrics endpiont not ready")
	}
	hostEndpoint := cluster.Status.Metrics.BaseURL
	client, err := newHttpClient(hostEndpoint, proxyCfg, tlsConfig)
	if err != nil {
		return nil, err
	}
	promAPI := promAPIv1.NewAPI(client)
	return &BasePromClient{promAPI: promAPI}, nil
}

// httpClient is a wrapper of prometheus `api.Client`.
// It just adds a `Authorization` header to the original http request.
type httpClient struct {
	client      api.Client
	bearerToken string
}

func newHttpClient(hostEndpoint string, proxyCfg config.ProxyConfig, tlsConfig *tls.Config) (api.Client, error) {
	client, err := api.NewClient(api.Config{
		Address: fmt.Sprintf("https://%s", hostEndpoint),
		RoundTripper: &http.Transport{
			TLSClientConfig: tlsConfig,
		},
	})
	if err != nil {
		return nil, err
	}
	token, err := auth0.GetClientToken(proxyCfg.Auth0Domain, proxyCfg.Auth0ClientID, proxyCfg.Auth0ClientSecret, proxyCfg.Auth0Audience)
	if err != nil {
		return nil, err
	}
	return &httpClient{
		client:      client,
		bearerToken: token,
	}, nil
}

func (c *httpClient) URL(ep string, args map[string]string) *url.URL {
	return c.client.URL(ep, args)
}

func (c *httpClient) Do(ctx context.Context, req *http.Request) (*http.Response, []byte, error) {
	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", c.bearerToken))
	return c.client.Do(ctx, req)
}