wedb_embed 0.1.0

use std::collections::BTreeSet;
use std::fmt::{self, Display};
use std::str::FromStr;

use rapidhash::RapidHashSet;
use serde::{Deserialize, Serialize};

pub use super::meta::{ChunkType, DuplicatePolicy, TimeSeriesMeta};

/// 聚合计算类型（对标 Apache Kvrocks TSAggregatorType + RedisTimeSeries）
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[repr(u8)]
pub enum AggregationType {
    #[default]
    None = 0,
    Sum = 1,
    Min = 2,
    Max = 3,
    Count = 4,
    First = 5,
    Last = 6,
    Avg = 7,
    Range = 8,
    StdP = 9,
    StdS = 10,
    VarP = 11,
    VarS = 12,
    Twa = 13,
}

impl FromStr for AggregationType {
    type Err = ();

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_ascii_uppercase().as_str() {
            "NONE" => Ok(Self::None),
            "SUM" => Ok(Self::Sum),
            "MIN" => Ok(Self::Min),
            "MAX" => Ok(Self::Max),
            "COUNT" => Ok(Self::Count),
            "FIRST" => Ok(Self::First),
            "LAST" => Ok(Self::Last),
            "AVG" => Ok(Self::Avg),
            "RANGE" => Ok(Self::Range),
            "STD.P" | "STDP" => Ok(Self::StdP),
            "STD.S" | "STDS" => Ok(Self::StdS),
            "VAR.P" | "VARP" => Ok(Self::VarP),
            "VAR.S" | "VARS" => Ok(Self::VarS),
            "TWA" => Ok(Self::Twa),
            _ => Err(()),
        }
    }
}

impl Display for AggregationType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

impl AggregationType {
    pub fn parse(s: &str) -> Option<Self> {
        s.parse().ok()
    }

    #[inline]
    pub const fn as_str(&self) -> &'static str {
        match self {
            Self::None => "none",
            Self::Sum => "sum",
            Self::Min => "min",
            Self::Max => "max",
            Self::Count => "count",
            Self::First => "first",
            Self::Last => "last",
            Self::Avg => "avg",
            Self::Range => "range",
            Self::StdP => "std.p",
            Self::StdS => "std.s",
            Self::VarP => "var.p",
            Self::VarS => "var.s",
            Self::Twa => "twa",
        }
    }

    /// 是否支持增量累加计算（对标 Kvrocks IsIncrementalAggregatorType）
    #[inline]
    pub const fn is_incremental(&self) -> bool {
        matches!(self, Self::Sum | Self::Min | Self::Max | Self::Count)
    }
}

/// 降采样桶时间戳类型（对标 Apache Kvrocks BucketTimestampType）
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[repr(u8)]
pub enum BucketTimestampType {
    #[default]
    Start = 0,
    End = 1,
    Mid = 2,
}

impl FromStr for BucketTimestampType {
    type Err = ();

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_ascii_uppercase().as_str() {
            "START" | "-" => Ok(Self::Start),
            "END" | "+" => Ok(Self::End),
            "MID" | "~" => Ok(Self::Mid),
            _ => Err(()),
        }
    }
}

impl Display for BucketTimestampType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

impl BucketTimestampType {
    pub fn parse(s: &str) -> Option<Self> {
        s.parse().ok()
    }

    #[inline]
    pub const fn as_str(&self) -> &'static str {
        match self {
            Self::Start => "start",
            Self::End => "end",
            Self::Mid => "mid",
        }
    }
}

/// 时序聚合计算器（对标 Apache Kvrocks TSAggregator）
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
pub struct Aggregator {
    pub agg_type: AggregationType,
    pub bucket_duration: u64,
    pub alignment: u64,
}

impl Aggregator {
    #[inline]
    pub fn new(agg_type: AggregationType, bucket_duration: u64, alignment: u64) -> Self {
        Self {
            agg_type,
            bucket_duration: bucket_duration.max(1),
            alignment,
        }
    }

