// SPDX-License-Identifier: AGPL-3.0-only

package ingest

import (
	"crypto/tls"
	"errors"
	"flag"
	"fmt"
	"math"
	"slices"
	"strconv"
	"strings"
	"time"

	"github.com/grafana/dskit/backoff"
	dskittls "github.com/grafana/dskit/crypto/tls"
	"github.com/grafana/dskit/flagext"

	"github.com/grafana/mimir/pkg/util"
)

const (
	consumeFromLastOffset = "last-offset"
	consumeFromStart      = "start"
	consumeFromEnd        = "end"
	consumeFromTimestamp  = "timestamp"
)

var (
	ErrMissingKafkaAddress               = errors.New("the Kafka address has not been configured")
	ErrMissingKafkaTopic                 = errors.New("the Kafka topic has not been configured")
	ErrInvalidConsumePosition            = errors.New("the configured consume position is invalid")
	ErrInvalidProducerMaxRecordSizeBytes = fmt.Errorf("the configured producer max record size bytes must be a value between %d and %d", minProducerRecordDataBytesLimit, maxProducerRecordDataBytesLimit)
	ErrInconsistentConsumerLagAtStartup  = fmt.Errorf("the target and max consumer lag at startup must be either both set to 0 or to a value greater than 0")
	ErrInvalidMaxConsumerLagAtStartup    = fmt.Errorf("the configured max consumer lag at startup must greater or equal than the configured target consumer lag")
	ErrInconsistentSASLCredentials       = fmt.Errorf("the SASL username and password must be both configured to enable SASL authentication")
	ErrSASLOauthbearerBadConfig          = fmt.Errorf("exactly one of OAuth token, file path, or HTTP socket path must be configured to enable SASL OAUTHBEARER authentication")
	ErrSASLMSKIAMBadConfig               = fmt.Errorf("exactly one of static credentials, file path, or HTTP socket path must be configured to enable SASL AWS_MSK_IAM authentication")
	ErrInvalidSASLMechanism              = fmt.Errorf("the configured SASL mechanism is invalid, must be one of: %s", util.JoinStrings(saslMechanismOptions, ", "))

	ErrInvalidIngestionConcurrencyMax    = errors.New("ingest-storage.kafka.ingestion-concurrency-max must either be set to 0 or to a value greater than 0")
	ErrInvalidIngestionConcurrencyParams = errors.New("ingest-storage.kafka.ingestion-concurrency-queue-capacity, ingest-storage.kafka.ingestion-concurrency-estimated-bytes-per-sample, ingest-storage.kafka.ingestion-concurrency-batch-size and ingest-storage.kafka.ingestion-concurrency-target-flushes-per-shard must be greater than 0")
	ErrInvalidAutoCreateTopicParams      = errors.New("ingest-storage.kafka.auto-create-topic-default-partitions must be -1 or greater than 0 when ingest-storage.kafka.auto-create-topic-default-partitions=true")
	ErrInvalidFetchMaxWait               = errors.New("the Kafka fetch max wait must be between 5s and 30s")
	ErrInvalidRecordVersion              = errors.New("invalid record format version")
	ErrInvalidWriteLogsFsyncConcurrency  = errors.New("the configured number of tenants to fsync concurrently before Kafka offsets are committed must be at least 1")

	consumeFromPositionOptions = []string{consumeFromLastOffset, consumeFromStart, consumeFromEnd, consumeFromTimestamp}

	defaultFetchBackoffConfig = backoff.Config{
		MinBackoff: 250 * time.Millisecond,
		MaxBackoff: 2 * time.Second,
		MaxRetries: 0, // Retry forever. Do NOT change!
	}
)

type Config struct {
	Enabled     bool            `yaml:"enabled"`
	KafkaConfig KafkaConfig     `yaml:"kafka"`
	Migration   MigrationConfig `yaml:"migration"`

	WriteLogsFsyncBeforeKafkaCommitConcurrency int `yaml:"write_logs_fsync_before_kafka_commit_concurrency" category:"advanced"`
}

func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
	f.BoolVar(&cfg.Enabled, "ingest-storage.enabled", false, "True to enable the ingestion via object storage.")
	f.IntVar(&cfg.WriteLogsFsyncBeforeKafkaCommitConcurrency, "ingest-storage.write-logs-fsync-before-kafka-commit-concurrency", 4, "Number of tenants to concurrently fsync WAL and WBL before Kafka offsets are committed, must be at least 1.")

	cfg.KafkaConfig.RegisterFlagsWithPrefix("ingest-storage.kafka.", f)
	cfg.Migration.RegisterFlagsWithPrefix("ingest-storage.migration.", f)
}

