// Licensed to Apache Software Foundation (ASF) under one or more contributor
// license agreements. See the NOTICE file distributed with
// this work for additional information regarding copyright
// ownership. Apache Software Foundation (ASF) licenses this file to you under
// the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

// Benchmarks for the shared sampler engine behind both first-party plugins. Decide
// runs inside the merge/finalize critical path — a filtering merge pays this per
// trace on top of the block decode it already forced — so the numbers that matter
// are ns/trace and allocs/trace, not wall time for the whole batch.
//
// Every case is parameterized by the two plugins' real Schema values
// (segmentSchema / zipkinSchema from sampler_test.go), so the sub-benchmark names
// read "sw/..." and "zipkin/..." and a per-plugin regression is attributable. The
// plugin mains add nothing measurable: each is a one-line delegate to New.
//
// Batches are built ONCE and reused across iterations, which is only sound because
// Decide treats the batch as read-only. That is load-bearing rather than incidental:
// the SDK's string-array decode rewrites its source in place, so an engine that
// decoded the shared buffer would return different verdicts on the second iteration
// (see TestDecide_EscapedEntriesSurviveChainedLinks). A benchmark whose ns/op drifts
// with -benchtime is a symptom of that bug returning.
package tracesampler

import (
	"fmt"
	"sync/atomic"
	"testing"

	"github.com/apache/skywalking-banyandb/pkg/pb/v1/valuetype"
	"github.com/apache/skywalking-banyandb/pkg/pipeline/sdk"
	"github.com/apache/skywalking-banyandb/pkg/pipeline/sdk/sdktest"
)

// benchSink keeps the compiler from eliminating the Decide call.
var benchSink sdk.Verdict

// benchPlugin describes one shipped plugin: its Schema plus the column names its
// fixtures have to populate.
type benchPlugin struct {
	name     string
	arrayCol string
	startCol string
	durCol   string
	errCol   string // "" when the error signal lives inside arrayCol
	scenario string // the Scenario 6.1 / 6.2 gating config as shipped
	// Field order below satisfies govet's fieldalignment: strings first, then the
	// pointer-bearing Schema, then scalars.
	schema     Schema
	durationMs int64
}

func benchPlugins() []benchPlugin {
	return []benchPlugin{
		{
			name: "sw", schema: segmentSchema,
			arrayCol: "tags", startCol: "start_time", durCol: "latency", errCol: "is_error",
			durationMs: 500,
			scenario: `{
				"durationThresholdMs": 500,
				"keepErrors": true,
				"healthySampleRate": 0.1,
				"keepTagRules": [
					{"tagKey":"db.type","equals":"PostgreSQL"},
					{"tagKey":"mq.queue","equals":"queue-songs-ping"}
				]
			}`,
		},
		{
			name: "zipkin", schema: zipkinSchema,
			arrayCol: "query", startCol: "timestamp_millis", durCol: "duration", errCol: "",
			durationMs: 1000,
			scenario: `{
				"durationThresholdMs": 1000,
				"keepErrors": true,
				"healthySampleRate": 0.05,
				"keepTagRules": [
					{"tagKey":"query","regex":"http\\.status_code=5\\d\\d"}
				]
			}`,
		},
	}
}

// benchEntries returns n realistic flattened searchable-tag entries. OAP writes each
// tag as "key=value", and Zipkin additionally writes the bare key, so the array a
// real trace carries is several entries deep even for a plain HTTP call.
func benchEntries(n int, escaped bool) []string {
	base := []string{
		"http.method=GET",
		"http.status_code=200",
		"url=http://frontend/api/orders",
		"rpc.system=grpc",
		"db.type=MySQL",
		"mq.queue=orders",
		"otel.scope=io.opentelemetry.tomcat",
		"thread.name=http-nio-8080-exec-3",
	}
	out := make([]string, 0, n)
	for i := 0; i < n; i++ {
		e := base[i%len(base)]
		if escaped {
			// A "|" forces the decoder's escape path, which is what makes the
			// defensive per-row copy in arrayEntries cost anything.
			e += "|shard" + fmt.Sprint(i)
		}
		out = append(out, fmt.Sprintf("%s-%d", e, i/len(base)))
	}
	return out
}

