// Copyright 2023 Greptime Team
//
// Licensed 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.

use std::num::NonZeroUsize;
use std::time::Duration;

use serde::{Deserialize, Serialize};
use servers::prom_remote_write::validation::PromValidationMode;

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct PromStoreOptions {
    pub enable: bool,
    pub with_metric_engine: bool,
    /// Validation mode while decoding Prometheus remote write requests.
    #[serde(default)]
    pub prom_validation_mode: PromValidationMode,
    /// Enables experimental Prometheus remote write v2 native histogram ingestion.
    #[serde(default)]
    pub experimental_enable_prometheus_native_histogram: bool,
    #[serde(default, with = "humantime_serde")]
    pub pending_rows_flush_interval: Duration,
    #[serde(default = "default_max_batch_rows")]
    pub max_batch_rows: usize,
    #[serde(default = "default_max_concurrent_flushes")]
    pub max_concurrent_flushes: usize,
    #[serde(default = "default_worker_channel_capacity")]
    pub worker_channel_capacity: usize,
    #[serde(default = "default_max_inflight_requests")]
    pub max_inflight_requests: usize,
    /// Maximum number of logical-table flow notifications waiting in the shared queue.
    #[serde(default = "default_flow_notification_queue_capacity")]
    pub flow_notification_queue_capacity: NonZeroUsize,
}

fn default_max_batch_rows() -> usize {
    100_000
}

fn default_max_concurrent_flushes() -> usize {
    256
}

fn default_worker_channel_capacity() -> usize {
    65526
}

fn default_max_inflight_requests() -> usize {
    3000
}

fn default_flow_notification_queue_capacity() -> NonZeroUsize {
    NonZeroUsize::new(1024).unwrap_or(NonZeroUsize::MIN)
}

impl PromStoreOptions {
    /// Returns whether the pending rows batcher can be enabled with these
    /// options. Mirrors the enablement conditions of
    /// `PendingRowsBatcher::try_new` in the servers crate, which returns
    /// `None` when any of these knobs is zero.
    pub fn pending_rows_batching_enabled(&self) -> bool {
        self.enable
            && self.with_metric_engine
            && !self.pending_rows_flush_interval.is_zero()
            && self.max_batch_rows > 0
            && self.max_concurrent_flushes > 0
            && self.worker_channel_capacity > 0
            && self.max_inflight_requests > 0
    }
}

impl Default for PromStoreOptions {
    fn default() -> Self {
        Self {
            enable: true,
            with_metric_engine: true,
            prom_validation_mode: PromValidationMode::Strict,
            experimental_enable_prometheus_native_histogram: false,
            pending_rows_flush_interval: Duration::ZERO,
            max_batch_rows: default_max_batch_rows(),
            max_concurrent_flushes: default_max_concurrent_flushes(),
            worker_channel_capacity: default_worker_channel_capacity(),
            max_inflight_requests: default_max_inflight_requests(),
            flow_notification_queue_capacity: default_flow_notification_queue_capacity(),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use super::{PromStoreOptions, PromValidationMode};
    use crate::service_config::prom_store::{
        default_flow_notification_queue_capacity, default_max_batch_rows,
        default_max_concurrent_flushes, default_max_inflight_requests,
        default_worker_channel_capacity,
    };

    #[test]
    fn test_prom_store_options() {
        let default = PromStoreOptions::default();
        assert!(default.enable);
        assert!(default.with_metric_engine);
        assert_eq!(default.prom_validation_mode, PromValidationMode::Strict);
        assert!(!default.experimental_enable_prometheus_native_histogram);
        assert_eq!(default.pending_rows_flush_interval, Duration::ZERO);
        assert_eq!(default.max_batch_rows, default_max_batch_rows());
        assert_eq!(
            default.max_concurrent_flushes,
            default_max_concurrent_flushes()
        );
        assert_eq!(
            default.worker_channel_capacity,
            default_worker_channel_capacity()
        );
        assert_eq!(
            default.max_inflight_requests,
            default_max_inflight_requests()
        );
        assert_eq!(
            default.flow_notification_queue_capacity,
            default_flow_notification_queue_capacity()
        );
    }
}
