package server

import (
	"encoding/base64"
	"fmt"
	"io"
	"mime"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/pingcap/log"
	"github.com/tidbcloud/serverless-metering/pkg/monitor"
	"github.com/tidbcloud/serverless-metering/pkg/store"
	"go.uber.org/zap"

	"github.com/matttproud/golang_protobuf_extensions/pbutil"
	dto "github.com/prometheus/client_model/go"
	"github.com/prometheus/common/expfmt"
	"github.com/prometheus/common/model"
	"github.com/prometheus/common/route"
)

const Base64Suffix = "@base64"

// Push returns an http.Handler which accepts samples over HTTP and stores them
// in the Storage. If replace is true, all metrics for the job and instance
// given by the request are deleted before new ones are stored. If check is
// true, the pushed metrics are immediately checked for consistency (with
// existing metrics and themselves), and an inconsistent push is rejected with
// http.StatusBadRequest.
//
// The returned handler is already instrumented for Prometheus.
func Push(storage store.Storage, forwarder *metricForwarder, replace, check, jobBase64Encoded, enableMetricForwarder bool) func(http.ResponseWriter, *http.Request) {
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var statusCode int
		start := time.Now()
		path := "/metrics"
		defer func() {
			monitor.MetricServerHttpRequestCounter.WithLabelValues(path, strconv.Itoa(statusCode)).Inc()
			monitor.MetricServerHttpDurationHistogram.WithLabelValues(path).Observe(time.Since(start).Seconds())
		}()

		job := route.Param(r.Context(), "job")
		if jobBase64Encoded {
			var err error
			if job, err = decodeBase64(job); err != nil {
				statusCode = http.StatusBadRequest
				http.Error(w, fmt.Sprintf("invalid base64 encoding in job name %q: %v", job, err), statusCode)
				log.Debug("invalid base64 encoding in job name", zap.String("job", job), zap.Error(err))
				return
			}
		}
		labelsString := route.Param(r.Context(), "labels")
		labels, err := splitLabels(labelsString)
		if err != nil {
			statusCode = http.StatusBadRequest
			http.Error(w, err.Error(), statusCode)
			log.Debug("failed to parse URL", zap.String("url", labelsString), zap.Error(err))
			return
		}
		if job == "" {
			statusCode = http.StatusBadRequest
			http.Error(w, "job name is required", statusCode)
			log.Debug("job name is required")
			return
		}
		labels["job"] = job

		var metricFamilies map[string]*dto.MetricFamily
		ctMediatype, ctParams, ctErr := mime.ParseMediaType(r.Header.Get("Content-Type"))
		if ctErr == nil && ctMediatype == "application/vnd.google.protobuf" &&
			ctParams["encoding"] == "delimited" &&
			ctParams["proto"] == "io.prometheus.client.MetricFamily" {
			metricFamilies = map[string]*dto.MetricFamily{}
			for {
				mf := &dto.MetricFamily{}
				if _, err = pbutil.ReadDelimited(r.Body, mf); err != nil {
					if err == io.EOF {
						err = nil
					}
					break
				}
				metricFamilies[mf.GetName()] = mf
			}
		} else {
			// We could do further content-type checks here, but the
			// fallback for now will anyway be the text format
			// version 0.0.4, so just go for it and see if it works.
			var parser expfmt.TextParser
			metricFamilies, err = parser.TextToMetricFamilies(r.Body)
		}
		if err != nil {
			statusCode = http.StatusBadRequest
			http.Error(w, err.Error(), statusCode)
			log.Debug("failed to parse text", zap.String("source", r.RemoteAddr), zap.Error(err))
			return
		}
		if enableMetricForwarder {
			forwarder.forwardMetrics(metricFamilies)
		}
		err = storage.Log(labels, metricFamilies)
		if err != nil {
			statusCode = http.StatusInternalServerError
			http.Error(w, err.Error(), statusCode)
			log.Error("failed to log metrics", zap.Error(err))
			return
		}

		statusCode = http.StatusOK
	})

	return func(w http.ResponseWriter, r *http.Request) {
		handler.ServeHTTP(w, r)
	}
}

// decodeBase64 decodes the provided string using the “Base 64 Encoding with URL
// and Filename Safe Alphabet” (RFC 4648). Padding characters (i.e. trailing
// '=') are ignored.
func decodeBase64(s string) (string, error) {
	b, err := base64.RawURLEncoding.DecodeString(strings.TrimRight(s, "="))
	return string(b), err
}

// splitLabels splits a labels string into a label map mapping names to values.
func splitLabels(labels string) (map[string]string, error) {
	result := map[string]string{}
	if len(labels) <= 1 {
		return result, nil
	}
	components := strings.Split(labels[1:], "/")
	if len(components)%2 != 0 {
		return nil, fmt.Errorf("odd number of components in label string %q", labels)
	}

	for i := 0; i < len(components)-1; i += 2 {
		name, value := components[i], components[i+1]
		trimmedName := strings.TrimSuffix(name, Base64Suffix)
		if !model.LabelNameRE.MatchString(trimmedName) ||
			strings.HasPrefix(trimmedName, model.ReservedLabelPrefix) {
			return nil, fmt.Errorf("improper label name %q", trimmedName)
		}
		if name == trimmedName {
			result[name] = value
			continue
		}
		decodedValue, err := decodeBase64(value)
		if err != nil {
			return nil, fmt.Errorf("invalid base64 encoding for label %s=%q: %v", trimmedName, value, err)
		}
		result[trimmedName] = decodedValue
	}
	return result, nil
}