// Copyright The OpenTelemetry Authors
// SPDX-License-Identifier: Apache-2.0

package tailsamplingprocessor // import "github.com/open-telemetry/opentelemetry-collector-contrib/processor/tailsamplingprocessor"

import (
	"errors"
	"fmt"
	"time"

	"go.opentelemetry.io/collector/component"

	"github.com/open-telemetry/opentelemetry-collector-contrib/pkg/ottl"
	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/tailsamplingprocessor/internal/tailstorageextension"
)

// PolicyType indicates the type of sampling policy.
type PolicyType string

const (
	// AlwaysSample samples all traces, typically used for debugging.
	AlwaysSample PolicyType = "always_sample"
	// Latency sample traces that are longer than a given threshold.
	Latency PolicyType = "latency"
	// NumericAttribute sample traces that have a given numeric attribute in a specified
	// range, e.g.: attribute "http.status_code" >= 399 and <= 999.
	NumericAttribute PolicyType = "numeric_attribute"
	// Probabilistic samples a given percentage of traces.
	Probabilistic PolicyType = "probabilistic"
	// StatusCode sample traces that have a given status code.
	StatusCode PolicyType = "status_code"
	// StringAttribute sample traces that an attribute, of type string, matching
	// one of the listed values.
	StringAttribute PolicyType = "string_attribute"
	// RateLimiting allows all traces until the specified limits are satisfied.
	RateLimiting PolicyType = "rate_limiting"
	// Composite allows defining a composite policy, combining the other policies in one
	Composite PolicyType = "composite"
	// And allows defining a And policy, combining the other policies in one
	And PolicyType = "and"
	// Not allows defining a Not policy, returning the opposite of the decision of a wrapped policy
	Not PolicyType = "not"
	// Drop allows defining a Drop policy, combining one or more policies to drop traces.
	Drop PolicyType = "drop"
	// SpanCount sample traces that are have more spans per Trace than a given threshold.
	SpanCount PolicyType = "span_count"
	// TraceState sample traces with specified values by the given key
	TraceState PolicyType = "trace_state"
	// BooleanAttribute sample traces having an attribute, of type bool, that matches
	// the specified boolean value [true|false].
	BooleanAttribute PolicyType = "boolean_attribute"
	// OTTLCondition sample traces which match user provided OpenTelemetry Transformation Language
	// conditions.
	OTTLCondition PolicyType = "ottl_condition"
	// BytesLimiting allows all traces until the specified byte limits are satisfied.
	BytesLimiting PolicyType = "bytes_limiting"
	// TraceFlags sample traces which have specific trace flags set.
	TraceFlags PolicyType = "trace_flags"
)

const (
	// samplingStrategyTraceComplete keeps the current tail-sampling behavior:
	// accumulate spans and decide on full trace data after decision timing.
	samplingStrategyTraceComplete samplingStrategy = "trace-complete"
	// samplingStrategySpanIngest evaluates each incoming span batch on ingest.
	// Non-terminal outcomes remain pending until cleanup finalization.
	samplingStrategySpanIngest samplingStrategy = "span-ingest"
)

type samplingStrategy string

// sharedPolicyCfg holds the common configuration to all policies that are used in derivative policy configurations
// such as the and & composite policies.
type sharedPolicyCfg struct {
	// Name given to the instance of the policy to make easy to identify it in metrics and logs.
	Name string `mapstructure:"name"`
	// Type of the policy this will be used to match the proper configuration of the policy.
	Type PolicyType `mapstructure:"type"`
	// Configs for latency filter sampling policy evaluator.
	LatencyCfg LatencyCfg `mapstructure:"latency"`
	// Configs for numeric attribute filter sampling policy evaluator.
	NumericAttributeCfg NumericAttributeCfg `mapstructure:"numeric_attribute"`
	// Configs for probabilistic sampling policy evaluator.
	ProbabilisticCfg ProbabilisticCfg `mapstructure:"probabilistic"`
	// Configs for status code filter sampling policy evaluator.
	StatusCodeCfg StatusCodeCfg `mapstructure:"status_code"`
	// Configs for string attribute filter sampling policy evaluator.
	StringAttributeCfg StringAttributeCfg `mapstructure:"string_attribute"`
	// Configs for rate limiting filter sampling policy evaluator.
	RateLimitingCfg RateLimitingCfg `mapstructure:"rate_limiting"`
	// Configs for bytes limiting filter sampling policy evaluator.
	BytesLimitingCfg BytesLimitingCfg `mapstructure:"bytes_limiting"`
	// Configs for span count filter sampling policy evaluator.
	SpanCountCfg SpanCountCfg `mapstructure:"span_count"`
	// Configs for defining trace_state policy
	TraceStateCfg TraceStateCfg `mapstructure:"trace_state"`
	// Configs for boolean attribute filter sampling policy evaluator.
	BooleanAttributeCfg BooleanAttributeCfg `mapstructure:"boolean_attribute"`
	// Configs for OTTL condition filter sampling policy evaluator
	OTTLConditionCfg OTTLConditionCfg `mapstructure:"ottl_condition"`
	// Configs for any extensions that are used.
	ExtensionCfg map[string]map[string]any `mapstructure:",remain"`
}