    /// 计算对齐时间桶左边界（对标 Kvrocks CalculateAlignedBucketLeft）
    #[inline]
    pub fn calculate_aligned_bucket_left(&self, ts: u64) -> u64 {
        if self.bucket_duration == 0 {
            return ts;
        }
        if ts >= self.alignment {
            let diff = ts - self.alignment;
            let k = diff / self.bucket_duration;
            self.alignment + k * self.bucket_duration
        } else {
            let diff = self.alignment - ts;
            let m0 = diff / self.bucket_duration
                + if diff.is_multiple_of(self.bucket_duration) {
                    0
                } else {
                    1
                };
            if m0 <= self.alignment / self.bucket_duration {
                self.alignment - m0 * self.bucket_duration
            } else {
                0
            }
        }
    }

    /// 计算对齐时间桶右边界（对标 Kvrocks CalculateAlignedBucketRight）
    #[inline]
    pub fn calculate_aligned_bucket_right(&self, ts: u64) -> u64 {
        if self.bucket_duration == 0 {
            return ts;
        }
        if ts < self.alignment {
            let diff = self.alignment - ts;
            let k = diff / self.bucket_duration;
            self.alignment.saturating_sub(k * self.bucket_duration)
        } else {
            let diff = ts - self.alignment;
            let m0 = diff / self.bucket_duration + 1;
            self.alignment
                .saturating_add(m0.saturating_mul(self.bucket_duration))
        }
    }

    /// 将采样点零拷贝切分到对齐时间桶切片列表
    pub fn split_samples_to_buckets<'a>(
        &self,
        samples: &'a [(u64, f64)],
    ) -> Vec<(u64, &'a [(u64, f64)])> {
        if samples.is_empty() {
            return Vec::new();
        }

        let mut buckets = Vec::new();
        let mut start_idx = 0;
        let mut current_left = self.calculate_aligned_bucket_left(samples[0].0);

        for (i, &(ts, _)) in samples.iter().enumerate() {
            let bucket_left = self.calculate_aligned_bucket_left(ts);
            if bucket_left != current_left {
                buckets.push((current_left, &samples[start_idx..i]));
                start_idx = i;
                current_left = bucket_left;
            }
        }

        if start_idx < samples.len() {
            buckets.push((current_left, &samples[start_idx..]));
        }

        buckets
    }

    /// 聚合计算单个时间桶内采样点值（采用 Welford 算法单次迭代计算在线方差）
    pub fn aggregate_samples(&self, samples: &[(u64, f64)]) -> f64 {
        if samples.is_empty() {
            return match self.agg_type {
                AggregationType::Sum | AggregationType::Count => 0.0,
                _ => f64::NAN,
            };
        }

        let count = samples.len();
        let first = samples[0].1;
        let last = samples[count - 1].1;

        let mut sum = 0.0;
        let mut min = f64::INFINITY;
        let mut max = f64::NEG_INFINITY;
        let mut mean = 0.0;
        let mut m2 = 0.0;

        for (i, &(_, val)) in samples.iter().enumerate() {
            sum += val;
            if val < min {
                min = val;
            }
            if val > max {
                max = val;
            }
            let delta = val - mean;
            mean += delta / ((i + 1) as f64);
            let delta2 = val - mean;
            m2 += delta * delta2;
        }

        match self.agg_type {
            AggregationType::None => last,
            AggregationType::Avg => mean,
            AggregationType::Sum => sum,
            AggregationType::Min => min,
            AggregationType::Max => max,
            AggregationType::Count => count as f64,
            AggregationType::First => first,
            AggregationType::Last => last,
            AggregationType::Range => max - min,
            AggregationType::StdP => (m2 / (count as f64)).sqrt(),
            AggregationType::StdS => {
                if count <= 1 {
                    0.0
                } else {
                    (m2 / ((count - 1) as f64)).sqrt()
                }
            }
            AggregationType::VarP => m2 / (count as f64),
            AggregationType::VarS => {
                if count <= 1 {
                    0.0
                } else {
                    m2 / ((count - 1) as f64)
                }
            }
            AggregationType::Twa => {
                if count == 1 {
                    first
                } else {
                    let total_time_span = samples[count - 1].0 - samples[0].0;
                    if total_time_span == 0 {
                        mean
                    } else {
                        let mut weighted_sum = 0.0;
                        for i in 0..count - 1 {
                            let dt = (samples[i + 1].0 - samples[i].0) as f64;
                            let avg_v = (samples[i].1 + samples[i + 1].1) / 2.0;
                            weighted_sum += avg_v * dt;
                        }
                        weighted_sum / (total_time_span as f64)
                    }
                }
            }
        }
    }

    /// 切分并聚合采样序列（默认 Start 对齐，不填充空桶）
    #[inline]
    pub fn split_and_aggregate(
        &self,
        samples: &[(u64, f64)],
        limit: Option<usize>,
    ) -> Vec<(u64, f64)> {
        self.split_and_aggregate_opt(samples, limit, false, BucketTimestampType::Start)
    }

    /// 切分并聚合采样序列（支持 limit、empty 桶填充以及时间戳对齐选项）
    pub fn split_and_aggregate_opt(
        &self,
        samples: &[(u64, f64)],
        limit: Option<usize>,
        is_return_empty: bool,
        ts_type: BucketTimestampType,
    ) -> Vec<(u64, f64)> {
        if samples.is_empty() {
            return Vec::new();
        }

        let get_bucket_ts = |left: u64| -> u64 {
            match ts_type {
                BucketTimestampType::Start => left,
                BucketTimestampType::End => left + self.bucket_duration,
                BucketTimestampType::Mid => left + self.bucket_duration / 2,
            }
        };

        let buckets = self.split_samples_to_buckets(samples);
        if buckets.is_empty() {
            return Vec::new();
        }

        if !is_return_empty {
            let cap = if let Some(l) = limit {
                l.min(buckets.len())
            } else {
                buckets.len()
            };
            let mut results = Vec::with_capacity(cap);
            for (bucket_left, bucket_samples) in buckets {
                let agg_val = self.aggregate_samples(bucket_samples);
                results.push((get_bucket_ts(bucket_left), agg_val));
                if let Some(l) = limit
                    && results.len() >= l
                {
                    break;
                }
            }
            return results;
        }

        // 处理 EMPTY 模式：双指针顺序填充空桶，避免哈希开销
        let mut results = Vec::new();
        let last_left = buckets[buckets.len() - 1].0;
        let mut curr_left = buckets[0].0;
        let mut bucket_idx = 0;
        let mut last_known_val = f64::NAN;

        while curr_left <= last_left {
            let out_ts = get_bucket_ts(curr_left);
            if bucket_idx < buckets.len() && buckets[bucket_idx].0 == curr_left {
                let b_samples = buckets[bucket_idx].1;
                let agg_val = self.aggregate_samples(b_samples);
                if let Some(last_s) = b_samples.last() {
                    last_known_val = last_s.1;
                }
                results.push((out_ts, agg_val));
                bucket_idx += 1;
            } else {
                let empty_val = match self.agg_type {
                    AggregationType::Sum | AggregationType::Count => 0.0,
                    AggregationType::Last => last_known_val,
                    _ => f64::NAN,
                };
                results.push((out_ts, empty_val));
            }

            if let Some(l) = limit
                && results.len() >= l
            {
                break;
            }

            curr_left = self.calculate_aligned_bucket_right(curr_left);
        }

        results
    }
}

