// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0.

use crate::{Error, Result, coding};

#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub(crate) struct BlockHandle {
    pub offset: u64,
    pub size: u64,
}

impl BlockHandle {
    pub fn encode(self, dst: &mut Vec<u8>) {
        coding::put_varint64(dst, self.offset);
        coding::put_varint64(dst, self.size);
    }

    pub fn decode(input: &[u8], pos: &mut usize) -> Result<Self> {
        Ok(Self {
            offset: coding::read_varint64(input, pos)?,
            size: coding::read_varint64(input, pos)?,
        })
    }
}

pub(crate) struct BlockBuilder {
    data: Vec<u8>,
    restarts: Vec<u32>,
    last_key: Vec<u8>,
    entries_since_restart: usize,
    restart_interval: usize,
    value_delta_encoding: bool,
}

impl BlockBuilder {
    pub fn new(restart_interval: usize, value_delta_encoding: bool) -> Self {
        Self {
            data: Vec::new(),
            restarts: vec![0],
            last_key: Vec::new(),
            entries_since_restart: 0,
            restart_interval: restart_interval.max(1),
            value_delta_encoding,
        }
    }

    pub fn add(&mut self, key: &[u8], value: &[u8], delta_value: Option<&[u8]>) -> Result<()> {
        let shared = if self.entries_since_restart >= self.restart_interval {
            self.restarts.push(self.data.len() as u32);
            self.entries_since_restart = 0;
            0
        } else {
            coding::common_prefix_len(&self.last_key, key)
        };
        let non_shared = key.len() - shared;
        coding::put_varint32(&mut self.data, shared as u32);
        coding::put_varint32(&mut self.data, non_shared as u32);
        if !self.value_delta_encoding {
            coding::put_varint32(&mut self.data, value.len() as u32);
        }
        self.data.extend_from_slice(&key[shared..]);
        if self.value_delta_encoding && shared != 0 {
            self.data
                .extend_from_slice(delta_value.ok_or_else(|| {
                    Error::Corruption("missing delta-encoded index value".into())
                })?);
        } else {
            self.data.extend_from_slice(value);
        }
        self.last_key.clear();
        self.last_key.extend_from_slice(key);
        self.entries_since_restart += 1;
        Ok(())
    }

    pub fn estimated_size(&self) -> usize {
        self.data.len() + (self.restarts.len() + 1) * 4
    }

    pub fn finish(mut self) -> Vec<u8> {
        let restart_count = self.restarts.len();
        for restart in self.restarts {
            self.data.extend_from_slice(&restart.to_le_bytes());
        }
        self.data
            .extend_from_slice(&(restart_count as u32).to_le_bytes());
        self.data
    }
}

pub(crate) fn parse_entries(input: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
    let entries_end = entries_end(input)?;
    let mut result = Vec::new();
    let mut pos = 0;
    let mut last_key = Vec::new();
    while pos < entries_end {
        let shared = coding::read_varint32(input, &mut pos)? as usize;
        let non_shared = coding::read_varint32(input, &mut pos)? as usize;
        let value_len = coding::read_varint32(input, &mut pos)? as usize;
        if shared > last_key.len() {
            return Err(Error::Corruption(
                "block shared prefix exceeds prior key".into(),
            ));
        }
        let key_end = pos
            .checked_add(non_shared)
            .ok_or_else(|| Error::Corruption("block key length overflow".into()))?;
        let value_end = key_end
            .checked_add(value_len)
            .ok_or_else(|| Error::Corruption("block value length overflow".into()))?;
        if value_end > entries_end {
            return Err(Error::Corruption(
                "block entry exceeds block boundary".into(),
            ));
        }
        let mut key = last_key[..shared].to_vec();
        key.extend_from_slice(&input[pos..key_end]);
        result.push((key.clone(), input[key_end..value_end].to_vec()));
        last_key = key;
        pos = value_end;
    }
    if pos != entries_end {
        return Err(Error::Corruption(
            "block entries overlap restart array".into(),
        ));
    }
    Ok(result)
}

pub(crate) fn parse_index_entries(
    input: &[u8],
    value_delta_encoding: bool,
) -> Result<Vec<(Vec<u8>, BlockHandle)>> {
    if !value_delta_encoding {
        return parse_entries(input)?
            .into_iter()
            .map(|(key, value)| {
                let mut pos = 0;
                let handle = BlockHandle::decode(&value, &mut pos)?;
                if pos != value.len() {
                    return Err(Error::Unsupported(
                        "index values containing first-key metadata".into(),
                    ));
                }
                Ok((key, handle))
            })
            .collect();
    }

    let entries_end = entries_end(input)?;
    let mut result = Vec::new();
    let mut pos = 0;
    let mut last_key = Vec::new();
    let mut previous_handle: Option<BlockHandle> = None;
    while pos < entries_end {
        let shared = coding::read_varint32(input, &mut pos)? as usize;
        let non_shared = coding::read_varint32(input, &mut pos)? as usize;
        if shared > last_key.len() {
            return Err(Error::Corruption(
                "index shared prefix exceeds prior key".into(),
            ));
        }
        let key_end = pos
            .checked_add(non_shared)
            .ok_or_else(|| Error::Corruption("index key length overflow".into()))?;
        if key_end > entries_end {
            return Err(Error::Corruption("index key exceeds block boundary".into()));
        }
        let mut key = last_key[..shared].to_vec();
        key.extend_from_slice(&input[pos..key_end]);
        pos = key_end;
        let handle = match previous_handle {
            Some(previous) if shared != 0 => {
                let delta = coding::read_zigzag64(input, &mut pos)?;
                let size = previous
                    .size
                    .checked_add_signed(delta)
                    .ok_or_else(|| Error::Corruption("delta-encoded block size overflow".into()))?;
                BlockHandle {
                    offset: previous
                        .offset
                        .checked_add(previous.size)
                        .and_then(|offset| offset.checked_add(5))
                        .ok_or_else(|| Error::Corruption("block offset overflow".into()))?,
                    size,
                }
            }
            _ => BlockHandle::decode(input, &mut pos)?,
        };
        result.push((key.clone(), handle));
        last_key = key;
        previous_handle = Some(handle);
    }
    Ok(result)
}

fn entries_end(input: &[u8]) -> Result<usize> {
    let count_bytes = input
        .get(input.len().saturating_sub(4)..)
        .filter(|bytes| bytes.len() == 4)
        .ok_or_else(|| Error::Corruption("block has no restart count".into()))?;
    let count = u32::from_le_bytes(count_bytes.try_into().unwrap()) as usize;
    let restart_bytes = count
        .checked_add(1)
        .and_then(|count| count.checked_mul(4))
        .ok_or_else(|| Error::Corruption("restart array length overflow".into()))?;
    let entries_end = input
        .len()
        .checked_sub(restart_bytes)
        .ok_or_else(|| Error::Corruption("restart array exceeds block".into()))?;
    for index in 0..count {
        let start = entries_end + index * 4;
        let restart = u32::from_le_bytes(input[start..start + 4].try_into().unwrap()) as usize;
        if restart > entries_end {
            return Err(Error::Corruption("restart offset exceeds entries".into()));
        }
    }
    Ok(entries_end)
}