package resource

import (
	"context"
	"fmt"
	"time"

	pkgErrors "github.com/pkg/errors"
	promapi "github.com/prometheus/client_golang/api"
	promv1 "github.com/prometheus/client_golang/api/prometheus/v1"
	prommodel "github.com/prometheus/common/model"

	apierror "github.com/pingcap/test-infra/common/errors"
	"github.com/pingcap/test-infra/common/model"
	"github.com/pingcap/test-infra/common/model/rms/request"
	"github.com/pingcap/test-infra/common/model/rms/response"
	"github.com/pingcap/test-infra/services/rms/config"
	"github.com/pingcap/test-infra/services/rms/utils"
)

const maxAllocatableRate = 0.8 // assume less usable resource as a buffer for TPS

type StatisticsService interface {
	GetResourcePoolStatus(ctx utils.Context, req request.GetResourcePoolStatusRequest) (*response.GetResourcePoolStatusResponse, *apierror.StatusError)
	GetTestbedsResource(ctx utils.Context, req request.GetTestbedsResourceRequest) (*response.GetTestbedsResourceResponse, *apierror.StatusError)
	GetTestbedResourceOvertime(ctx utils.Context, req request.GetTestbedResourceOvertimeRequest) (*response.GetTestbedResourceOvertimeResponse, *apierror.StatusError)
}

var _ StatisticsService = &PrometheusStatistics{}

type PrometheusStatistics struct {
	promAPI promv1.API
}

func NewDefaultStatisticsService() StatisticsService {
	promClient, err := promapi.NewClient(promapi.Config{Address: config.DefaultRMSConfiguration.StatisticsService.PrometheusURL})
	if err != nil {
		panic(fmt.Errorf("NewDefaultStatisticsService: %s", err))
	}

	return &PrometheusStatistics{promAPI: promv1.NewAPI(promClient)}
}

func (s *PrometheusStatistics) GetResourcePoolStatus(ctx utils.Context, req request.GetResourcePoolStatusRequest) (*response.GetResourcePoolStatusResponse, *apierror.StatusError) {
	resp := &response.GetResourcePoolStatusResponse{}
	if config.DefaultRMSConfiguration.StatisticsService.StaticResourcePoolMetrics.UseStaticMetrics {
		staticMetrics := config.DefaultRMSConfiguration.StatisticsService.StaticResourcePoolMetrics
		resp.TotalResource.CPU.Core = staticMetrics.CPUCores
		resp.TotalResource.Memory.Byte = staticMetrics.MemoryBytes
		return resp, nil
	}

	var query string
	queryExecutor := shortCircuitQueryExecutor{promAPI: s.promAPI}
	{ // query prometheus
		query = fmt.Sprintf(`%.2f * sum(kube_node_status_allocatable{resource="%s"} and on (node) kube_node_labels{label_resourcepool="%s"} and on(node) available_node)`, maxAllocatableRate, "cpu", req.ResourcePool)
		resp.TotalResource.CPU.Core = float64(queryExecutor.QuerySingle(ctx.Request().Context(), query).Value)

		query = fmt.Sprintf(`%.2f * sum(kube_node_status_allocatable{resource="%s"} and on (node) kube_node_labels{label_resourcepool="%s"} and on(node) available_node)`, maxAllocatableRate, "memory", req.ResourcePool)
		resp.TotalResource.Memory.Byte = int(queryExecutor.QuerySingle(ctx.Request().Context(), query).Value)

		query = fmt.Sprintf(`sum(kube_pod_container_resource_requests{resource="%s"} and on(node) kube_node_labels{label_resourcepool="%s"} unless on(namespace) living_testbed)`, "cpu", req.ResourcePool)
		resp.NonTestbedOccupation.CPU.Core = float64(queryExecutor.QuerySingle(ctx.Request().Context(), query).Value)

		query = fmt.Sprintf(`sum(kube_pod_container_resource_requests{resource="%s"} and on(node) kube_node_labels{label_resourcepool="%s"} unless on(namespace) living_testbed)`, "memory", req.ResourcePool)
		resp.NonTestbedOccupation.Memory.Byte = int(queryExecutor.QuerySingle(ctx.Request().Context(), query).Value)
	}

	if queryExecutor.Error() != nil {
		return nil, apierror.NewInternalError(fmt.Errorf("query prometheus: %s", queryExecutor.Error()))
	}

	return resp, nil
}

func (s *PrometheusStatistics) GetTestbedsResource(ctx utils.Context, req request.GetTestbedsResourceRequest) (*response.GetTestbedsResourceResponse, *apierror.StatusError) {
	resp := &response.GetTestbedsResourceResponse{TestbedsOccupation: make(map[string]*model.PhysicalResource)}

	var query string
	var testbed string
	queryExecutor := shortCircuitQueryExecutor{promAPI: s.promAPI}
	{ // query prometheus
		query = fmt.Sprintf(`sum(kube_pod_container_resource_requests{resource="%s"} and on (node) kube_node_labels{label_resourcepool="%s"} and on(namespace) living_testbed) by(namespace)`, "cpu", req.ResourcePool)
		result := queryExecutor.Query(ctx.Request().Context(), query)
		for _, r := range result {
			testbed = string(r.Metric["namespace"])
			if resp.TestbedsOccupation[testbed] == nil {
				resp.TestbedsOccupation[testbed] = &model.PhysicalResource{}
			}

			resp.TestbedsOccupation[testbed].CPU.Core = float64(r.Value)
		}

		query = fmt.Sprintf(`sum(kube_pod_container_resource_requests{resource="%s"} and on (node) kube_node_labels{label_resourcepool="%s"} and on(namespace) living_testbed) by(namespace)`, "memory", req.ResourcePool)
		result = queryExecutor.Query(ctx.Request().Context(), query)
		for _, r := range result {
			testbed = string(r.Metric["namespace"])
			if resp.TestbedsOccupation[testbed] == nil {
				resp.TestbedsOccupation[testbed] = &model.PhysicalResource{}
			}

			resp.TestbedsOccupation[testbed].Memory.Byte = int(r.Value)
		}
	}

	if queryExecutor.Error() != nil {
		return nil, apierror.NewInternalError(fmt.Errorf("query prometheus: %s", queryExecutor.Error()))
	}

	return resp, nil
}