/// 创建 TS 选项（对标 Apache Kvrocks TSCreateOption）
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct TSCreateOption {
    pub retention_time: u64,
    pub chunk_size: u64,
    pub chunk_type: ChunkType,
    pub duplicate_policy: DuplicatePolicy,
    pub source_key: String,
    pub labels: Vec<(String, String)>,
}

/// TS.RANGE 检索选项（对标 Apache Kvrocks TSRangeOption）
#[derive(Debug, Clone, PartialEq, Default)]
pub struct TSRangeOption {
    pub start_ts: u64,
    pub end_ts: u64,
    pub count_limit: Option<usize>,
    pub filter_by_ts: BTreeSet<u64>,
    pub filter_by_value: Option<(f64, f64)>,
    pub aggregator: Option<Aggregator>,
    pub is_return_latest: bool,
    pub is_return_empty: bool,
    pub bucket_timestamp_type: BucketTimestampType,
}

/// 多序列聚合组 Reducer 类型（对标 Apache Kvrocks GroupReducerType）
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[repr(u8)]
pub enum GroupReducerType {
    #[default]
    None = 0,
    Avg = 1,
    Sum = 2,
    Min = 3,
    Max = 4,
    Range = 5,
    Count = 6,
    StdP = 7,
    StdS = 8,
    VarP = 9,
    VarS = 10,
    Twa = 11,
}

