openvariant 0.1.0

use std::borrow::Cow;

use anyhow::{Result, bail, ensure};

#[derive(Debug)]
pub struct VariantMetadata<'a> {
    pub(crate) header: VariantHeader,
    pub(crate) offsets: Vec<u32>,
    pub(crate) strings: Vec<&'a str>,
}

impl<'a> VariantMetadata<'a> {
    pub fn find_by_field_id(&self, field_id: u32) -> Option<&'a str> {
        match self.header.layout {
            DictionaryLayout::SortedUnique => todo!(),
            DictionaryLayout::Naive => {
                if field_id >= self.dictionary_len() as u32 {
                    return None;
                }

                let field_id = field_id as usize;
                self.strings.get(field_id).map(|v| &**v)
            }
        }
    }

    pub fn find_by_field_name(&self, field_name: &str) -> Option<&'a str> {
        match self.header.layout {
            DictionaryLayout::SortedUnique => self.search_sorted_unique(field_name),
            DictionaryLayout::Naive => self.search_naive(field_name),
        }
    }

    pub const fn dictionary_len(&self) -> usize {
        self.offsets.len() - 1
    }

    pub const fn is_empty(&self) -> bool {
        self.dictionary_len() == 0
    }

    pub fn dictionary_strings(&self) -> &[&str] {
        &self.strings
    }

    fn search_sorted_unique(&self, field_name: &str) -> Option<&'a str> {
        if self.dictionary_len() < 6 {
            return self.search_naive(field_name);
        }

        self.strings
            .binary_search(&field_name)
            .map(|i| self.strings[i])
            .ok()
    }

    fn search_naive(&self, field_name: &str) -> Option<&'a str> {
        self.strings
            .iter()
            .find(|&&n| n == field_name)
            .map(|v| &**v)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DictionaryLayout {
    /// dictionary strings are unique and sorted in lexicographic order
    SortedUnique,
    /// dictionary strings may contain duplicates and have no ordering
    Naive,
}

#[derive(Debug, Clone, Copy)]
pub struct VariantHeader {
    pub layout: DictionaryLayout,
    pub offset_width: OffsetWidth,
}

impl VariantHeader {
    pub const VERSION: u8 = 1;
}

impl TryFrom<u8> for VariantHeader {
    type Error = anyhow::Error;

    fn try_from(b: u8) -> Result<Self, Self::Error> {
        ensure!(b & 0b1111 == 1, "version must be 1");

        Ok(Self {
            layout: if ((b >> 4) & 0b1) == 1 {
                DictionaryLayout::SortedUnique
            } else {
                DictionaryLayout::Naive
            },
            offset_width: OffsetWidth::try_from(b >> 6)?,
        })
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OffsetWidth {
    U8 = 0,
    U16 = 1,
    U24 = 2,
    U32 = 3,
}

impl OffsetWidth {
    pub const fn find_offset(n: u32) -> Self {
        if n <= u8::MAX as u32 {
            return Self::U8;
        }

        if n <= u16::MAX as u32 {
            return Self::U16;
        }

        if n <= 0xFFFFFF_u32 {
            return Self::U24;
        }

        Self::U32
    }
}

impl TryFrom<u8> for OffsetWidth {
    type Error = anyhow::Error;

    fn try_from(b: u8) -> Result<Self> {
        let w = match b {
            0 => Self::U8,
            1 => Self::U16,
            2 => Self::U24,
            3 => Self::U32,
            _ => bail!("Offset width can not be larger than 3"),
        };

        Ok(w)
    }
}

/// The Variant Physical Type
#[derive(Debug, PartialEq)]
pub enum Variant<'a> {
    // Primitive values
    Null,                    // 0
    BooleanTrue,             // 1
    BooleanFalse,            // 2
    Int8(i8),                // 3
    Int16(i16),              // 4
    Int32(i32),              // 5
    Int64(i64),              // 6
    Double(f64),             // 7
    Decimal4(u8, i32),       // 8
    Decimal8(u8, i64),       // 9
    Decimal16(u8, i128),     // 10
    Date(u32),               // 11
    TimestampMicros(u64),    // 12
    TimestampNTZMicros(u64), // 13
    Float(f32),              // 14
    Binary(Cow<'a, [u8]>),   // 15
    String(Cow<'a, str>),    // 16
    TimeNTZMicros(u64),      // 17
    TimestampNanos(u64),     // 18
    TimestampNTZNanos(u64),  // 19
    UUID(u128),              // 20

    // Experimental types
    // you could theoretically have up to 64 tags
    Int16Leb128(i16),   // 21
    Int32Leb128(i32),   // 22
    Int64Leb128(i64),   // 23
    Int128Leb128(i128), // 24

    // Complex types
    ShortString(ShortString<'a>),
    Object(VariantObject<'a>),
    Array(VariantArray<'a>),
}

impl From<()> for Variant<'_> {
    fn from(_: ()) -> Self {
        Self::Null
    }
}

impl From<bool> for Variant<'_> {
    fn from(value: bool) -> Self {
        if value {
            Self::BooleanTrue
        } else {
            Self::BooleanFalse
        }
    }
}

impl From<String> for Variant<'_> {
    fn from(value: String) -> Self {
        if value.len() < 64 {
            Self::ShortString(ShortString(Cow::Owned(value)))
        } else {
            Self::String(Cow::Owned(value))
        }
    }
}

impl<'a> From<&'a str> for Variant<'a> {
    fn from(value: &'a str) -> Self {
        if value.len() < 64 {
            Self::ShortString(ShortString(Cow::Borrowed(value)))
        } else {
            Self::String(Cow::Borrowed(value))
        }
    }
}

impl From<Vec<u8>> for Variant<'_> {
    fn from(value: Vec<u8>) -> Self {
        Self::Binary(Cow::Owned(value))
    }
}

impl<'a> From<&'a [u8]> for Variant<'a> {
    fn from(value: &'a [u8]) -> Self {
        Self::Binary(Cow::Borrowed(value))
    }
}

