// Package query implements the decision metrics Provider with Prometheus's
// PromQL engine and an embedded TSDB.
package query

import (
	"context"
	"errors"
	"fmt"
	"maps"
	"sort"
	"time"

	commonmodel "github.com/prometheus/common/model"
	"github.com/prometheus/common/promslog"
	"github.com/prometheus/prometheus/model/labels"
	promengine "github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/promql/parser"
	"github.com/prometheus/prometheus/storage"

	decisionmetrics "github.com/tidbcloud/cluster-service-ng/internal/capacity-planner/pkg/metrics"
	metricstarget "github.com/tidbcloud/cluster-service-ng/internal/capacity-planner/pkg/target"
)

const (
	defaultLookbackDelta = 5 * time.Minute
	defaultMaxSamples    = 50000
	defaultQueryTimeout  = 10 * time.Second
)

var (
	ErrInvalidQuery     = errors.New("invalid metrics query")
	ErrUnsupportedQuery = errors.New("unsupported metrics query")
)

type Clock interface {
	Now() time.Time
}

// Config configures a PromQL-backed Provider.
type Config struct {
	Queryable     storage.Queryable
	Clock         Clock
	LookbackDelta time.Duration
	MaxSamples    int
	QueryTimeout  time.Duration
}

// Provider evaluates the supported selector-only PromQL subset.
type Provider struct {
	queryable storage.Queryable
	clock     Clock
	engine    *promengine.Engine
	parser    parser.Parser
	queryOpts promengine.QueryOpts
}

var _ decisionmetrics.Provider = (*Provider)(nil)

func New(cfg *Config) (*Provider, error) {
	if cfg == nil {
		return nil, fmt.Errorf("promql provider: config is required")
	}
	if cfg.Queryable == nil {
		return nil, fmt.Errorf("promql provider: queryable is required")
	}
	if cfg.Clock == nil {
		return nil, fmt.Errorf("promql provider: clock is required")
	}
	lookbackDelta := cfg.LookbackDelta
	if lookbackDelta <= 0 {
		lookbackDelta = defaultLookbackDelta
	}
	maxSamples := cfg.MaxSamples
	if maxSamples <= 0 {
		maxSamples = defaultMaxSamples
	}
	queryTimeout := cfg.QueryTimeout
	if queryTimeout <= 0 {
		queryTimeout = defaultQueryTimeout
	}

	engine := promengine.NewEngine(promengine.EngineOpts{
		Logger:               promslog.NewNopLogger(),
		MaxSamples:           maxSamples,
		Timeout:              queryTimeout,
		LookbackDelta:        lookbackDelta,
		EnableAtModifier:     false,
		EnableNegativeOffset: false,
		NoStepSubqueryIntervalFn: func(int64) int64 {
			return time.Minute.Milliseconds()
		},
	})
	return &Provider{
		queryable: cfg.Queryable,
		clock:     cfg.Clock,
		engine:    engine,
		parser:    parser.NewParser(parser.Options{}),
		queryOpts: promengine.NewPrometheusQueryOpts(false, lookbackDelta),
	}, nil
}

// Query implements metrics.Provider. Every expression is parsed and validated
// before any query runs, and any execution failure discards the whole batch.
func (p *Provider) Query(ctx context.Context, target metricstarget.Target, promQLs ...string) ([]decisionmetrics.Metrics, error) {
	if len(promQLs) == 0 {
		return []decisionmetrics.Metrics{}, nil
	}
	queries := make([]preparedQuery, len(promQLs))
	for i := range promQLs {
		prepared, err := p.prepare(promQLs[i])
		if err != nil {
			return nil, fmt.Errorf("promql provider: queries[%d]: %w", i, err)
		}
		queries[i] = prepared
	}

	evaluationTime := p.clock.Now()
	results := make([]decisionmetrics.Metrics, len(queries))
	targetFilter := targetSeriesFilter{target: target}
	for i := range queries {
		result, err := p.execute(ctx, &queries[i], evaluationTime, &targetFilter)
		if err != nil {
			return nil, fmt.Errorf("promql provider: queries[%d]: %w", i, err)
		}
		results[i] = *result
	}
	return results, nil
}

type preparedQuery struct {
	promQL string
}

func (p *Provider) prepare(input string) (preparedQuery, error) {
	expr, err := p.parser.ParseExpr(input)
	if err != nil {
		return preparedQuery{}, fmt.Errorf("%w: %v", ErrInvalidQuery, err)
	}
	if _, err := selectorForExpression(expr); err != nil {
		return preparedQuery{}, err
	}
	return preparedQuery{promQL: input}, nil
}

func selectorForExpression(expr parser.Expr) (*parser.VectorSelector, error) {
	var selector *parser.VectorSelector
	switch typed := expr.(type) {
	case *parser.VectorSelector:
		selector = typed
	case *parser.MatrixSelector:
		if typed.Range <= 0 {
			return nil, fmt.Errorf("%w: range must be greater than zero", ErrUnsupportedQuery)
		}
		var ok bool
		selector, ok = typed.VectorSelector.(*parser.VectorSelector)
		if !ok {
			return nil, fmt.Errorf("%w: range must apply directly to a vector selector", ErrUnsupportedQuery)
		}
	default:
		return nil, fmt.Errorf("%w: only vector and range selectors are allowed", ErrUnsupportedQuery)
	}

	if selector.Name == "" {
		return nil, fmt.Errorf("%w: an explicit metric name is required", ErrUnsupportedQuery)
	}
	if selector.OriginalOffset != 0 || selector.OriginalOffsetExpr != nil {
		return nil, fmt.Errorf("%w: offset is not allowed", ErrUnsupportedQuery)
	}
	if selector.Timestamp != nil || selector.StartOrEnd != 0 {
		return nil, fmt.Errorf("%w: @ modifier is not allowed", ErrUnsupportedQuery)
	}
	if selector.Anchored || selector.Smoothed {
		return nil, fmt.Errorf("%w: extended range selectors are not allowed", ErrUnsupportedQuery)
	}
	for _, matcher := range selector.LabelMatchers {
		if matcher.Name == labels.MetricName && (matcher.Type != labels.MatchEqual || matcher.Value != selector.Name) {
			return nil, fmt.Errorf("%w: metric-name matchers are not allowed", ErrUnsupportedQuery)
		}
	}
	return selector, nil
}