impl FromStr for GroupReducerType {
    type Err = ();

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_ascii_uppercase().as_str() {
            "NONE" => Ok(Self::None),
            "AVG" => Ok(Self::Avg),
            "SUM" => Ok(Self::Sum),
            "MIN" => Ok(Self::Min),
            "MAX" => Ok(Self::Max),
            "RANGE" => Ok(Self::Range),
            "COUNT" => Ok(Self::Count),
            "STD.P" | "STDP" => Ok(Self::StdP),
            "STD.S" | "STDS" => Ok(Self::StdS),
            "VAR.P" | "VARP" => Ok(Self::VarP),
            "VAR.S" | "VARS" => Ok(Self::VarS),
            "TWA" => Ok(Self::Twa),
            _ => Err(()),
        }
    }
}

impl Display for GroupReducerType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

impl GroupReducerType {
    pub fn parse(s: &str) -> Option<Self> {
        s.parse().ok()
    }

    #[inline]
    pub const fn as_str(&self) -> &'static str {
        match self {
            Self::None => "none",
            Self::Avg => "avg",
            Self::Sum => "sum",
            Self::Min => "min",
            Self::Max => "max",
            Self::Range => "range",
            Self::Count => "count",
            Self::StdP => "std.p",
            Self::StdS => "std.s",
            Self::VarP => "var.p",
            Self::VarS => "var.s",
            Self::Twa => "twa",
        }
    }
}

/// TS.MGET 选项
#[derive(Debug, Clone, Default)]
pub struct TSMGetOption {
    pub with_labels: bool,
    pub selected_labels: RapidHashSet<String>,
    pub filters: Vec<String>,
}

/// TS.MGET 结果
#[derive(Debug, Clone, PartialEq)]
pub struct TSMGetResult {
    pub name: String,
    pub labels: Vec<(String, String)>,
    pub sample: Option<(u64, f64)>,
}

/// TS.MRANGE 选项
#[derive(Debug, Clone, Default)]
pub struct TSMRangeOption {
    pub mget: TSMGetOption,
    pub range: TSRangeOption,
    pub reducer: GroupReducerType,
    pub group_by_label: Option<String>,
}

/// TS.MRANGE 结果
#[derive(Debug, Clone, PartialEq)]
pub struct TSMRangeResult {
    pub name: String,
    pub labels: Vec<(String, String)>,
    pub samples: Vec<(u64, f64)>,
    pub source_keys: Vec<String>,
}

/// 下游降采样元数据（对标 Apache Kvrocks TSDownStreamMeta）
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
pub struct TSDownStreamMeta {
    pub aggregator: Aggregator,
    pub latest_bucket_idx: u64,
    pub u64_auxs: Vec<u64>,
    pub f64_auxs: Vec<f64>,
}

impl TSDownStreamMeta {
    #[inline]
    pub fn new(aggregator: Aggregator) -> Self {
        let mut meta = Self {
            aggregator,
            latest_bucket_idx: 0,
            u64_auxs: Vec::new(),
            f64_auxs: Vec::new(),
        };
        meta.reset_auxs();
        meta
    }

    pub fn reset_auxs(&mut self) {
        match self.aggregator.agg_type {
            AggregationType::Sum | AggregationType::Count => {
                self.f64_auxs = vec![0.0];
                self.u64_auxs.clear();
            }
            AggregationType::Min | AggregationType::Max => {
                self.f64_auxs = vec![f64::NAN];
                self.u64_auxs.clear();
            }
            AggregationType::First => {
                self.u64_auxs = vec![u64::MAX];
                self.f64_auxs = vec![f64::NAN];
            }
            AggregationType::Last => {
                self.u64_auxs = vec![0];
                self.f64_auxs = vec![f64::NAN];
            }
            AggregationType::Avg => {
                self.u64_auxs = vec![0];
                self.f64_auxs = vec![0.0];
            }
            AggregationType::StdP
            | AggregationType::StdS
            | AggregationType::VarP
            | AggregationType::VarS => {
                self.u64_auxs = vec![0];
                self.f64_auxs = vec![0.0, 0.0];
            }
            AggregationType::Range => {
                self.f64_auxs = vec![f64::NAN, f64::NAN];
                self.u64_auxs.clear();
            }
            _ => {
                self.u64_auxs.clear();
                self.f64_auxs.clear();
            }
        }
    }