// Validate the config.
func (cfg *Config) Validate() error {
	// Skip validation if disabled.
	if !cfg.Enabled {
		return nil
	}

	if cfg.WriteLogsFsyncBeforeKafkaCommitConcurrency < 1 {
		return fmt.Errorf("%w, is %d", ErrInvalidWriteLogsFsyncConcurrency, cfg.WriteLogsFsyncBeforeKafkaCommitConcurrency)
	}

	if err := cfg.KafkaConfig.Validate(); err != nil {
		return err
	}

	return nil
}

// KafkaConfig holds the generic config for the Kafka backend.
type KafkaConfig struct {
	// Address is a list of seed brokers. The config name is singular for backward compatibility.
	Address      flagext.StringSliceCSV `yaml:"address"`
	Topic        string                 `yaml:"topic"`
	ClientID     string                 `yaml:"client_id"`
	ClientRack   string                 `yaml:"client_rack"`
	DialTimeout  time.Duration          `yaml:"dial_timeout"`
	WriteTimeout time.Duration          `yaml:"write_timeout"`
	WriteClients int                    `yaml:"write_clients" category:"deprecated"` // TODO Remove in Mimir 3.3.

	SASL       KafkaAuthConfig `yaml:",inline"`
	TLSEnabled bool            `yaml:"tls_enabled"`
	TLS        TLSClientConfig `yaml:",inline"`

	ConsumerGroup                         string        `yaml:"consumer_group"`
	ConsumerGroupOffsetCommitInterval     time.Duration `yaml:"consumer_group_offset_commit_interval"`
	ConsumerGroupOffsetCommitFileEnforced bool          `yaml:"consumer_group_offset_commit_file_enforced" category:"experimental"`

	LastProducedOffsetPollInterval time.Duration `yaml:"last_produced_offset_poll_interval"`
	LastProducedOffsetRetryTimeout time.Duration `yaml:"last_produced_offset_retry_timeout"`

	ConsumeFromPositionAtStartup  string        `yaml:"consume_from_position_at_startup"`
	ConsumeFromTimestampAtStartup int64         `yaml:"consume_from_timestamp_at_startup"`
	TargetConsumerLagAtStartup    time.Duration `yaml:"target_consumer_lag_at_startup"`
	MaxConsumerLagAtStartup       time.Duration `yaml:"max_consumer_lag_at_startup"`

	// MaxReplayPeriod is used when file-based offset enforcement is enabled but no offset file exists. It specifies
	//how far back to replay data; typically set to the TSDB retention period.
	MaxReplayPeriod time.Duration `yaml:"-"`

	AutoCreateTopicEnabled           bool `yaml:"auto_create_topic_enabled"`
	AutoCreateTopicDefaultPartitions int  `yaml:"auto_create_topic_default_partitions"`

	ProducerMaxRecordSizeBytes int   `yaml:"producer_max_record_size_bytes"`
	ProducerMaxBufferedBytes   int64 `yaml:"producer_max_buffered_bytes"`

	WaitStrongReadConsistencyTimeout time.Duration `yaml:"wait_strong_read_consistency_timeout"`

	ProducerRecordVersion int `yaml:"producer_record_version"`

	// Used when logging unsampled client errors. Set from ingester's ErrorSampleRate.
	FallbackClientErrorSampleRate int64 `yaml:"-"`

	FetchMaxWait                      time.Duration `yaml:"fetch_max_wait"`
	FetchConcurrencyMax               int           `yaml:"fetch_concurrency_max"`
	UseCompressedBytesAsFetchMaxBytes bool          `yaml:"use_compressed_bytes_as_fetch_max_bytes"`
	MaxBufferedBytes                  int           `yaml:"max_buffered_bytes"`

	IngestionConcurrencyMax       int `yaml:"ingestion_concurrency_max"`
	IngestionConcurrencyBatchSize int `yaml:"ingestion_concurrency_batch_size"`

	// IngestionConcurrencyQueueCapacity controls how many batches can be enqueued for flushing series to the TSDB HEAD.
	// We don't want to push any batches in parallel and instead want to prepare the next ones while the current one finishes, hence a small buffer.
	// For example, if we flush 1 batch/sec, then batching 2 batches/sec doesn't make us faster.
	// If there's high variability in the time to flush or in the time to batch, then this buffer might need to be increased.
	IngestionConcurrencyQueueCapacity int `yaml:"ingestion_concurrency_queue_capacity"`

	// IngestionConcurrencyTargetFlushesPerShard is the number of flushes we want to target per shard.
	// There is some overhead in the parallelization. With fewer flushes, the overhead of splitting up the work is higher than the benefit of parallelization.
	// The default was devised experimentally to keep the memory and CPU usage low during ongoing consumption, while keeping replay speed high during cold replay.
	IngestionConcurrencyTargetFlushesPerShard int `yaml:"ingestion_concurrency_target_flushes_per_shard"`

	// IngestionConcurrencyEstimatedBytesPerSample is the estimated number of bytes per sample.
	// Our data indicates that the average sample size is somewhere between ~100 and ~200 bytes. We'll use 200 bytes as a conservative estimate.
	IngestionConcurrencyEstimatedBytesPerSample int `yaml:"ingestion_concurrency_estimated_bytes_per_sample"`

	// The fetch backoff config to use in the concurrent fetchers (when enabled). This setting
	// is just used to change the default backoff in tests.
	concurrentFetchersFetchBackoffConfig backoff.Config `yaml:"-"`

	// DisableLinger disables producer linger. This setting is just used in tests.
	DisableLinger bool `yaml:"-"`

	// MaxInflightProduceRequests is the max number of in-flight Produce requests per broker.
	// This setting is just used in tests. 0 uses the default.
	MaxInflightProduceRequests int `yaml:"-"`
}

