wedb_embed 0.1.0

use super::gorilla::{TSSample, gorilla_compress_samples, gorilla_decompress_samples};
use super::meta::{ChunkType, DuplicatePolicy};
use crate::error::{Error, Result};

/// TSChunk 头部元数据（8 字节：4 字节 Flag + 4 字节 Count，对标 Apache Kvrocks TSChunk::MetaData）
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct ChunkHeader {
    pub is_compressed: bool,
    pub count: u32,
}

impl ChunkHeader {
    pub const ENCODED_SIZE: usize = 8;

    #[inline]
    pub fn encode(&self) -> [u8; Self::ENCODED_SIZE] {
        let mut buf = [0u8; Self::ENCODED_SIZE];
        let flag = if self.is_compressed { 1u32 } else { 0u32 };
        buf[0..4].copy_from_slice(&flag.to_be_bytes());
        buf[4..8].copy_from_slice(&self.count.to_be_bytes());
        buf
    }

    #[inline]
    pub fn decode(bytes: &[u8]) -> Option<Self> {
        if bytes.len() < Self::ENCODED_SIZE {
            return None;
        }
        let mut b4 = [0u8; 4];
        b4.copy_from_slice(&bytes[0..4]);
        let flag = u32::from_be_bytes(b4);
        b4.copy_from_slice(&bytes[4..8]);
        let count = u32::from_be_bytes(b4);
        Some(Self {
            is_compressed: (flag & 1) != 0,
            count,
        })
    }
}

/// 样本合并统计结果
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct MergeStats {
    pub inserted: usize,
    pub updated: usize,
    pub skipped: usize,
}

/// TSChunk 操作封装（支持 Uncompressed 及 Gorilla Compressed 切片与合并）
pub struct TSChunk;

impl TSChunk {
    /// 编码为未压缩块（Header + [TSSample (16B)] * count）
    #[inline]
    pub fn encode_uncompressed(samples: &[TSSample]) -> Vec<u8> {
        let header = ChunkHeader {
            is_compressed: false,
            count: samples.len() as u32,
        };
        let mut buf = Vec::with_capacity(ChunkHeader::ENCODED_SIZE + samples.len() * 16);
        buf.extend_from_slice(&header.encode());
        for s in samples {
            buf.extend_from_slice(&s.ts.to_be_bytes());
            buf.extend_from_slice(&s.v.to_be_bytes());
        }
        buf
    }

    /// 编码为 Gorilla 压缩块（Header + Gorilla Bitstream）
    #[inline]
    pub fn encode_compressed(samples: &[TSSample]) -> Vec<u8> {
        let compressed_payload = gorilla_compress_samples(samples);
        let header = ChunkHeader {
            is_compressed: true,
            count: samples.len() as u32,
        };
        let mut buf = Vec::with_capacity(ChunkHeader::ENCODED_SIZE + compressed_payload.len());
        buf.extend_from_slice(&header.encode());
        buf.extend_from_slice(&compressed_payload);
        buf
    }

    /// 根据 ChunkType 自动编码
    #[inline]
    pub fn encode_with_type(samples: &[TSSample], chunk_type: ChunkType) -> Vec<u8> {
        match chunk_type {
            ChunkType::Compressed => Self::encode_compressed(samples),
            ChunkType::Uncompressed => Self::encode_uncompressed(samples),
        }
    }

    /// 解码 Chunk 字节流为采样点向量
    pub fn decode_samples(chunk_data: &[u8]) -> Result<Vec<TSSample>> {
        if chunk_data.is_empty() {
            return Ok(Vec::new());
        }
        if chunk_data.len() < ChunkHeader::ENCODED_SIZE {
            if chunk_data.len() == 8 {
                let mut b = [0u8; 8];
                b.copy_from_slice(chunk_data);
                return Ok(vec![TSSample::new(0, f64::from_be_bytes(b))]);
            }
            return Err(Error::invalid_data(
                "ERR TSDB: TSChunk payload data too short",
            ));
        }

        let header = ChunkHeader::decode(chunk_data)
            .ok_or_else(|| Error::invalid_data("ERR TSDB: invalid TSChunk header"))?;

        if header.count == 0 {
            return Ok(Vec::new());
        }

        let payload = &chunk_data[ChunkHeader::ENCODED_SIZE..];

        if header.is_compressed {
            gorilla_decompress_samples(payload, header.count as usize)
        } else {
            let count = header.count as usize;
            if payload.len() < count * 16 {
                return Err(Error::invalid_data(
                    "ERR TSDB: uncompressed TSChunk payload too short",
                ));
            }
            let mut samples = Vec::with_capacity(count);
            for i in 0..count {
                let offset = i * 16;
                let mut b8 = [0u8; 8];
                b8.copy_from_slice(&payload[offset..offset + 8]);
                let ts = u64::from_be_bytes(b8);
                b8.copy_from_slice(&payload[offset + 8..offset + 16]);
                let v = f64::from_be_bytes(b8);
                samples.push(TSSample::new(ts, v));
            }
            Ok(samples)
        }
    }

