// Package httpquery exposes the embedded Prometheus TSDB through a read-only,
// Prometheus-compatible HTTP query API.
package httpquery

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"math"
	"net/http"
	"runtime"
	"slices"
	"strconv"
	"strings"
	"time"

	commonmodel "github.com/prometheus/common/model"
	"github.com/prometheus/common/promslog"
	"github.com/prometheus/prometheus/model/labels"
	"github.com/prometheus/prometheus/model/timestamp"
	promengine "github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/promql/parser"
	"github.com/prometheus/prometheus/storage"
	"github.com/prometheus/prometheus/util/annotations"
)

const (
	defaultLookbackDelta = 5 * time.Minute
	defaultMaxSamples    = 50000
	defaultQueryTimeout  = 10 * time.Second
	defaultMaxPoints     = 11000
	defaultSubqueryStep  = time.Minute
	maxAnnotationCount   = 10
	// Keep this aligned with the Prometheus module version in go.mod. The
	// v0.314.0 Go module contains the Prometheus 3.14.0 query engine.
	prometheusCompatibilityVersion = "3.14.0"
)

type api struct {
	queryable storage.Queryable
	engine    *promengine.Engine
	parser    parser.Parser
	now       func() time.Time
	maxPoints int
}

func newAPI(cfg *Config) *api {
	lookbackDelta := cfg.LookbackDelta
	if lookbackDelta <= 0 {
		lookbackDelta = defaultLookbackDelta
	}
	maxSamples := cfg.MaxSamples
	if maxSamples <= 0 {
		maxSamples = defaultMaxSamples
	}
	queryTimeout := cfg.QueryTimeout
	if queryTimeout <= 0 {
		queryTimeout = defaultQueryTimeout
	}
	maxPoints := cfg.MaxPoints
	if maxPoints <= 0 {
		maxPoints = defaultMaxPoints
	}
	now := cfg.Now
	if now == nil {
		now = time.Now
	}

	return &api{
		queryable: cfg.Queryable,
		engine: promengine.NewEngine(promengine.EngineOpts{
			Logger:               promslog.NewNopLogger(),
			MaxSamples:           maxSamples,
			Timeout:              queryTimeout,
			LookbackDelta:        lookbackDelta,
			EnableAtModifier:     true,
			EnableNegativeOffset: true,
			NoStepSubqueryIntervalFn: func(int64) int64 {
				return defaultSubqueryStep.Milliseconds()
			},
		}),
		parser:    parser.NewParser(parser.Options{}),
		now:       now,
		maxPoints: maxPoints,
	}
}

func (a *api) handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /api/v1/query", a.query)
	mux.HandleFunc("POST /api/v1/query", a.query)
	mux.HandleFunc("GET /api/v1/query_range", a.queryRange)
	mux.HandleFunc("POST /api/v1/query_range", a.queryRange)
	mux.HandleFunc("GET /api/v1/series", a.series)
	mux.HandleFunc("POST /api/v1/series", a.series)
	mux.HandleFunc("GET /api/v1/labels", a.labelNames)
	mux.HandleFunc("POST /api/v1/labels", a.labelNames)
	mux.HandleFunc("GET /api/v1/label/{name}/values", a.labelValues)
	mux.HandleFunc("GET /api/v1/metadata", a.metadata)
	mux.HandleFunc("GET /api/v1/status/buildinfo", a.buildInfo)
	mux.HandleFunc("GET /-/healthy", healthy)
	mux.HandleFunc("GET /-/ready", ready)
	return mux
}

type response struct {
	Status    string   `json:"status"`
	Data      any      `json:"data,omitempty"`
	ErrorType string   `json:"errorType,omitempty"`
	Error     string   `json:"error,omitempty"`
	Warnings  []string `json:"warnings,omitempty"`
	Infos     []string `json:"infos,omitempty"`
}

type queryData struct {
	ResultType parser.ValueType `json:"resultType"`
	Result     parser.Value     `json:"result"`
}