// CompositeSubPolicyCfg holds the common configuration to all policies under composite policy.
type CompositeSubPolicyCfg struct {
	sharedPolicyCfg `mapstructure:",squash"` // squash ensures fields are correctly decoded in embedded struct

	// Configs for and policy evaluator.
	AndCfg AndCfg `mapstructure:"and"`
}

// AndSubPolicyCfg holds the common configuration to all policies under and policy.
type AndSubPolicyCfg struct {
	sharedPolicyCfg `mapstructure:",squash"` // squash ensures fields are correctly decoded in embedded struct

	// Configs for defining not policy under and policy.
	NotCfg NotCfg `mapstructure:"not"`
}

// NotSubPolicyCfg holds the common configuration to the policy under the not policy.
type NotSubPolicyCfg struct {
	sharedPolicyCfg `mapstructure:",squash"` // squash ensures fields are correctly decoded in embedded struct
}

// TraceStateCfg holds the common configuration for trace states.
type TraceStateCfg struct {
	// Tag that the filter is going to be matching against.
	Key string `mapstructure:"key"`
	// Values indicate the set of values to use when matching against trace_state values.
	Values []string `mapstructure:"values"`
}

// AndCfg holds the common configuration to all and policies.
type AndCfg struct {
	SubPolicyCfg []AndSubPolicyCfg `mapstructure:"and_sub_policy"`
	// prevent unkeyed literal initialization
	_ struct{}
}

// NotCfg holds the configuration for the not policy.
type NotCfg struct {
	SubPolicy NotSubPolicyCfg `mapstructure:"not_sub_policy"`
}

// DropCfg holds the common configuration to all policies under drop policy.
type DropCfg struct {
	SubPolicyCfg []AndSubPolicyCfg `mapstructure:"drop_sub_policy"`
	// prevent unkeyed literal initialization
	_ struct{}
}

// CompositeCfg holds the configurable settings to create a composite
// sampling policy evaluator.
type CompositeCfg struct {
	MaxTotalSpansPerSecond int64                   `mapstructure:"max_total_spans_per_second"`
	PolicyOrder            []string                `mapstructure:"policy_order"`
	SubPolicyCfg           []CompositeSubPolicyCfg `mapstructure:"composite_sub_policy"`
	RateAllocation         []RateAllocationCfg     `mapstructure:"rate_allocation"`
}

// RateAllocationCfg used within composite policy
type RateAllocationCfg struct {
	Policy  string `mapstructure:"policy"`
	Percent int64  `mapstructure:"percent"`
	// prevent unkeyed literal initialization
	_ struct{}
}

// PolicyCfg holds the common configuration to all policies.
type PolicyCfg struct {
	sharedPolicyCfg `mapstructure:",squash"` // squash ensures fields are correctly decoded in embedded struct

	// Configs for defining composite policy
	CompositeCfg CompositeCfg `mapstructure:"composite"`
	// Configs for defining and policy
	AndCfg AndCfg `mapstructure:"and"`
	// Configs for defining not policy
	NotCfg NotCfg `mapstructure:"not"`
	// Configs for defining drop policy
	DropCfg DropCfg `mapstructure:"drop"`
}