func (s *PrometheusStatistics) GetTestbedResourceOvertime(ctx utils.Context, req request.GetTestbedResourceOvertimeRequest) (*response.GetTestbedResourceOvertimeResponse, *apierror.StatusError) {
	resp := &response.GetTestbedResourceOvertimeResponse{}

	queryDuration := req.End.Sub(req.Start).Round(time.Second).String()
	queryOffset := time.Since(req.End).Round(time.Second).String()
	if queryOffset == "0s" {
		queryOffset = ""
	} else {
		queryOffset = "offset " + queryOffset
	}

	var query string
	queryExecutor := shortCircuitQueryExecutor{promAPI: s.promAPI}
	{ // query prometheus
		query = fmt.Sprintf(`avg_over_time((sum(kube_pod_container_resource_requests{resource="%s"} and on(namespace) living_testbed {namespace="%s"}))[%s:30s] %s)`,
			"cpu", req.Testbed, queryDuration, queryOffset)
		resp.AverageTestBedPhysicalResource.CPU.Core = float64(queryExecutor.QuerySingle(ctx.Request().Context(), query).Value)

		query = fmt.Sprintf(`max_over_time((sum(kube_pod_container_resource_requests{resource="%s"} and on(namespace) living_testbed {namespace="%s"}))[%s:30s] %s)`,
			"cpu", req.Testbed, queryDuration, queryOffset)
		resp.MaxTestBedPhysicalResource.CPU.Core = float64(queryExecutor.QuerySingle(ctx.Request().Context(), query).Value)

		query = fmt.Sprintf(`avg_over_time((sum(kube_pod_container_resource_requests{resource="%s"} and on(namespace) living_testbed {namespace="%s"}))[%s:30s] %s)`,
			"memory", req.Testbed, queryDuration, queryOffset)
		resp.AverageTestBedPhysicalResource.Memory.Byte = int(queryExecutor.QuerySingle(ctx.Request().Context(), query).Value)

		query = fmt.Sprintf(`max_over_time((sum(kube_pod_container_resource_requests{resource="%s"} and on(namespace) living_testbed {namespace="%s"}))[%s:30s] %s)`,
			"memory", req.Testbed, queryDuration, queryOffset)
		resp.MaxTestBedPhysicalResource.Memory.Byte = int(queryExecutor.QuerySingle(ctx.Request().Context(), query).Value)
	}

	if queryExecutor.Error() != nil {
		return nil, apierror.NewInternalError(fmt.Errorf("query prometheus: %s", queryExecutor.Error()))
	}

	return resp, nil
}

// shortCircuitQueryExecutor is used to execute multiple queries in a row
// following queries won't run(i.e. short-circuited) if previous queries returned error
type shortCircuitQueryExecutor struct {
	err     error
	promAPI promv1.API
}

func (s *shortCircuitQueryExecutor) Error() error {
	return s.err
}

// QuerySingle is a wrapper for Query, it runs query with exactly 1 value returned
// value count more than 1 will be regarded as error
func (s *shortCircuitQueryExecutor) QuerySingle(ctx context.Context, query string) prommodel.Sample {
	queryResult := s.Query(ctx, query)
	if s.err != nil {
		return prommodel.ZeroSample
	}

	switch len(queryResult) {
	case 0:
		return prommodel.ZeroSample
	case 1:
		return *queryResult[0]
	default:
		s.err = pkgErrors.Errorf("expect result length 1, got %d, query: %s", len(queryResult), query)
		return prommodel.ZeroSample
	}
}

func (s *shortCircuitQueryExecutor) Query(ctx context.Context, query string) []*prommodel.Sample {
	if s.err != nil {
		return nil
	}

	queryResult, err := s.do(ctx, query)
	if err != nil {
		s.err = pkgErrors.Wrapf(err, "query: %s", query)
	}
	return queryResult
}

func (s *shortCircuitQueryExecutor) do(ctx context.Context, query string) ([]*prommodel.Sample, error) {
	result, _, err := s.promAPI.Query(ctx, query, time.Now())
	if err != nil {
		return nil, err
	}

	switch r := result.(type) {
	case *prommodel.Scalar:
		return extractPromScalarValue(r)
	case prommodel.Vector:
		return r, nil
	default:
		return nil, fmt.Errorf("unsupported prometheus type: %s", r.Type().String())
	}
}

func extractPromScalarValue(record *prommodel.Scalar) ([]*prommodel.Sample, error) {
	if record == nil {
		return nil, fmt.Errorf("nil scalar record")
	}

	return []*prommodel.Sample{{Value: record.Value, Timestamp: record.Timestamp}}, nil
}