impl From<i8> for Variant<'_> {
    fn from(value: i8) -> Self {
        Self::Int8(value)
    }
}

impl From<i16> for Variant<'_> {
    fn from(value: i16) -> Self {
        Self::Int16(value)
    }
}

impl From<i32> for Variant<'_> {
    fn from(value: i32) -> Self {
        Self::Int32(value)
    }
}

impl From<i64> for Variant<'_> {
    fn from(value: i64) -> Self {
        Self::Int64(value)
    }
}

impl From<f64> for Variant<'_> {
    fn from(value: f64) -> Self {
        Self::Double(value)
    }
}

#[derive(Debug, PartialEq, Eq)]
pub struct ShortString<'a>(Cow<'a, str>);

impl<'a> ShortString<'a> {
    pub const fn new(str: &'a str) -> Self {
        Self(Cow::Borrowed(str))
    }

    pub fn inner(&self) -> &str {
        &self.0
    }

    pub fn len(&self) -> u8 {
        self.0.len() as u8
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

impl<'a> TryFrom<&'a str> for ShortString<'a> {
    type Error = anyhow::Error;

    fn try_from(s: &'a str) -> Result<Self, Self::Error> {
        ensure!(s.len() < 64, "short string must be less than 64 bytes");

        Ok(Self(Cow::Borrowed(s)))
    }
}

impl<'a> TryFrom<String> for ShortString<'a> {
    type Error = anyhow::Error;

    fn try_from(s: String) -> Result<Self, Self::Error> {
        ensure!(s.len() < 64, "short string must be less than 64 bytes");

        Ok(Self(Cow::Owned(s)))
    }
}

#[derive(Debug, PartialEq)]
pub struct VariantArray<'a> {
    pub(crate) field_offsets: Vec<u32>,
    pub(crate) values: Vec<Variant<'a>>,
}

impl<'a> VariantArray<'a> {
    pub const fn len(&self) -> u32 {
        self.values.len() as u32
    }

    pub const fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

#[derive(Debug, PartialEq)]
pub struct VariantObject<'a> {
    pub(crate) field_ids: Vec<u32>,
    pub(crate) field_offsets: Vec<u32>,
    pub(crate) values: Vec<Variant<'a>>,
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn test_offset_width() {
        assert_eq!(OffsetWidth::try_from(0).unwrap(), OffsetWidth::U8);
        assert_eq!(OffsetWidth::try_from(1).unwrap(), OffsetWidth::U16);
        assert_eq!(OffsetWidth::try_from(2).unwrap(), OffsetWidth::U24);
        assert_eq!(OffsetWidth::try_from(3).unwrap(), OffsetWidth::U32);

        // error cases
        assert!(OffsetWidth::try_from(4).is_err());
        assert!(OffsetWidth::try_from(std::u8::MAX).is_err());
    }

    #[test]
    fn test_sorted_header() {
        let b = 0b1101_0001; // sorted string, offset width is u32
        let header = VariantHeader::try_from(b).unwrap();

        assert_eq!(header.layout, DictionaryLayout::SortedUnique);
        assert_eq!(header.offset_width, OffsetWidth::U32);
    }
}