func (cfg *KafkaConfig) RegisterFlags(f *flag.FlagSet) {
	cfg.RegisterFlagsWithPrefix("", f)
}

func (cfg *KafkaConfig) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
	cfg.concurrentFetchersFetchBackoffConfig = defaultFetchBackoffConfig

	f.Var(&cfg.Address, prefix+"address", "The Kafka seed broker address, or a comma-separated list of seed broker addresses.")
	f.StringVar(&cfg.Topic, prefix+"topic", "", "The Kafka topic name.")
	f.StringVar(&cfg.ClientID, prefix+"client-id", "", "The Kafka client ID.")
	f.StringVar(&cfg.ClientRack, prefix+"client-rack", "", "The rack identifier for this Kafka client. Corresponds to the Kafka client.rack setting.")
	f.DurationVar(&cfg.DialTimeout, prefix+"dial-timeout", 2*time.Second, "The maximum time allowed to open a connection to a Kafka broker.")
	f.DurationVar(&cfg.WriteTimeout, prefix+"write-timeout", 10*time.Second, "How long to wait for an incoming write request to be successfully committed to the Kafka backend.")
	f.IntVar(&cfg.WriteClients, prefix+"write-clients", 1, "The number of Kafka clients used by producers. When the configured number of clients is greater than 1, partitions are sharded among Kafka clients. A higher number of clients may provide higher write throughput at the cost of additional Metadata requests pressure to Kafka. Deprecated: has no effect (Mimir always uses a single Kafka write client).")

	f.StringVar(&cfg.ConsumerGroup, prefix+"consumer-group", "", "The consumer group used by the consumer to track the last consumed offset. The consumer group must be different for each ingester. If the configured consumer group contains the '<partition>' placeholder, it is replaced with the actual partition ID owned by the ingester. When empty (recommended), Mimir uses the ingester instance ID to guarantee uniqueness.")
	f.DurationVar(&cfg.ConsumerGroupOffsetCommitInterval, prefix+"consumer-group-offset-commit-interval", time.Second, "How frequently a consumer should commit the consumed offset to Kafka. The last committed offset is used at startup to continue the consumption from where it was left.")
	f.BoolVar(&cfg.ConsumerGroupOffsetCommitFileEnforced, prefix+"consumer-group-offset-commit-file-enforced", false, "When true, the file-based offset stored in the TSDB directory is enforced at startup, taking precedence over Kafka consumer group offset. When false, offsets are still written to the file (in the TSDB directory) but the Kafka consumer group offset is used at startup.")

	f.DurationVar(&cfg.LastProducedOffsetPollInterval, prefix+"last-produced-offset-poll-interval", time.Second, "How frequently to poll the last produced offset, used to enforce strong read consistency.")
	f.DurationVar(&cfg.LastProducedOffsetRetryTimeout, prefix+"last-produced-offset-retry-timeout", 10*time.Second, "How long to retry a failed request to get the last produced offset.")

	f.StringVar(&cfg.ConsumeFromPositionAtStartup, prefix+"consume-from-position-at-startup", consumeFromLastOffset, fmt.Sprintf("From which position to start consuming the partition at startup. Supported options: %s.", strings.Join(consumeFromPositionOptions, ", ")))
	f.Int64Var(&cfg.ConsumeFromTimestampAtStartup, prefix+"consume-from-timestamp-at-startup", 0, fmt.Sprintf("Milliseconds timestamp after which the consumption of the partition starts at startup. Only applies when consume-from-position-at-startup is %s", consumeFromTimestamp))

	targetConsumerLagAtStartupFlag := prefix + "target-consumer-lag-at-startup"
	maxConsumerLagAtStartupFlag := prefix + "max-consumer-lag-at-startup"
	howToDisableConsumerLagAtStartup := fmt.Sprintf("Set both -%s and -%s to 0 to disable waiting for maximum consumer lag being honored at startup.", targetConsumerLagAtStartupFlag, maxConsumerLagAtStartupFlag)
	f.DurationVar(&cfg.TargetConsumerLagAtStartup, targetConsumerLagAtStartupFlag, 2*time.Second, "The best-effort maximum lag a consumer tries to achieve at startup. "+howToDisableConsumerLagAtStartup)
	f.DurationVar(&cfg.MaxConsumerLagAtStartup, maxConsumerLagAtStartupFlag, 15*time.Second, "The guaranteed maximum lag before a consumer is considered to have caught up reading from a partition at startup, becomes ACTIVE in the hash ring and passes the readiness check. "+howToDisableConsumerLagAtStartup)

	f.BoolVar(&cfg.AutoCreateTopicEnabled, prefix+"auto-create-topic-enabled", true, "Enable auto-creation of Kafka topic on startup if it doesn't exist. If creating the topic fails and the topic doesn't already exist, Mimir fails to start.")
	f.IntVar(&cfg.AutoCreateTopicDefaultPartitions, prefix+"auto-create-topic-default-partitions", -1, "When auto-creation of Kafka topic is enabled and this value is positive, Mimir creates the topic with this number of partitions. When the value is -1 the Kafka broker uses the default number of partitions (num.partitions configuration).")

	f.IntVar(&cfg.ProducerMaxRecordSizeBytes, prefix+"producer-max-record-size-bytes", maxProducerRecordDataBytesLimit, "The maximum size of a Kafka record data that should be generated by the producer. An incoming write request larger than this size is split into multiple Kafka records. We strongly recommend to not change this setting unless for testing purposes.")
	f.Int64Var(&cfg.ProducerMaxBufferedBytes, prefix+"producer-max-buffered-bytes", 1024*1024*1024, "The maximum size of (uncompressed) buffered and unacknowledged produced records sent to Kafka. The produce request fails once this limit is reached. This limit is per Kafka client. 0 to disable the limit.")

	f.DurationVar(&cfg.WaitStrongReadConsistencyTimeout, prefix+"wait-strong-read-consistency-timeout", 20*time.Second, "The maximum allowed for a read requests processed by an ingester to wait until strong read consistency is enforced. 0 to disable the timeout.")

	f.IntVar(&cfg.ProducerRecordVersion, prefix+"producer-record-version", 2, "The record version that this producer sends.")

	f.DurationVar(&cfg.FetchMaxWait, prefix+"fetch-max-wait", 5*time.Second, "The maximum amount of time a Kafka broker waits for some records before a Fetch response is returned.")
	f.IntVar(&cfg.FetchConcurrencyMax, prefix+"fetch-concurrency-max", 12, "The maximum number of concurrent fetch requests that the ingester makes when reading data from Kafka during startup. Concurrent fetch requests are issued only when there is sufficient backlog of records to consume. Set to 0 to disable.")
	f.BoolVar(&cfg.UseCompressedBytesAsFetchMaxBytes, prefix+"use-compressed-bytes-as-fetch-max-bytes", true, "When enabled, the fetch request MaxBytes field is computed using the compressed size of previous records. When disabled, MaxBytes is computed using uncompressed bytes. Different Kafka implementations interpret MaxBytes differently.")
	f.IntVar(&cfg.MaxBufferedBytes, prefix+"max-buffered-bytes", 1_000_000_000, "The maximum number of buffered records ready to be processed. This limit applies to the sum of all inflight requests. Set to 0 to disable the limit.")

	f.IntVar(&cfg.IngestionConcurrencyMax, prefix+"ingestion-concurrency-max", 8, "The maximum number of concurrent ingestion streams to the TSDB head. Every tenant has their own set of streams. 0 to disable.")
	f.IntVar(&cfg.IngestionConcurrencyBatchSize, prefix+"ingestion-concurrency-batch-size", 150, "The number of timeseries to batch together before ingesting to the TSDB head. Only use this setting when -ingest-storage.kafka.ingestion-concurrency-max is greater than 0.")
	f.IntVar(&cfg.IngestionConcurrencyQueueCapacity, prefix+"ingestion-concurrency-queue-capacity", 3, "The number of batches to prepare and queue to ingest to the TSDB head. Only use this setting when -ingest-storage.kafka.ingestion-concurrency-max is greater than 0.")
	f.IntVar(&cfg.IngestionConcurrencyTargetFlushesPerShard, prefix+"ingestion-concurrency-target-flushes-per-shard", 40, "The expected number of times to ingest timeseries to the TSDB head after batching. With fewer flushes, the overhead of splitting up the work is higher than the benefit of parallelization. Only use this setting when -ingest-storage.kafka.ingestion-concurrency-max is greater than 0.")
	f.IntVar(&cfg.IngestionConcurrencyEstimatedBytesPerSample, prefix+"ingestion-concurrency-estimated-bytes-per-sample", 200, "The estimated number of bytes a sample has at time of ingestion. This value is used to estimate the timeseries without decompressing them. Only use this setting when -ingest-storage.kafka.ingestion-concurrency-max is greater than 0.")

	cfg.SASL.RegisterFlagsWithPrefix(prefix+"sasl-", f)
	f.BoolVar(&cfg.TLSEnabled, prefix+"tls-enabled", false, "Enable TLS for the Kafka client connection.")
	cfg.TLS.RegisterFlagsWithPrefix(prefix, f)
}