func validateTargetLabels(targetLabels map[string]string) error {
	for name := range targetLabels {
		if !commonmodel.LegacyValidation.IsValidLabelName(name) || name == labels.MetricName {
			return fmt.Errorf("promql provider: target has invalid label name %q", name)
		}
	}
	return nil
}

func (p *Provider) execute(ctx context.Context, prepared *preparedQuery, evaluationTime time.Time, targetFilter *targetSeriesFilter) (*decisionmetrics.Metrics, error) {
	query, err := p.engine.NewInstantQuery(ctx, p.queryable, p.queryOpts, prepared.promQL, evaluationTime)
	if err != nil {
		return nil, fmt.Errorf("build query: %w", err)
	}
	defer query.Close()

	result := query.Exec(ctx)
	if result.Err != nil {
		return nil, fmt.Errorf("execute query: %w", result.Err)
	}
	return convertResult(result.Value, targetFilter)
}

type targetSeriesFilter struct {
	target      metricstarget.Target
	initialized bool
	labels      map[string]string
	err         error
}

func (f *targetSeriesFilter) keep(metricLabels labels.Labels) (bool, error) {
	if metricLabels.Get(decisionmetrics.LabelPod) == "" {
		return true, nil
	}
	if !f.initialized {
		f.initialized = true
		if f.target == nil {
			f.err = fmt.Errorf("promql provider: target is required for Pod metrics")
			return false, f.err
		}
		namespace := f.target.Namespace()
		if namespace == "" {
			f.err = fmt.Errorf("promql provider: target namespace is required for Pod metrics")
			return false, f.err
		}
		f.labels = maps.Clone(f.target.Labels())
		if len(f.labels) == 0 {
			f.err = fmt.Errorf("promql provider: target labels are required for Pod metrics")
			return false, f.err
		}
		f.err = validateTargetLabels(f.labels)
		f.labels[decisionmetrics.LabelNamespace] = namespace
	}
	if f.err != nil {
		return false, f.err
	}
	for name, value := range f.labels {
		if metricLabels.Get(name) != value {
			return false, nil
		}
	}
	return true, nil
}

type convertedSeries struct {
	key    string
	series decisionmetrics.TimeSeries
}

func convertResult(value parser.Value, targetFilter *targetSeriesFilter) (*decisionmetrics.Metrics, error) {
	var converted []convertedSeries
	switch typed := value.(type) {
	case promengine.Vector:
		converted = make([]convertedSeries, 0, len(typed))
		for i := range typed {
			if typed[i].H != nil {
				return nil, fmt.Errorf("native histogram samples are not supported")
			}
			keep, err := targetFilter.keep(typed[i].Metric)
			if err != nil {
				return nil, err
			}
			if !keep {
				continue
			}
			series := newTimeSeries(typed[i].Metric, []decisionmetrics.Sample{{
				Timestamp: time.UnixMilli(typed[i].T),
				Value:     typed[i].F,
			}})
			converted = append(converted, convertedSeries{key: typed[i].Metric.String(), series: series})
		}
	case promengine.Matrix:
		converted = make([]convertedSeries, 0, len(typed))
		for i := range typed {
			if len(typed[i].Histograms) != 0 {
				return nil, fmt.Errorf("native histogram samples are not supported")
			}
			samples := make([]decisionmetrics.Sample, len(typed[i].Floats))
			for j := range typed[i].Floats {
				samples[j] = decisionmetrics.Sample{
					Timestamp: time.UnixMilli(typed[i].Floats[j].T),
					Value:     typed[i].Floats[j].F,
				}
			}
			keep, err := targetFilter.keep(typed[i].Metric)
			if err != nil {
				return nil, err
			}
			if !keep {
				continue
			}
			series := newTimeSeries(typed[i].Metric, samples)
			converted = append(converted, convertedSeries{key: typed[i].Metric.String(), series: series})
		}
	default:
		return nil, fmt.Errorf("unexpected PromQL result type %s", value.Type())
	}

	sort.Slice(converted, func(i, j int) bool { return converted[i].key < converted[j].key })
	series := make([]decisionmetrics.TimeSeries, len(converted))
	for i := range converted {
		series[i] = converted[i].series
	}
	return &decisionmetrics.Metrics{Series: series}, nil
}

func newTimeSeries(metricLabels labels.Labels, samples []decisionmetrics.Sample) decisionmetrics.TimeSeries {
	series := decisionmetrics.TimeSeries{Samples: samples}
	if name := metricLabels.Get(decisionmetrics.LabelPod); name != "" {
		series.Ref = &decisionmetrics.Reference{Name: name}
	}
	return series
}