// LatencyCfg holds the configurable settings to create a latency filter sampling policy
// evaluator
type LatencyCfg struct {
	// Lower bound in milliseconds. Retaining original name for compatibility
	ThresholdMs int64 `mapstructure:"threshold_ms"`
	// Upper bound in milliseconds.
	UpperThresholdMs int64 `mapstructure:"upper_threshold_ms"`
	// prevent unkeyed literal initialization
	_ struct{}
}

// NumericAttributeCfg holds the configurable settings to create a numeric attribute filter
// sampling policy evaluator.
type NumericAttributeCfg struct {
	// Tag that the filter is going to be matching against.
	Key string `mapstructure:"key"`
	// MinValue is the minimum value of the attribute to be considered a match.
	MinValue int64 `mapstructure:"min_value"`
	// MaxValue is the maximum value of the attribute to be considered a match.
	MaxValue int64 `mapstructure:"max_value"`
	// InvertMatch indicates that values must not match against attribute values.
	// If InvertMatch is true and Values is equal to '123', all other values will be sampled except '123'.
	// Also, if the specified Key does not match any resource or span attributes, data will be sampled.
	InvertMatch bool `mapstructure:"invert_match"`
}

// ProbabilisticCfg holds the configurable settings to create a probabilistic
// sampling policy evaluator.
type ProbabilisticCfg struct {
	// HashSalt allows one to configure the hashing salts. This is important in scenarios where multiple layers of collectors
	// have different sampling rates: if they use the same salt all passing one layer may pass the other even if they have
	// different sampling rates, configuring different salts avoids that.
	HashSalt string `mapstructure:"hash_salt"`
	// SamplingPercentage is the percentage rate at which traces are going to be sampled. Defaults to zero, i.e.: no sample.
	// Values greater or equal 100 are treated as "sample all traces".
	SamplingPercentage float64 `mapstructure:"sampling_percentage"`
	// prevent unkeyed literal initialization
	_ struct{}
}

// StatusCodeCfg holds the configurable settings to create a status code filter sampling
// policy evaluator.
type StatusCodeCfg struct {
	StatusCodes []string `mapstructure:"status_codes"`
	// prevent unkeyed literal initialization
	_ struct{}
}

// StringAttributeCfg holds the configurable settings to create a string attribute filter
// sampling policy evaluator.
type StringAttributeCfg struct {
	// Tag that the filter is going to be matching against.
	Key string `mapstructure:"key"`
	// Values indicate the set of values or regular expressions to use when matching against attribute values.
	// StringAttribute Policy will apply exact value match on Values unless EnabledRegexMatching is true.
	Values []string `mapstructure:"values"`
	// EnabledRegexMatching determines whether match attribute values by regexp string.
	EnabledRegexMatching bool `mapstructure:"enabled_regex_matching"`
	// CacheMaxSize is the maximum number of attribute entries of LRU Cache that stores the matched result
	// from the regular expressions defined in Values.
	// CacheMaxSize will not be used if EnabledRegexMatching is set to false.
	CacheMaxSize int `mapstructure:"cache_max_size"`
	// InvertMatch indicates that values or regular expressions must not match against attribute values.
	// If InvertMatch is true and Values is equal to 'acme', all other values will be sampled except 'acme'.
	// Also, if the specified Key does not match on any resource or span attributes, data will be sampled.
	InvertMatch bool `mapstructure:"invert_match"`
}

// RateLimitingCfg holds the configurable settings to create a rate limiting
// sampling policy evaluator.
type RateLimitingCfg struct {
	// SpansPerSecond sets the limit on the maximum number of spans that can be processed each second.
	SpansPerSecond int64 `mapstructure:"spans_per_second"`
	// BurstCapacity sets the maximum burst capacity in spans. If not specified, defaults to 2x SpansPerSecond.
	// This allows for short bursts of traffic above the sustained rate. It also acts as a
	// limit for individual trace span counts, a single trace with more spans than the burst size will not pass.
	BurstCapacity int64 `mapstructure:"burst_capacity"`
	// prevent unkeyed literal initialization
	_ struct{}
}