func (cfg *KafkaConfig) Validate() error {
	if cfg.Address.String() == "" {
		return ErrMissingKafkaAddress
	}
	if cfg.Topic == "" {
		return ErrMissingKafkaTopic
	}
	if !slices.Contains(consumeFromPositionOptions, cfg.ConsumeFromPositionAtStartup) {
		return ErrInvalidConsumePosition
	}
	if cfg.ConsumeFromPositionAtStartup == consumeFromTimestamp {
		// We only do a simple soundness check for the value be a millisecond precision timestamp.
		if cfg.ConsumeFromTimestampAtStartup < 1e12 {
			return fmt.Errorf("%w: configured timestamp must be a millisecond timestamp", ErrInvalidConsumePosition)
		}
	} else {
		if cfg.ConsumeFromTimestampAtStartup > 0 {
			return fmt.Errorf("%w: configured consume position must be set to %q", ErrInvalidConsumePosition, consumeFromTimestamp)
		}
	}
	if cfg.ProducerMaxRecordSizeBytes < minProducerRecordDataBytesLimit || cfg.ProducerMaxRecordSizeBytes > maxProducerRecordDataBytesLimit {
		return ErrInvalidProducerMaxRecordSizeBytes
	}
	if (cfg.TargetConsumerLagAtStartup != 0) != (cfg.MaxConsumerLagAtStartup != 0) {
		return ErrInconsistentConsumerLagAtStartup
	}
	if cfg.MaxConsumerLagAtStartup < cfg.TargetConsumerLagAtStartup {
		return ErrInvalidMaxConsumerLagAtStartup
	}

	if err := ValidateRecordVersion(cfg.ProducerRecordVersion); err != nil {
		return fmt.Errorf("%w: %w", ErrInvalidRecordVersion, err)
	}

	if cfg.FetchMaxWait < 5*time.Second || cfg.FetchMaxWait > 30*time.Second {
		return ErrInvalidFetchMaxWait
	}

	if cfg.FetchConcurrencyMax < 0 {
		return fmt.Errorf("ingest-storage.kafka.fetch-concurrency-max must be greater than or equal to 0")
	}

	if cfg.MaxBufferedBytes >= math.MaxInt32 {
		return fmt.Errorf("ingest-storage.kafka.max-buffered-bytes must be less than %d", math.MaxInt32)
	}

	if cfg.IngestionConcurrencyMax < 0 {
		return ErrInvalidIngestionConcurrencyMax
	}

	if cfg.IngestionConcurrencyMax >= 1 {
		if cfg.IngestionConcurrencyBatchSize <= 0 || cfg.IngestionConcurrencyQueueCapacity <= 0 || cfg.IngestionConcurrencyEstimatedBytesPerSample <= 0 || cfg.IngestionConcurrencyTargetFlushesPerShard <= 0 {
			return ErrInvalidIngestionConcurrencyParams
		}
	}

	if partitions := cfg.AutoCreateTopicDefaultPartitions; cfg.AutoCreateTopicEnabled && (partitions != -1 && partitions < 1) {
		return ErrInvalidAutoCreateTopicParams
	}

	if err := cfg.SASL.Validate(); err != nil {
		return err
	}

	if cfg.TLSEnabled {
		if _, err := cfg.TLS.GetTLSConfig(); err != nil {
			return fmt.Errorf("invalid Kafka TLS config: %w", err)
		}
	}

	return nil
}