type apiError struct {
	typ    string
	err    error
	status int
}

func badData(parameter string, err error) *apiError {
	return &apiError{
		typ:    "bad_data",
		err:    fmt.Errorf("invalid parameter %q: %w", parameter, err),
		status: http.StatusBadRequest,
	}
}

func executionError(err error) *apiError {
	var canceled promengine.ErrQueryCanceled
	var timeout promengine.ErrQueryTimeout
	var storageErr promengine.ErrStorage
	switch {
	case errors.As(err, &canceled), errors.Is(err, context.Canceled):
		return &apiError{typ: "canceled", err: err, status: 499}
	case errors.As(err, &timeout), errors.Is(err, context.DeadlineExceeded):
		return &apiError{typ: "timeout", err: err, status: http.StatusServiceUnavailable}
	case errors.As(err, &storageErr):
		return &apiError{typ: "internal", err: err, status: http.StatusInternalServerError}
	default:
		return &apiError{typ: "execution", err: err, status: http.StatusUnprocessableEntity}
	}
}

func (a *api) query(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		writeError(w, badData("form", err))
		return
	}

	evaluationTime := a.now()
	if value := r.FormValue("time"); value != "" {
		parsed, err := parseTime(value)
		if err != nil {
			writeError(w, badData("time", err))
			return
		}
		evaluationTime = parsed
	}

	ctx, cancel, err := requestContext(r)
	if err != nil {
		writeError(w, badData("timeout", err))
		return
	}
	defer cancel()

	queryOpts, err := queryOptions(r)
	if err != nil {
		writeError(w, badData("lookback_delta", err))
		return
	}
	query, err := a.engine.NewInstantQuery(ctx, a.queryable, queryOpts, r.FormValue("query"), evaluationTime)
	if err != nil {
		writeError(w, badData("query", err))
		return
	}
	defer query.Close()

	result := query.Exec(ctx)
	if result.Err != nil {
		writeError(w, executionError(result.Err))
		return
	}
	value, warnings, err := applyResultLimit(result.Value, result.Warnings, r.FormValue("limit"))
	if err != nil {
		writeError(w, badData("limit", err))
		return
	}
	writeQueryResult(w, r.FormValue("query"), value, warnings)
}

func (a *api) queryRange(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		writeError(w, badData("form", err))
		return
	}
	start, err := parseTime(r.FormValue("start"))
	if err != nil {
		writeError(w, badData("start", err))
		return
	}
	end, err := parseTime(r.FormValue("end"))
	if err != nil {
		writeError(w, badData("end", err))
		return
	}
	if end.Before(start) {
		writeError(w, badData("end", errors.New("end timestamp must not be before start time")))
		return
	}
	step, err := parseDuration(r.FormValue("step"))
	if err != nil {
		writeError(w, badData("step", err))
		return
	}
	if step <= 0 {
		writeError(w, badData("step", errors.New("zero or negative query resolution step widths are not accepted")))
		return
	}
	if end.Sub(start)/step > time.Duration(a.maxPoints) {
		writeError(w, badData("step", fmt.Errorf("exceeded maximum resolution of %d points per timeseries", a.maxPoints)))
		return
	}

	ctx, cancel, err := requestContext(r)
	if err != nil {
		writeError(w, badData("timeout", err))
		return
	}
	defer cancel()
	queryOpts, err := queryOptions(r)
	if err != nil {
		writeError(w, badData("lookback_delta", err))
		return
	}
	query, err := a.engine.NewRangeQuery(ctx, a.queryable, queryOpts, r.FormValue("query"), start, end, step)
	if err != nil {
		writeError(w, badData("query", err))
		return
	}
	defer query.Close()

	result := query.Exec(ctx)
	if result.Err != nil {
		writeError(w, executionError(result.Err))
		return
	}
	value, warnings, err := applyResultLimit(result.Value, result.Warnings, r.FormValue("limit"))
	if err != nil {
		writeError(w, badData("limit", err))
		return
	}
	writeQueryResult(w, r.FormValue("query"), value, warnings)
}