// BytesLimitingCfg holds the configurable settings to create a bytes limiting
// sampling policy evaluator using a token bucket algorithm.
type BytesLimitingCfg struct {
	// BytesPerSecond sets the limit on the maximum number of bytes that can be processed each second.
	BytesPerSecond int64 `mapstructure:"bytes_per_second"`
	// BurstCapacity sets the maximum burst capacity in bytes. If not specified, defaults to 2x BytesPerSecond.
	// This allows for short bursts of traffic above the sustained rate. It also acts as a
	// limit for individual trace sizes, a single trace larger than the burst size will not pass.
	BurstCapacity int64 `mapstructure:"burst_capacity"`
}

// SpanCountCfg holds the configurable settings to create a Span Count filter sampling
// policy evaluator
type SpanCountCfg struct {
	// Minimum number of spans in a Trace
	MinSpans int32 `mapstructure:"min_spans"`
	MaxSpans int32 `mapstructure:"max_spans"`
	// prevent unkeyed literal initialization
	_ struct{}
}

// BooleanAttributeCfg holds the configurable settings to create a boolean attribute filter
// sampling policy evaluator.
type BooleanAttributeCfg struct {
	// Tag that the filter is going to be matching against.
	Key string `mapstructure:"key"`
	// Value indicate the bool value, either true or false to use when matching against attribute values.
	// BooleanAttribute Policy will apply exact value match on Value
	Value bool `mapstructure:"value"`
	// InvertMatch indicates that values must not match against attribute values.
	// If InvertMatch is true and Values is equal to 'true', all other values will be sampled except 'true'.
	// Also, if the specified Key does not match any resource or span attributes, data will be sampled.
	InvertMatch bool `mapstructure:"invert_match"`
}

// OTTLConditionCfg holds the configurable setting to create a OTTL condition filter
// sampling policy evaluator.
type OTTLConditionCfg struct {
	ErrorMode           ottl.ErrorMode `mapstructure:"error_mode"`
	SpanConditions      []string       `mapstructure:"span"`
	SpanEventConditions []string       `mapstructure:"spanevent"`
	// prevent unkeyed literal initialization
	_ struct{}
}

type DecisionCacheConfig struct {
	// SampledCacheSize specifies the size of the cache that holds the sampled trace IDs.
	// This value will be the maximum amount of trace IDs that the cache can hold before overwriting previous IDs.
	// For effective use, this value should be at least an order of magnitude greater than Config.NumTraces.
	// If left as default 0, a no-op DecisionCache will be used.
	SampledCacheSize int `mapstructure:"sampled_cache_size"`
	// NonSampledCacheSize specifies the size of the cache that holds the non-sampled trace IDs.
	// This value will be the maximum amount of trace IDs that the cache can hold before overwriting previous IDs.
	// For effective use, this value should be at least an order of magnitude greater than Config.NumTraces.
	// If left as default 0, a no-op DecisionCache will be used.
	NonSampledCacheSize int `mapstructure:"non_sampled_cache_size"`
	// prevent unkeyed literal initialization
	_ struct{}
}