// GetConsumerGroup returns the consumer group to use for the given instanceID and partitionID.
func (cfg *KafkaConfig) GetConsumerGroup(instanceID string, partitionID int32) string {
	if cfg.ConsumerGroup == "" {
		return instanceID
	}

	return strings.ReplaceAll(cfg.ConsumerGroup, "<partition>", strconv.Itoa(int(partitionID)))
}

// MigrationConfig holds the configuration used to migrate Mimir to ingest storage. This config shouldn't be
// set for any other reason.
type MigrationConfig struct {
	DistributorSendToIngestersEnabled bool          `yaml:"distributor_send_to_ingesters_enabled"`
	IgnoreIngestStorageErrors         bool          `yaml:"ignore_ingest_storage_errors"`
	IngestStorageMaxWaitTime          time.Duration `yaml:"ingest_storage_max_wait_time"`
}

func (cfg *MigrationConfig) RegisterFlags(f *flag.FlagSet) {
	cfg.RegisterFlagsWithPrefix("", f)
}

func (cfg *MigrationConfig) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
	f.BoolVar(&cfg.DistributorSendToIngestersEnabled, prefix+"distributor-send-to-ingesters-enabled", false, "When both this option and ingest storage are enabled, distributors write to both Kafka and ingesters. A write request is considered successful only when written to both backends.")
	f.BoolVar(&cfg.IgnoreIngestStorageErrors, prefix+"ignore-ingest-storage-errors", false, "When enabled, errors writing to ingest storage are logged but do not affect write success or quorum. When disabled, write requests fail if ingest storage write fails.")
	f.DurationVar(&cfg.IngestStorageMaxWaitTime, prefix+"ingest-storage-max-wait-time", 0, "The maximum time a write request that goes through the ingest storage waits before it times out. Set to `0` to disable the timeout.")
}