    /// 提取 Chunk 首个时间戳（零全量解压轻量级提取）
    pub fn get_first_timestamp(chunk_data: &[u8]) -> Option<u64> {
        if chunk_data.len() < ChunkHeader::ENCODED_SIZE {
            return None;
        }
        let header = ChunkHeader::decode(chunk_data)?;
        if header.count == 0 {
            return None;
        }
        let payload = &chunk_data[ChunkHeader::ENCODED_SIZE..];
        if payload.len() >= 8 {
            let mut b8 = [0u8; 8];
            b8.copy_from_slice(&payload[0..8]);
            Some(u64::from_be_bytes(b8))
        } else {
            None
        }
    }

    /// 提取 Chunk 末尾时间戳
    pub fn get_last_timestamp(chunk_data: &[u8]) -> Option<u64> {
        let samples = Self::decode_samples(chunk_data).ok()?;
        samples.last().map(|s| s.ts)
    }

    /// 提取 Chunk 采样点总数
    #[inline]
    pub fn get_count(chunk_data: &[u8]) -> u32 {
        if chunk_data.len() < ChunkHeader::ENCODED_SIZE {
            return 0;
        }
        ChunkHeader::decode(chunk_data)
            .map(|h| h.count)
            .unwrap_or(0)
    }

    /// 合并新样本点并应用 DuplicatePolicy（若 Block 策略冲突则返回错误）
    pub fn merge_samples(
        existing: &mut Vec<TSSample>,
        new_samples: &[TSSample],
        policy: DuplicatePolicy,
    ) -> Result<MergeStats> {
        let mut stats = MergeStats::default();
        if new_samples.is_empty() {
            return Ok(stats);
        }

        // 单样本优化快路径
        if new_samples.len() == 1 {
            let new_s = new_samples[0];
            match existing.binary_search_by_key(&new_s.ts, |s| s.ts) {
                Ok(idx) => {
                    let old_v = existing[idx].v;
                    match policy {
                        DuplicatePolicy::Block => {
                            return Err(Error::invalid_data(
                                "ERR TSDB: Error at upsert, update is not supported when DUPLICATE_POLICY is set to BLOCK mode",
                            ));
                        }
                        DuplicatePolicy::First => {
                            stats.skipped += 1;
                        }
                        DuplicatePolicy::Last => {
                            if (existing[idx].v - new_s.v).abs() < f64::EPSILON {
                                stats.skipped += 1;
                            } else {
                                existing[idx].v = new_s.v;
                                stats.updated += 1;
                            }
                        }
                        DuplicatePolicy::Min => {
                            if new_s.v < old_v {
                                existing[idx].v = new_s.v;
                                stats.updated += 1;
                            } else {
                                stats.skipped += 1;
                            }
                        }
                        DuplicatePolicy::Max => {
                            if new_s.v > old_v {
                                existing[idx].v = new_s.v;
                                stats.updated += 1;
                            } else {
                                stats.skipped += 1;
                            }
                        }
                        DuplicatePolicy::Sum => {
                            if new_s.v == 0.0 {
                                stats.skipped += 1;
                            } else {
                                existing[idx].v = old_v + new_s.v;
                                stats.updated += 1;
                            }
                        }
                    }
                }
                Err(idx) => {
                    existing.insert(idx, new_s);
                    stats.inserted += 1;
                }
            }
            return Ok(stats);
        }

        // 多样本双指针归并
        let mut sorted_new = new_samples.to_vec();
        sorted_new.sort_by_key(|s| s.ts);

        let mut merged = Vec::with_capacity(existing.len() + sorted_new.len());
        let mut i = 0;
        let mut j = 0;

        while i < existing.len() && j < sorted_new.len() {
            let e = existing[i];
            let n = sorted_new[j];
            if e.ts < n.ts {
                merged.push(e);
                i += 1;
            } else if e.ts > n.ts {
                merged.push(n);
                stats.inserted += 1;
                j += 1;
            } else {
                let old_v = e.v;
                let final_v = match policy {
                    DuplicatePolicy::Block => {
                        return Err(Error::invalid_data(
                            "ERR TSDB: Error at upsert, update is not supported when DUPLICATE_POLICY is set to BLOCK mode",
                        ));
                    }
                    DuplicatePolicy::First => {
                        stats.skipped += 1;
                        old_v
                    }
                    DuplicatePolicy::Last => {
                        if (old_v - n.v).abs() < f64::EPSILON {
                            stats.skipped += 1;
                        } else {
                            stats.updated += 1;
                        }
                        n.v
                    }
                    DuplicatePolicy::Min => {
                        if n.v < old_v {
                            stats.updated += 1;
                            n.v
                        } else {
                            stats.skipped += 1;
                            old_v
                        }
                    }
                    DuplicatePolicy::Max => {
                        if n.v > old_v {
                            stats.updated += 1;
                            n.v
                        } else {
                            stats.skipped += 1;
                            old_v
                        }
                    }
                    DuplicatePolicy::Sum => {
                        if n.v == 0.0 {
                            stats.skipped += 1;
                        } else {
                            stats.updated += 1;
                        }
                        old_v + n.v
                    }
                };
                merged.push(TSSample::new(e.ts, final_v));
                i += 1;
                j += 1;
            }
        }

        while i < existing.len() {
            merged.push(existing[i]);
            i += 1;
        }

        while j < sorted_new.len() {
            merged.push(sorted_new[j]);
            stats.inserted += 1;
            j += 1;
        }

        *existing = merged;
        Ok(stats)
    }