func requestContext(r *http.Request) (context.Context, context.CancelFunc, error) {
	if value := r.FormValue("timeout"); value != "" {
		timeout, err := parseDuration(value)
		if err != nil {
			return nil, nil, err
		}
		ctx, cancel := context.WithTimeout(r.Context(), timeout)
		return ctx, cancel, nil
	}
	return r.Context(), func() {}, nil
}

func queryOptions(r *http.Request) (promengine.QueryOpts, error) {
	var lookbackDelta time.Duration
	if value := r.FormValue("lookback_delta"); value != "" {
		parsed, err := parseDuration(value)
		if err != nil {
			return nil, err
		}
		lookbackDelta = parsed
	}
	return promengine.NewPrometheusQueryOpts(r.FormValue("stats") == "all", lookbackDelta), nil
}

func applyResultLimit(value parser.Value, warnings annotations.Annotations, input string) (parser.Value, annotations.Annotations, error) {
	if input == "" {
		return value, warnings, nil
	}
	limit, err := strconv.Atoi(input)
	if err != nil || limit < 0 {
		return nil, warnings, errors.New("limit must be a non-negative integer")
	}
	if limit == 0 {
		return value, warnings, nil
	}

	truncated := false
	switch typed := value.(type) {
	case promengine.Vector:
		if len(typed) > limit {
			value = typed[:limit]
			truncated = true
		}
	case promengine.Matrix:
		if len(typed) > limit {
			value = typed[:limit]
			truncated = true
		}
	}
	if truncated {
		(&warnings).Add(errors.New("results truncated due to limit"))
	}
	return value, warnings, nil
}

func writeQueryResult(w http.ResponseWriter, expression string, value parser.Value, queryWarnings annotations.Annotations) {
	warnings, infos := queryWarnings.AsStrings(expression, maxAnnotationCount, maxAnnotationCount)
	writeJSON(w, http.StatusOK, &response{
		Status: "success",
		Data: &queryData{
			ResultType: value.Type(),
			Result:     value,
		},
		Warnings: warnings,
		Infos:    infos,
	})
}

func (a *api) labelNames(w http.ResponseWriter, r *http.Request) {
	querier, matcherSets, limit, apiErr := a.labelQuery(r)
	if apiErr != nil {
		writeError(w, apiErr)
		return
	}
	defer querier.Close()

	names := map[string]struct{}{}
	var queryWarnings annotations.Annotations
	if len(matcherSets) == 0 {
		matcherSets = []matcherSet{nil}
	}
	for _, matchers := range matcherSets {
		values, warnings, err := querier.LabelNames(r.Context(), &storage.LabelHints{}, matchers...)
		if err != nil {
			writeError(w, executionError(err))
			return
		}
		queryWarnings.Merge(warnings)
		for _, value := range values {
			names[value] = struct{}{}
		}
	}
	result := sortedKeys(names, limit)
	writeData(w, result, queryWarnings, "")
}

func (a *api) labelValues(w http.ResponseWriter, r *http.Request) {
	name := r.PathValue("name")
	if strings.HasPrefix(name, "U__") {
		name = commonmodel.UnescapeName(name, commonmodel.ValueEncodingEscaping)
	}
	if !commonmodel.UTF8Validation.IsValidLabelName(name) {
		writeError(w, badData("name", fmt.Errorf("invalid label name %q", name)))
		return
	}
	querier, matcherSets, limit, apiErr := a.labelQuery(r)
	if apiErr != nil {
		writeError(w, apiErr)
		return
	}
	defer querier.Close()

	valuesSet := map[string]struct{}{}
	var queryWarnings annotations.Annotations
	if len(matcherSets) == 0 {
		matcherSets = []matcherSet{nil}
	}
	for _, matchers := range matcherSets {
		values, warnings, err := querier.LabelValues(r.Context(), name, &storage.LabelHints{}, matchers...)
		if err != nil {
			writeError(w, executionError(err))
			return
		}
		queryWarnings.Merge(warnings)
		for _, value := range values {
			valuesSet[value] = struct{}{}
		}
	}
	writeData(w, sortedKeys(valuesSet, limit), queryWarnings, "")
}