type SASLMechanism string

// Set implements flag.Value.
func (s *SASLMechanism) Set(v string) error {
	if !slices.Contains(saslMechanismOptions, SASLMechanism(v)) {
		return ErrInvalidSASLMechanism
	}
	*s = SASLMechanism(v)
	return nil
}

// String implements flag.Value.
func (s *SASLMechanism) String() string {
	return string(*s)
}

// FlagType implements usage.FlagTyper.
func (SASLMechanism) FlagType() string {
	return "string"
}

const (
	SASLMechanismPlain       SASLMechanism = "PLAIN"
	SASLMechanismScramSHA256 SASLMechanism = "SCRAM-SHA-256"
	SASLMechanismScramSHA512 SASLMechanism = "SCRAM-SHA-512"
	SASLMechanismOauthbearer SASLMechanism = "OAUTHBEARER"
	SASLMechanismMSKIAM      SASLMechanism = "AWS_MSK_IAM"
)

var saslMechanismOptions = []SASLMechanism{SASLMechanismPlain, SASLMechanismScramSHA256, SASLMechanismScramSHA512, SASLMechanismOauthbearer, SASLMechanismMSKIAM}

type KafkaAuthConfig struct {
	Mechanism SASLMechanism `yaml:"sasl_mechanism"`

	// For PLAIN and SCRAM-SHA-* mechanisms

	Username string         `yaml:"sasl_username"`
	Password flagext.Secret `yaml:"sasl_password"`

	Oauthbearer KafkaAuthOauthbearerConfig `yaml:",inline"`

	MSKIAM KafkaAuthMSKIAMConfig `yaml:",inline"`
}

func (cfg *KafkaAuthConfig) RegisterFlags(f *flag.FlagSet) {
	cfg.RegisterFlagsWithPrefix("", f)
}

func (cfg *KafkaAuthConfig) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
	cfg.Mechanism = "PLAIN"
	f.Var(&cfg.Mechanism, prefix+"mechanism", fmt.Sprintf("The SASL mechanism used to authenticate to Kafka. Supported values: %s. For backwards-compatibility, PLAIN with no username nor password disables SASL.", util.JoinStrings(saslMechanismOptions, ", ")))
	f.StringVar(&cfg.Username, prefix+"username", "", "The username used to authenticate to Kafka using SASL. To enable SASL, configure both the username and password.")
	f.Var(&cfg.Password, prefix+"password", "The password used to authenticate to Kafka using SASL. To enable SASL, configure both the username and password.")
	cfg.Oauthbearer.RegisterFlagsWithPrefix(prefix+"oauthbearer-", f)
	cfg.MSKIAM.RegisterFlagsWithPrefix(prefix+"msk-iam-", f)
}

func (cfg *KafkaAuthConfig) Validate() error {
	switch cfg.Mechanism {
	case SASLMechanismPlain:
		if (cfg.Username == "") != (cfg.Password.String() == "") {
			return ErrInconsistentSASLCredentials
		}

	case SASLMechanismScramSHA256, SASLMechanismScramSHA512:
		if cfg.Username == "" || cfg.Password.String() == "" {
			return ErrInconsistentSASLCredentials
		}

	case SASLMechanismOauthbearer:
		return cfg.Oauthbearer.Validate()

	case SASLMechanismMSKIAM:
		return cfg.MSKIAM.Validate()

	default:
		return ErrInvalidSASLMechanism
	}

	return nil
}