// benchBatch builds traces that match NOTHING, the honest upper bound: every
// sure-keep predicate is evaluated before the verdict is known. rows is the segment
// (or span) count per trace; entries is the depth of the flattened tag array.
func benchBatch(b *testing.B, p benchPlugin, traces, rows, entries int, escaped bool) *sdk.TraceBatch {
	b.Helper()
	blocks := make([]sdk.TraceBlock, 0, traces)
	for t := 0; t < traces; t++ {
		tb := sdktest.NewTrace(fmt.Sprintf("bench-trace-%d", t))
		for r := 0; r < rows; r++ {
			// Fast rows: the envelope stays far below the threshold, so the duration
			// rule is evaluated in full and still misses.
			tb.TagAs(p.startCol, valuetype.ValueTypeTimestamp, int64(r)*int64(1_000_000))
			tb.Tag(p.durCol, int64(3))
			if p.errCol != "" {
				tb.Tag(p.errCol, int64(0))
			}
			tb.Tag(p.arrayCol, benchEntries(entries, escaped))
		}
		blk, err := tb.Build()
		if err != nil {
			b.Fatalf("fixture: %v", err)
		}
		blocks = append(blocks, blk)
	}
	return sdktest.Batch(blocks...)
}

func mustSampler(b *testing.B, cfg string, schema Schema) sdk.Sampler {
	b.Helper()
	s, err := New([]byte(cfg), schema)
	if err != nil {
		b.Fatalf("config %s: %v", cfg, err)
	}
	return s
}

// runDecide reports per-trace cost, which is the unit that scales with ingest.
func runDecide(b *testing.B, s sdk.Sampler, batch *sdk.TraceBatch) {
	b.Helper()
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		v, err := s.Decide(batch)
		if err != nil {
			b.Fatalf("decide: %v", err)
		}
		benchSink = v
	}
	b.StopTimer()
	elapsed := b.Elapsed()
	traceCount := float64(b.N * len(batch.Traces))
	b.ReportMetric(float64(elapsed.Nanoseconds())/traceCount, "ns/trace")
	b.ReportMetric(traceCount/elapsed.Seconds(), "traces/s")
}

// BenchmarkDecide isolates what each keep rule costs, so a config change's price is
// visible: the duration envelope reads two numeric columns, the tag rules force the
// flattened-array decode, and the healthy sample is a single hash.
func BenchmarkDecide(b *testing.B) {
	const (
		traces  = 64
		rows    = 3
		entries = 8
	)
	for _, p := range benchPlugins() {
		configs := []struct{ name, cfg string }{
			{"sampleOnly", `{"healthySampleRate":0.1}`},
			{"durationOnly", fmt.Sprintf(`{"durationThresholdMs":%d}`, p.durationMs)},
			{"errorsOnly", `{"keepErrors":true}`},
			{"tagRulesOnly", `{"keepTagRules":[{"tagKey":"db.type","equals":"PostgreSQL"}]}`},
			{"scenario", p.scenario},
		}
		for _, c := range configs {
			b.Run(p.name+"/"+c.name, func(b *testing.B) {
				s := mustSampler(b, c.cfg, p.schema)
				runDecide(b, s, benchBatch(b, p, traces, rows, entries, false))
			})
		}
	}
}

// BenchmarkDecide_EarlyExit quantifies the OR short-circuit: rules are evaluated in
// order and the first match wins, so a trace kept by the first rule never pays for
// the rest. Worth knowing when ordering a config for a mostly-kept workload.
func BenchmarkDecide_EarlyExit(b *testing.B) {
	for _, p := range benchPlugins() {
		// durationThresholdMs is evaluated first, so a threshold of 1ms is matched by
		// every trace immediately; the scenario threshold is never reached.
		hit := `{"durationThresholdMs":1,"keepErrors":true,` +
			`"keepTagRules":[{"tagKey":"db.type","equals":"PostgreSQL"}],"healthySampleRate":0.1}`
		miss := p.scenario
		for _, c := range []struct{ name, cfg string }{{"firstRuleHits", hit}, {"allRulesMiss", miss}} {
			b.Run(p.name+"/"+c.name, func(b *testing.B) {
				s := mustSampler(b, c.cfg, p.schema)
				runDecide(b, s, benchBatch(b, p, 64, 3, 8, false))
			})
		}
	}
}