type matcherSet []*labels.Matcher

func (a *api) labelQuery(r *http.Request) (storage.Querier, []matcherSet, int, *apiError) {
	if err := r.ParseForm(); err != nil {
		return nil, nil, 0, badData("form", err)
	}
	mint, maxt, err := queryBounds(r)
	if err != nil {
		return nil, nil, 0, err
	}
	matcherSets, parseErr := a.parseMatcherSets(r.Form["match[]"])
	if parseErr != nil {
		return nil, nil, 0, badData("match[]", parseErr)
	}
	limit, parseErr := parseLimit(r.FormValue("limit"))
	if parseErr != nil {
		return nil, nil, 0, badData("limit", parseErr)
	}
	querier, openErr := a.queryable.Querier(mint, maxt)
	if openErr != nil {
		return nil, nil, 0, executionError(openErr)
	}
	return querier, matcherSets, limit, nil
}

func (a *api) series(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		writeError(w, badData("form", err))
		return
	}
	if len(r.Form["match[]"]) == 0 {
		writeError(w, badData("match[]", errors.New("no match[] parameter provided")))
		return
	}
	mint, maxt, apiErr := queryBounds(r)
	if apiErr != nil {
		writeError(w, apiErr)
		return
	}
	matcherSets, err := a.parseMatcherSets(r.Form["match[]"])
	if err != nil {
		writeError(w, badData("match[]", err))
		return
	}
	limit, err := parseLimit(r.FormValue("limit"))
	if err != nil {
		writeError(w, badData("limit", err))
		return
	}
	querier, err := a.queryable.Querier(mint, maxt)
	if err != nil {
		writeError(w, executionError(err))
		return
	}
	defer querier.Close()

	unique := map[string]labels.Labels{}
	var queryWarnings annotations.Annotations
	for _, matchers := range matcherSets {
		set := querier.Select(r.Context(), true, &storage.SelectHints{
			Start: mint,
			End:   maxt,
			Func:  "series",
		}, matchers...)
		queryWarnings.Merge(set.Warnings())
		for set.Next() {
			metric := set.At().Labels()
			unique[metric.String()] = metric
		}
		if err := set.Err(); err != nil {
			writeError(w, executionError(err))
			return
		}
	}

	metrics := make([]labels.Labels, 0, len(unique))
	for _, metric := range unique {
		metrics = append(metrics, metric)
	}
	slices.SortFunc(metrics, labels.Compare)
	if limit > 0 && len(metrics) > limit {
		metrics = metrics[:limit]
		queryWarnings.Add(errors.New("results truncated due to limit"))
	}
	writeData(w, metrics, queryWarnings, "")
}

func (a *api) parseMatcherSets(inputs []string) ([]matcherSet, error) {
	if len(inputs) == 0 {
		return nil, nil
	}
	parsed, err := a.parser.ParseMetricSelectors(inputs)
	if err != nil {
		return nil, err
	}
	sets := make([]matcherSet, len(parsed))
	for i, matchers := range parsed {
		valid := false
		for _, matcher := range matchers {
			if matcher != nil && !matcher.Matches("") {
				valid = true
				break
			}
		}
		if !valid {
			return nil, errors.New("match[] must contain at least one non-empty matcher")
		}
		sets[i] = matchers
	}
	return sets, nil
}

func queryBounds(r *http.Request) (int64, int64, *apiError) {
	mint := int64(math.MinInt64)
	maxt := int64(math.MaxInt64)
	if value := r.FormValue("start"); value != "" {
		parsed, err := parseTime(value)
		if err != nil {
			return 0, 0, badData("start", err)
		}
		mint = timestamp.FromTime(parsed)
	}
	if value := r.FormValue("end"); value != "" {
		parsed, err := parseTime(value)
		if err != nil {
			return 0, 0, badData("end", err)
		}
		maxt = timestamp.FromTime(parsed)
	}
	if maxt < mint {
		return 0, 0, badData("end", errors.New("end timestamp must not be before start time"))
	}
	return mint, maxt, nil
}