// kafkaSASLConfig defines how to get a SASL secret: either from a static secret
// provided in config, from a file, or with HTTP through a domain socket.
//
// Exactly one source must be configured; call Validate to enforce this.
//
// Because this struct is generic, it cannot have distinct YAML struct tags
// for each kind of staticSecretConfig. Concrete structs with the same shape (so
// they can be converted to kafkaSASLConfig) are used instead.
type kafkaSASLConfig[T saslSecretConfig] struct {
	Secret            T
	FilePath          string
	HTTPSocketPath    string
	HTTPSocketTimeout time.Duration
}

var errNoSecret = errors.New("no static credentials provided")

var errIncompleteMSKIAMSecret = errors.New("both access key and secret key must be configured for static AWS_MSK_IAM credentials")

// Validate returns errMultipleSources unless exactly one of the static secret,
// file path, or HTTP socket path is configured.
func (cfg kafkaSASLConfig[T]) Validate(errNoSingleSource error) error {
	err := cfg.Secret.Validate()
	if err != nil && !errors.Is(err, errNoSecret) {
		return err
	}
	hasStaticSecret := err == nil
	sourceFound := false
	for _, source := range []bool{
		hasStaticSecret,
		cfg.FilePath != "",
		cfg.HTTPSocketPath != "",
	} {
		if source {
			if sourceFound {
				return errNoSingleSource
			}
			sourceFound = true
		}
	}
	if !sourceFound {
		return errNoSingleSource
	}
	return nil
}

// KafkaAuthOauthbearerConfig holds OAUTHBEARER-specific SASL configuration.
type KafkaAuthOauthbearerConfig struct {
	Secret            KafkaOauthbearerStaticConfig `yaml:",inline"`
	FilePath          string                       `yaml:"sasl_oauthbearer_file_path"`
	HTTPSocketPath    string                       `yaml:"sasl_oauthbearer_http_socket_path"`
	HTTPSocketTimeout time.Duration                `yaml:"sasl_oauthbearer_http_socket_timeout"`
}

func (cfg *KafkaAuthOauthbearerConfig) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
	cfg.Secret.RegisterFlagsWithPrefix(prefix, f)
	f.StringVar(&cfg.FilePath, prefix+"file-path", "", `Path to a file containing an OAuth token to authenticate to Kafka. Mutually exclusive with `+prefix+`http-socket-path. The file is read anew on every reauthentication, so it can be updated with fresh tokens. The file must be in JSON format, adhering to this JSON schema: {"type": "object", "required": ["token"], "properties": {"token": {"type": "string"}, "zid": {"type": "string"}, "extensions": {"type": "object", "additionalProperties": {"type": "string"}}}`)
	f.StringVar(&cfg.HTTPSocketPath, prefix+"http-socket-path", "", `Path to a Unix domain socket to fetch an OAuth token from via HTTP. Mutually exclusive with `+prefix+`file-path. On every authentication or reauthentication, an HTTP GET / request is made to the socket and the response body is read as JSON. The JSON schema is the same as for `+prefix+`file-path.`)
	f.DurationVar(&cfg.HTTPSocketTimeout, prefix+"http-socket-timeout", 10*time.Second, "Timeout for requesting the token from the HTTP socket. Effective when "+prefix+"http-socket-path is set.")
}

func (cfg KafkaAuthOauthbearerConfig) Validate() error {
	return (kafkaSASLConfig[KafkaOauthbearerStaticConfig])(cfg).Validate(ErrSASLOauthbearerBadConfig)
}

// KafkaOauthbearerStaticConfig holds static OAUTHBEARER credentials.
type KafkaOauthbearerStaticConfig struct {
	Token      flagext.Secret            `yaml:"sasl_oauthbearer_token"`
	Zid        string                    `yaml:"sasl_oauthbearer_zid"`
	Extensions flagext.LimitsMap[string] `yaml:"sasl_oauthbearer_extensions"`
}

// Validate returns ErrSecretNotProvided when no token has been set.
func (s KafkaOauthbearerStaticConfig) Validate() error {
	if s.Token.String() == "" {
		return errNoSecret
	}
	return nil
}