    /// Upsert 并按 chunk_size 拆分（对标 Apache Kvrocks UpsertSampleAndSplit）
    pub fn upsert_and_split(
        existing_data: &[u8],
        new_samples: &[TSSample],
        policy: DuplicatePolicy,
        preferred_chunk_size: usize,
        chunk_type: ChunkType,
    ) -> Result<Vec<Vec<u8>>> {
        let mut samples = if existing_data.is_empty() {
            Vec::new()
        } else {
            Self::decode_samples(existing_data)?
        };

        Self::merge_samples(&mut samples, new_samples, policy)?;

        if samples.is_empty() {
            return Ok(Vec::new());
        }

        let chunk_size = preferred_chunk_size.max(1);
        let mut chunks = Vec::new();

        for chunk_slice in samples.chunks(chunk_size) {
            chunks.push(Self::encode_with_type(chunk_slice, chunk_type));
        }

        Ok(chunks)
    }

    /// 删除指定时间范围内的样本点 [from_ts, to_ts]
    pub fn remove_samples_between(
        chunk_data: &[u8],
        from_ts: u64,
        to_ts: u64,
        chunk_type: ChunkType,
    ) -> Result<(Vec<u8>, usize)> {
        if from_ts > to_ts || chunk_data.is_empty() {
            return Ok((chunk_data.to_vec(), 0));
        }
        let samples = Self::decode_samples(chunk_data)?;
        let orig_len = samples.len();
        let filtered: Vec<TSSample> = samples
            .into_iter()
            .filter(|s| s.ts < from_ts || s.ts > to_ts)
            .collect();
        let deleted = orig_len - filtered.len();
        if deleted == 0 {
            return Ok((chunk_data.to_vec(), 0));
        }
        let encoded = Self::encode_with_type(&filtered, chunk_type);
        Ok((encoded, deleted))
    }

    /// 更新指定时间戳的样本点数值
    pub fn update_sample_value(
        chunk_data: &[u8],
        ts: u64,
        value: f64,
        is_add_on: bool,
        chunk_type: ChunkType,
    ) -> Result<Option<Vec<u8>>> {
        let mut samples = Self::decode_samples(chunk_data)?;
        if let Ok(idx) = samples.binary_search_by_key(&ts, |s| s.ts) {
            if is_add_on {
                samples[idx].v += value;
            } else {
                samples[idx].v = value;
            }
            Ok(Some(Self::encode_with_type(&samples, chunk_type)))
        } else {
            Ok(None)
        }
    }
}