// BenchmarkDecide_ArrayEntries shows how tag-rule cost scales with the depth of the
// flattened array, and what the escape path costs. The escaped variant is the one to
// watch: arrayEntries decodes each row from a copy because the SDK's decode mutates
// its source, and only values containing "|" or "\" actually shift bytes.
func BenchmarkDecide_ArrayEntries(b *testing.B) {
	for _, p := range benchPlugins() {
		cfg := `{"keepTagRules":[{"tagKey":"db.type","equals":"PostgreSQL"}]}`
		for _, entries := range []int{2, 8, 32} {
			for _, esc := range []bool{false, true} {
				name := fmt.Sprintf("%s/entries=%d/escaped=%v", p.name, entries, esc)
				b.Run(name, func(b *testing.B) {
					s := mustSampler(b, cfg, p.schema)
					runDecide(b, s, benchBatch(b, p, 64, 1, entries, esc))
				})
			}
		}
	}
}

// BenchmarkDecide_BatchSize checks the per-trace cost stays flat as the batch grows.
// The engine loops over batch.Traces with no cross-trace state, so a rising ns/trace
// would mean an accidental per-batch allocation or quadratic step crept in.
//
// 1 and 16 bracket the small end, where a per-batch cost would show up amplified;
// 64/128/256 span the range a real merge stages, since the filter's chunk budget is
// derived from the memory protector rather than a trace count. Measured across those
// three, ns/trace holds within 3% and allocs/trace is constant — only the keep mask
// is per-batch, so its share halves as the batch doubles.
func BenchmarkDecide_BatchSize(b *testing.B) {
	for _, p := range benchPlugins() {
		for _, traces := range []int{1, 16, 64, 128, 256} {
			b.Run(fmt.Sprintf("%s/traces=%d", p.name, traces), func(b *testing.B) {
				s := mustSampler(b, p.scenario, p.schema)
				runDecide(b, s, benchBatch(b, p, traces, 3, 8, false))
			})
		}
	}
}

type benchmarkCounter struct{ total atomic.Uint64 }

func (c *benchmarkCounter) Inc(delta float64, _ ...string) { c.total.Add(uint64(delta)) }

type benchmarkMeter struct {
	decisions   *benchmarkCounter
	rows        *benchmarkCounter
	droppedRows *benchmarkCounter
	unavailable *benchmarkCounter
}

func (m *benchmarkMeter) Counter(name string, _ ...string) sdk.Counter {
	switch name {
	case decisionMetricName:
		return m.decisions
	case rowMetricName:
		return m.rows
	case rowDroppedMetricName:
		return m.droppedRows
	case rowCountUnavailableMetricName:
		return m.unavailable
	default:
		return &benchmarkCounter{}
	}
}
func (m *benchmarkMeter) Gauge(_ string, _ ...string) sdk.Gauge { return nil }
func (m *benchmarkMeter) Histogram(_ string, _ []float64, _ ...string) sdk.Histogram {
	return nil
}

type benchmarkHost struct{ meter *benchmarkMeter }

func (h *benchmarkHost) Meter() sdk.Meter   { return h.meter }
func (h *benchmarkHost) Logger() sdk.Logger { return nil }