func (s *KafkaOauthbearerStaticConfig) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
	f.Var(&s.Token, prefix+"token", "The OAuth token to use to authenticate to Kafka. Consider "+prefix+"file-path instead.")
	f.StringVar(&s.Zid, prefix+"zid", "", "Optional authorization ID to use when authenticating to Kafka using SASL OAUTHBEARER.")
	if !s.Extensions.IsInitialized() {
		s.Extensions = flagext.NewLimitsMap[string](nil)
	}
	f.Var(&s.Extensions, prefix+"extensions", "Optional additional OAuth extensions to include when authenticating to Kafka using SASL OAUTHBEARER as a JSON object.")
}

// KafkaAuthMSKIAMConfig holds AWS_MSK_IAM-specific SASL configuration.
type KafkaAuthMSKIAMConfig struct {
	Secret            KafkaMSKIAMStaticConfig `yaml:",inline"`
	FilePath          string                  `yaml:"sasl_msk_iam_file_path"`
	HTTPSocketPath    string                  `yaml:"sasl_msk_iam_http_socket_path"`
	HTTPSocketTimeout time.Duration           `yaml:"sasl_msk_iam_http_socket_timeout"`
}

func (cfg *KafkaAuthMSKIAMConfig) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
	cfg.Secret.RegisterFlagsWithPrefix(prefix, f)
	f.StringVar(&cfg.FilePath, prefix+"file-path", "", `Path to a file containing AWS credentials to authenticate to Kafka using SASL AWS_MSK_IAM. Mutually exclusive with `+prefix+`http-socket-path. The file is read anew on every reauthentication, so it can be updated with fresh credentials. The file must be in JSON format, adhering to this JSON schema: {"type": "object", "required": ["AccessKey", "SecretKey"], "properties": {"AccessKey": {"type": "string"}, "SecretKey": {"type": "string"}, "SessionToken": {"type": "string"}, "UserAgent": {"type": "string"}}}`)
	f.StringVar(&cfg.HTTPSocketPath, prefix+"http-socket-path", "", `Path to a Unix domain socket to fetch AWS credentials from via HTTP. Mutually exclusive with `+prefix+`file-path. On every authentication or reauthentication, an HTTP GET / request is made to the socket and the response body is read as JSON. The JSON schema is the same as for `+prefix+`file-path.`)
	f.DurationVar(&cfg.HTTPSocketTimeout, prefix+"http-socket-timeout", 10*time.Second, "Timeout for requesting AWS credentials from the HTTP socket. Effective when "+prefix+"http-socket-path is set.")
}

func (cfg KafkaAuthMSKIAMConfig) Validate() error {
	return (kafkaSASLConfig[KafkaMSKIAMStaticConfig])(cfg).Validate(ErrSASLMSKIAMBadConfig)
}

// KafkaMSKIAMStaticConfig holds static AWS_MSK_IAM credentials.
type KafkaMSKIAMStaticConfig struct {
	AccessKey    flagext.Secret `yaml:"sasl_msk_iam_access_key"`
	SecretKey    flagext.Secret `yaml:"sasl_msk_iam_secret_key"`
	SessionToken flagext.Secret `yaml:"sasl_msk_iam_session_token"`
	UserAgent    string         `yaml:"sasl_msk_iam_user_agent"`
}

// Validate returns errNoSecret when no access key and secret key have been set.
func (s KafkaMSKIAMStaticConfig) Validate() error {
	hasAccessKey := s.AccessKey.String() != ""
	hasSecretKey := s.SecretKey.String() != ""

	if !hasAccessKey && !hasSecretKey {
		return errNoSecret
	}
	if !hasAccessKey || !hasSecretKey {
		return errIncompleteMSKIAMSecret
	}

	return nil
}

func (s *KafkaMSKIAMStaticConfig) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
	f.Var(&s.AccessKey, prefix+"access-key", "The AWS access key ID to authenticate to Kafka using SASL AWS_MSK_IAM. Consider "+prefix+"file-path instead.")
	f.Var(&s.SecretKey, prefix+"secret-key", "The AWS secret access key to authenticate to Kafka using SASL AWS_MSK_IAM. Consider "+prefix+"file-path instead.")
	f.Var(&s.SessionToken, prefix+"session-token", "Optional AWS session token to authenticate to Kafka using SASL AWS_MSK_IAM.")
	f.StringVar(&s.UserAgent, prefix+"user-agent", "", "Optional user agent to use when authenticating to Kafka using SASL AWS_MSK_IAM.")
}

type TLSClientConfig struct {
	dskittls.ClientConfig `yaml:",inline"`
	resolved              *tls.Config `yaml:"-"`
}

func (c *TLSClientConfig) GetTLSConfig() (*tls.Config, error) {
	if c.resolved != nil {
		return c.resolved, nil
	}
	var err error
	c.resolved, err = c.ClientConfig.GetTLSConfig()
	return c.resolved, err
}