    pub fn encode(&self) -> Vec<u8> {
        let cap = 1 + 8 + 8 + 8 + 4 + self.u64_auxs.len() * 8 + 4 + self.f64_auxs.len() * 8;
        let mut buf = Vec::with_capacity(cap);
        buf.push(self.aggregator.agg_type as u8);
        buf.extend_from_slice(&self.aggregator.bucket_duration.to_be_bytes());
        buf.extend_from_slice(&self.aggregator.alignment.to_be_bytes());
        buf.extend_from_slice(&self.latest_bucket_idx.to_be_bytes());
        buf.extend_from_slice(&(self.u64_auxs.len() as u32).to_be_bytes());
        for &u in &self.u64_auxs {
            buf.extend_from_slice(&u.to_be_bytes());
        }
        buf.extend_from_slice(&(self.f64_auxs.len() as u32).to_be_bytes());
        for &f in &self.f64_auxs {
            buf.extend_from_slice(&f.to_be_bytes());
        }
        buf
    }

    pub fn decode(bytes: &[u8]) -> Option<Self> {
        if bytes.len() < 1 + 8 + 8 + 8 + 4 {
            return None;
        }
        let agg_type = match bytes[0] {
            1 => AggregationType::Sum,
            2 => AggregationType::Min,
            3 => AggregationType::Max,
            4 => AggregationType::Count,
            5 => AggregationType::First,
            6 => AggregationType::Last,
            7 => AggregationType::Avg,
            8 => AggregationType::Range,
            9 => AggregationType::StdP,
            10 => AggregationType::StdS,
            11 => AggregationType::VarP,
            12 => AggregationType::VarS,
            13 => AggregationType::Twa,
            _ => AggregationType::None,
        };
        let mut offset = 1;
        let mut b8 = [0u8; 8];
        b8.copy_from_slice(&bytes[offset..offset + 8]);
        let bucket_duration = u64::from_be_bytes(b8);
        offset += 8;

        b8.copy_from_slice(&bytes[offset..offset + 8]);
        let alignment = u64::from_be_bytes(b8);
        offset += 8;

        b8.copy_from_slice(&bytes[offset..offset + 8]);
        let latest_bucket_idx = u64::from_be_bytes(b8);
        offset += 8;

        let mut b4 = [0u8; 4];
        b4.copy_from_slice(&bytes[offset..offset + 4]);
        let u64_len = u32::from_be_bytes(b4) as usize;
        offset += 4;

        let mut u64_auxs = Vec::with_capacity(u64_len);
        for _ in 0..u64_len {
            if offset + 8 > bytes.len() {
                break;
            }
            b8.copy_from_slice(&bytes[offset..offset + 8]);
            u64_auxs.push(u64::from_be_bytes(b8));
            offset += 8;
        }

        let mut f64_auxs = Vec::new();
        if offset + 4 <= bytes.len() {
            b4.copy_from_slice(&bytes[offset..offset + 4]);
            let f64_len = u32::from_be_bytes(b4) as usize;
            offset += 4;

            f64_auxs.reserve(f64_len);
            for _ in 0..f64_len {
                if offset + 8 > bytes.len() {
                    break;
                }
                b8.copy_from_slice(&bytes[offset..offset + 8]);
                f64_auxs.push(f64::from_be_bytes(b8));
                offset += 8;
            }
        }

        Some(Self {
            aggregator: Aggregator::new(agg_type, bucket_duration, alignment),
            latest_bucket_idx,
            u64_auxs,
            f64_auxs,
        })
    }
}

/// TS.INFO 响应结果（对标 Apache Kvrocks TSInfoResult）
#[derive(Debug, Clone, PartialEq)]
pub struct TSInfoResult {
    pub total_samples: u64,
    pub memory_usage: u64,
    pub first_timestamp: u64,
    pub last_timestamp: u64,
    pub retention_time: u64,
    pub chunk_size: u64,
    pub chunk_type: ChunkType,
    pub duplicate_policy: DuplicatePolicy,
    pub labels: Vec<(String, String)>,
    pub source_key: String,
    pub downstream_rules: Vec<(String, Aggregator)>,
}