// BenchmarkDecide_TelemetryRegression compares the unchanged no-host path with
// batch-aggregated decision telemetry over both shipped schemas. It is intended
// for repeated benchstat runs when changing the sampler hot path.
func BenchmarkDecide_TelemetryRegression(b *testing.B) {
	for _, plugin := range benchPlugins() {
		config, batch := benchmarkDecisionMetricBatch(b, plugin)
		b.Run(plugin.name+"/disabled", func(b *testing.B) {
			sampler := mustSampler(b, config, plugin.schema)
			runDecide(b, sampler, batch)
		})
		b.Run(plugin.name+"/enabled", func(b *testing.B) {
			sampler := mustSampler(b, config, plugin.schema).(*Sampler)
			meter := &benchmarkMeter{
				decisions: &benchmarkCounter{}, rows: &benchmarkCounter{}, droppedRows: &benchmarkCounter{}, unavailable: &benchmarkCounter{},
			}
			sampler.UseHost(&benchmarkHost{meter: meter})
			runDecide(b, sampler, batch)
			expectedTotal := uint64(b.N * len(batch.Traces))
			if actualTotal := meter.decisions.total.Load(); actualTotal != expectedTotal {
				b.Fatalf("telemetry counter total %d, want %d", actualTotal, expectedTotal)
			}
			if totalRows := meter.rows.total.Load(); totalRows > 0 {
				b.ReportMetric(100*float64(meter.droppedRows.total.Load())/float64(totalRows), "row-drop-%")
			}
		})
	}
}

// benchmarkDecisionMetricBatch exercises every decision family and both tag
// decode paths in the telemetry regression benchmark.
func benchmarkDecisionMetricBatch(b *testing.B, plugin benchPlugin) (string, *sdk.TraceBatch) {
	b.Helper()
	config := fmt.Sprintf(`{
		"durationThresholdMs":%d,"keepErrors":true,"healthySampleRate":0.5,
		"keepTagRules":[
			{"tagKey":"missing","exists":true},
			{"tagKey":"db.type","equals":"PostgreSQL"},
			{"tagKey":"http.method","in":["POST","PUT"]}
		]}`, plugin.durationMs)
	healthyID := benchmarkTraceIDForSampleResult(true, 0.5)
	rejectedID := benchmarkTraceIDForSampleResult(false, 0.5)
	durationHit := plugin.durationMs*nanosPerMillis/plugin.schema.DurationTagNanosPerUnit + 1
	blocks := []sdk.TraceBlock{
		benchmarkDecisionMetricBlock(b, plugin, "duration", durationHit, false, nil),
		benchmarkDecisionMetricBlock(b, plugin, "error", 1, true, nil),
		benchmarkDecisionMetricBlock(b, plugin, "tag-equals", 1, false, []string{"db.type=PostgreSQL"}),
		benchmarkDecisionMetricBlock(b, plugin, "tag-in", 1, false, []string{"note=a|b", "http.method=PUT"}),
		benchmarkDecisionMetricBlock(b, plugin, healthyID, 1, false, nil),
		benchmarkDecisionMetricBlock(b, plugin, rejectedID, 1, false, nil),
	}
	const realisticBatchSize = 64
	seedCount := len(blocks)
	for len(blocks) < realisticBatchSize {
		blocks = append(blocks, blocks[len(blocks)%seedCount])
	}
	return config, sdktest.Batch(blocks...)
}

func benchmarkDecisionMetricBlock(b *testing.B, plugin benchPlugin, traceID string, duration int64, hasError bool, entries []string) sdk.TraceBlock {
	b.Helper()
	builder := sdktest.NewTrace(traceID).
		TagAs(plugin.startCol, valuetype.ValueTypeTimestamp, int64(0)).
		Tag(plugin.durCol, duration)
	if plugin.errCol != "" {
		errorValue := int64(0)
		if hasError {
			errorValue = 1
		}
		builder.Tag(plugin.errCol, errorValue)
	} else if hasError {
		entries = append(entries, "error")
	}
	if entries != nil {
		builder.Tag(plugin.arrayCol, entries)
	}
	block, buildErr := builder.Build()
	if buildErr != nil {
		b.Fatalf("build decision metric fixture: %v", buildErr)
	}
	return block
}

func benchmarkTraceIDForSampleResult(keep bool, rate float64) string {
	for candidateIndex := 0; ; candidateIndex++ {
		traceID := fmt.Sprintf("benchmark-sample-%d", candidateIndex)
		if (sampleFraction(traceID) < rate) == keep {
			return traceID
		}
	}
}