func parseLimit(value string) (int, error) {
	if value == "" {
		return 0, nil
	}
	limit, err := strconv.Atoi(value)
	if err != nil || limit < 0 {
		return 0, errors.New("limit must be a non-negative integer")
	}
	return limit, nil
}

func sortedKeys(values map[string]struct{}, limit int) []string {
	result := make([]string, 0, len(values))
	for value := range values {
		result = append(result, value)
	}
	slices.Sort(result)
	if limit > 0 && len(result) > limit {
		return result[:limit]
	}
	return result
}

func (a *api) metadata(w http.ResponseWriter, _ *http.Request) {
	// Scraped samples do not currently carry Prometheus metadata. An empty map
	// is the Prometheus-compatible response and lets clients discover that fact.
	writeJSON(w, http.StatusOK, &response{Status: "success", Data: map[string]any{}})
}

func (a *api) buildInfo(w http.ResponseWriter, _ *http.Request) {
	writeJSON(w, http.StatusOK, &response{Status: "success", Data: map[string]string{
		"version":   prometheusCompatibilityVersion,
		"revision":  "unknown",
		"branch":    "unknown",
		"buildUser": "unknown",
		"buildDate": "unknown",
		"goVersion": runtime.Version(),
	}})
}

func healthy(w http.ResponseWriter, _ *http.Request) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.WriteHeader(http.StatusOK)
	_, _ = w.Write([]byte("Capacity Planner Query Server is Healthy.\n"))
}

func ready(w http.ResponseWriter, _ *http.Request) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.WriteHeader(http.StatusOK)
	_, _ = w.Write([]byte("Capacity Planner Query Server is Ready.\n"))
}

func writeData(w http.ResponseWriter, data any, queryWarnings annotations.Annotations, expression string) {
	warnings, infos := queryWarnings.AsStrings(expression, maxAnnotationCount, maxAnnotationCount)
	writeJSON(w, http.StatusOK, &response{
		Status:   "success",
		Data:     data,
		Warnings: warnings,
		Infos:    infos,
	})
}

func writeError(w http.ResponseWriter, apiErr *apiError) {
	writeJSON(w, apiErr.status, &response{
		Status:    "error",
		ErrorType: apiErr.typ,
		Error:     apiErr.err.Error(),
	})
}

func writeJSON(w http.ResponseWriter, status int, payload *response) {
	encoded, err := json.Marshal(payload)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_, _ = w.Write(encoded)
}

func parseTime(value string) (time.Time, error) {
	if parsed, err := strconv.ParseFloat(value, 64); err == nil {
		seconds, fraction := math.Modf(parsed)
		fraction = math.Round(fraction*1000) / 1000
		return time.Unix(int64(seconds), int64(fraction*float64(time.Second))).UTC(), nil
	}
	if parsed, err := time.Parse(time.RFC3339Nano, value); err == nil {
		return parsed, nil
	}
	return time.Time{}, fmt.Errorf("cannot parse %q to a valid timestamp", value)
}

func parseDuration(value string) (time.Duration, error) {
	if parsed, err := strconv.ParseFloat(value, 64); err == nil {
		nanoseconds := parsed * float64(time.Second)
		if nanoseconds > float64(math.MaxInt64) || nanoseconds < float64(math.MinInt64) {
			return 0, fmt.Errorf("cannot parse %q to a valid duration: value overflows int64", value)
		}
		return time.Duration(nanoseconds), nil
	}
	if parsed, err := commonmodel.ParseDuration(value); err == nil {
		return time.Duration(parsed), nil
	}
	return 0, fmt.Errorf("cannot parse %q to a valid duration", value)
}