// Config holds the configuration for tail-based sampling.
type Config struct {
	// DecisionWait is the time before timer handling for a trace.
	// When sampling_strategy is "trace-complete", this controls decision timing.
	// When sampling_strategy is "span-ingest", this controls pending cleanup finalization timing.
	DecisionWait time.Duration `mapstructure:"decision_wait"`
	// DecisionWaitAfterRootReceived adds root-span-based acceleration for timer handling.
	// When sampling_strategy is "trace-complete", this can make decisions earlier.
	// When sampling_strategy is "span-ingest", this can finalize pending traces earlier on cleanup.
	DecisionWaitAfterRootReceived time.Duration `mapstructure:"decision_wait_after_root_received"`
	// NumTraces is the number of traces kept on memory. Typically most of the data
	// of a trace is released after a sampling decision is taken.
	NumTraces uint64 `mapstructure:"num_traces"`
	// BlockOnOverflow determines the behavior when the component's NumTraces limit is reached.
	// If true, the component will wait for space; otherwise, old traces will be evicted to make space.
	BlockOnOverflow bool `mapstructure:"block_on_overflow"`
	// ExpectedNewTracesPerSec sets the expected number of new traces sending to the tail sampling processor
	// per second. This helps with allocating data structures with closer to actual usage size.
	ExpectedNewTracesPerSec uint64 `mapstructure:"expected_new_traces_per_sec"`
	// PolicyCfgs sets the tail-based sampling policy which makes a sampling decision
	// for a given trace when requested.
	PolicyCfgs []PolicyCfg `mapstructure:"policies"`
	// DecisionCache holds configuration for the decision cache(s)
	DecisionCache DecisionCacheConfig `mapstructure:"decision_cache"`
	// TailStorageID specifies an optional tail storage extension to use for buffering spans.
	// If not set, in-memory tail storage is used.
	// It is behind feature gate `processor.tailsamplingprocessor.tailstorageextension`.
	TailStorageID *component.ID `mapstructure:"tail_storage"`
	// Options allows for additional configuration of the tail-based sampling processor in code.
	Options []Option `mapstructure:"-"`
	// Make decision as soon as a policy matches
	SampleOnFirstMatch bool `mapstructure:"sample_on_first_match"`
	// SamplingStrategy controls how/when sampling decisions are made.
	// "trace-complete" (default) evaluates accumulated trace data on timer handling.
	// "span-ingest" evaluates each incoming batch on ingest; terminal outcomes
	// finalize immediately, and non-terminal traces are finalized on cleanup.
	SamplingStrategy samplingStrategy `mapstructure:"sampling_strategy"`
	// DropPendingTracesOnShutdown will drop all traces that are part of batches that have not yet reached the decision
	// wait when the processor is shutdown.
	DropPendingTracesOnShutdown bool `mapstructure:"drop_pending_traces_on_shutdown"`
	// MaximumTraceSizeBytes is the largest size of a trace a decision will be made for.
	// If the trace size exceeds this it will be dropped before the decision period to keep memory more predictable.
	// A 0 value disables dropping large traces early.
	MaximumTraceSizeBytes uint64 `mapstructure:"maximum_trace_size_bytes"`
	// NumShards controls the number of parallel goroutine loops processing
	// traces. Each shard runs an independent event loop with its own trace
	// storage and decision batcher. Traces are routed to shards by a hash of
	// the trace ID, ensuring all spans for a given trace are processed by
	// the same shard. Higher values reduce contention between trace
	// ingestion and sampling decision evaluation under high load. NumTraces,
	// ExpectedNewTracesPerSec, decision cache sizes, and per-second rate
	// limits in policies (rate_limiting, bytes_limiting, and composite
	// max_total_spans_per_second) are divided evenly across shards so
	// aggregate behavior matches the configured values. Limiter burst
	// capacities are not divided so single large traces stay admissible
	// regardless of the shard count. Must not exceed 256,
	// and values greater than 1 are not supported with tail_storage.
	// Defaults to 1 (single event loop, original behavior).
	NumShards uint32 `mapstructure:"num_shards"`
}

// maxNumShards bounds num_shards to catch configuration mistakes: each shard
// materializes a full event loop with its own goroutine, storage and batcher,
// so values beyond CPU-count scale only add overhead.
const maxNumShards = 256

func (cfg *Config) Validate() error {
	switch cfg.SamplingStrategy {
	case samplingStrategyTraceComplete, samplingStrategySpanIngest:
		// valid sampling strategies
	default:
		return fmt.Errorf(
			"invalid sampling_strategy %q, expected one of %q or %q",
			cfg.SamplingStrategy,
			samplingStrategyTraceComplete,
			samplingStrategySpanIngest,
		)
	}

	if cfg.NumShards > maxNumShards {
		return fmt.Errorf("num_shards (%d) must not exceed %d", cfg.NumShards, maxNumShards)
	}

	// The TailStorage contract makes the caller responsible for serializing
	// access, which multiple shard event loops cannot guarantee for a shared
	// extension instance. Sharding support belongs in the storage layer.
	if cfg.NumShards > 1 && cfg.TailStorageID != nil {
		return errors.New("num_shards greater than 1 is not supported with tail_storage")
	}

	if cfg.TailStorageID != nil && !tailstorageextension.IsFeatureGateEnabled() {
		return fmt.Errorf(
			"'tail_storage' requires the %q feature gate to be enabled, use --feature-gates=+%s",
			tailstorageextension.FeatureGateID,
			tailstorageextension.FeatureGateID,
		)
	}

	return nil
}