/*
    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
    Copyright (C) 2026 metamuffin <metamuffin.org>
*/
#![feature(iter_array_chunks, strip_circumfix)]

mod buffer;
pub mod debug;
#[cfg(feature = "json")]
pub mod json;
mod path;
pub mod tag;
#[cfg(test)]
mod tests;
mod value;

pub use buffer::*;
pub use path::*;
pub use tag::*;
pub use value::*;

use std::{collections::BTreeSet, hash::Hash, marker::PhantomData};

#[derive(Hash, PartialEq, Eq)]
#[repr(transparent)]
pub struct Object([u32]);

impl ToOwned for Object {
    type Owned = Box<Object>;
    fn to_owned(&self) -> Self::Owned {
        vec_to_ob(self.0.to_vec())
    }
}

pub const EMPTY: &Object = slice_to_ob(&[0]);

impl Object {
    #[inline]
    fn len(&self) -> usize {
        self.0[0] as usize
    }
    #[inline]
    fn tags(&self) -> &[u32] {
        &self.0[1..1 + self.len()]
    }
    #[inline]
    fn offsets(&self) -> &[u32] {
        let nf = self.len();
        &self.0[1 + nf..1 + nf + nf]
    }
    #[inline]
    fn values(&self) -> &[u32] {
        let nf = self.len();
        &self.0[1 + nf + nf..]
    }
    pub fn export(&self) -> &[u32] {
        &self.0
    }

    pub fn find_field(&self, tag: Tag) -> Option<usize> {
        // using partition as binary search for the first field (instead of regular binary_search that returns any)
        let first = self.tags().partition_point(|&x| x < tag.0);
        self.tags()
            .get(first)
            .is_some_and(|&x| x == tag.0)
            .then_some(first)
    }
    pub fn has(&self, tag: Tag) -> bool {
        self.find_field(tag).is_some()
    }
    #[inline]
    fn offset(&self, i: usize) -> usize {
        self.offsets()
            .get(i)
            .map(|&v| v >> 5)
            .unwrap_or(self.values().len() as u32) as usize
    }
    fn offset_type(&self, i: usize) -> ValueType {
        let raw = (self.offsets()[i] >> 2) & 0b111;
        ValueType::from_num(raw)
    }
    fn size(&self, i: usize) -> u32 {
        let start_raw = self.offsets()[i];
        let end_raw = self
            .offsets()
            .get(i + 1)
            .copied()
            .unwrap_or((self.values().len() as u32) << 5);

        let u32_len = (end_raw >> 5) - (start_raw >> 5);
        let padding = start_raw & 0b11;
        u32_len * 4 - padding
    }
    fn get_aligned(&self, index: usize) -> Option<&[u32]> {
        let start_raw = self.offsets()[index];
        let end_raw = self
            .offsets()
            .get(index + 1)
            .copied()
            .unwrap_or((self.values().len() as u32) << 5);

        let start = start_raw >> 5;
        let end = end_raw >> 5;

        Some(&self.values()[start as usize..end as usize])
    }

    fn get_unaligned(&self, index: usize) -> Option<&[u8]> {
        let start_raw = self.offsets()[index];
        let end_raw = self
            .offsets()
            .get(index + 1)
            .copied()
            .unwrap_or((self.values().len() as u32) << 5);

        let start = (start_raw >> 5) * 4;
        let padding = start_raw & 0b11;
        let end = (end_raw >> 5) * 4 - padding;

        let values_u8: &[u8] = bytemuck::cast_slice(self.values());
        Some(&values_u8[start as usize..end as usize])
    }
    #[inline]
    pub fn get_typed<'a, T: ValueMarker<'a> + ?Sized>(&'a self, index: usize) -> Option<T::Inner> {
        if T::Inner::TYPE.is_aligned() {
            T::Inner::load_aligned(self.get_aligned(index)?)
        } else {
            T::Inner::load_unaligned(self.get_unaligned(index)?)
        }
    }
    pub fn get<'a, T: ValueMarker<'a> + ?Sized>(&'a self, tag: TypedTag<T>) -> Option<T::Inner> {
        self.get_typed::<T>(self.find_field(tag.0)?)
    }
    pub fn keys<'a>(&'a self) -> KeysIter<'a> {
        KeysIter {
            object: self,
            index: 0,
        }
    }
    pub fn entries<'a, T: ?Sized>(&'a self) -> EntriesIter<'a, T> {
        EntriesIter {
            object: self,
            index: 0,
            ty: PhantomData,
        }
    }
    pub fn iter<'a, T: ?Sized>(&'a self, tag: TypedTag<T>) -> FieldIter<'a, T> {
        FieldIter {
            object: self,
            index: self.tags().partition_point(|&x| x < tag.0.0),
            tag: tag.0.0,
            ty: PhantomData,
        }
    }
    #[must_use]
    pub fn extend<'a, T: ValueMarker<'a> + ?Sized>(
        &'a self,
        tag: TypedTag<T>,
        values: impl IntoIterator<Item = T::Inner>,
    ) -> Box<Object>
    where
        T::Inner: Eq + Ord,
    {
        self.insert_multi(
            tag,
            &self
                .iter(tag)
                .chain(values)
                .collect::<BTreeSet<_>>()
                .into_iter()
                .collect::<Vec<_>>(),
        )
    }
    #[must_use]
    pub fn extend_object(
        &self,
        tag: TypedTag<Object>,
        ident: Tag,
        values: &[Box<Object>],
    ) -> Box<Object> {
        let ident = TypedTag::<[u8]>::new(ident);
        let mut new_vals = Vec::new();
        for ob in self.iter(tag) {
            new_vals.push(ob);
        }
        let mut any_new = false;
        for val in values {
            if new_vals.iter().all(|rhs| rhs.get(ident) != val.get(ident)) {
                any_new = true;
                new_vals.push(val);
            }
        }
        if any_new {
            self.insert_multi(TypedTag::<Object>::new(tag.0), &new_vals)
        } else {
            self.to_owned()
        }
    }
    #[must_use]
    #[inline]
    pub fn insert<'a, T: ValueMarker<'a> + ?Sized>(
        &'a self,
        tag: TypedTag<T>,
        value: T::Inner,
    ) -> Box<Object> {
        self.insert_multi(tag, &[value])
    }
    #[must_use]
    #[inline]
    pub fn remove<'a, T: ValueMarker<'a>>(&'a self, tag: TypedTag<T>) -> Box<Object> {
        self.insert_multi(tag, &[])
    }
    #[must_use]
    pub fn insert_multi<'a, T: ValueMarker<'a> + ?Sized>(
        &'a self,
        tag: TypedTag<T>,
        values: &[T::Inner],
    ) -> Box<Object> {
        let prefix = self.tags().partition_point(|&x| x < tag.0.0);
        let suffix = self.tags().partition_point(|&x| x <= tag.0.0);
        let values_prefix = self.offset(prefix);
        let values_suffix = self.offset(suffix);

        let mut buf = Vec::new();
        let cut_size = suffix - prefix;
        buf.push((self.tags().len() - cut_size + values.len()) as u32);

        buf.extend(&self.tags()[..prefix]);
        buf.extend(values.iter().map(|_| tag.0.0));
        buf.extend(&self.tags()[suffix..]);

        buf.extend(&self.offsets()[..prefix]);
        let new_offs = buf.len();
        buf.extend(values.iter().map(|_| 0)); // placeholder
        let suffix_offs = buf.len();
        buf.extend(&self.offsets()[suffix..]); // need offsetting later
        let suffix_end = buf.len();

        let values_start = buf.len() as u32;
        buf.extend(&self.values()[..values_prefix]);
        let mut temp = Vec::new();
        let values_new = buf.len();
        for (i, val) in values.iter().enumerate() {
            let off = (buf.len() as u32 - values_start) << 5 | (T::Inner::TYPE as u32) << 2;
            if T::Inner::TYPE.is_aligned() {
                buf[new_offs + i] = off;
                val.store_aligned(&mut buf);
            } else {
                temp.clear();
                val.store_unaligned(&mut temp);
                let pad = pad_vec(&mut temp);
                buf[new_offs + i] = off | pad;
                buf.extend(&*slice_u8_to_u32(&temp));
            }
        }
        let values_insert_size = buf.len() - values_new;
        let values_cut_size = values_suffix - values_prefix;
        let suffix_offset = (values_insert_size as i32 - values_cut_size as i32) << 5;
        if suffix_offset != 0 {
            buf[suffix_offs..suffix_end]
                .iter_mut()
                .for_each(|e| *e = e.strict_add_signed(suffix_offset));
        }
        buf.extend(&self.values()[values_suffix..]);

        vec_to_ob(buf)
    }

    #[must_use]
    pub fn update(
        &self,
        tag: TypedTag<Object>,
        update: impl FnOnce(&Object) -> Box<Object>,
    ) -> Box<Object> {
        let ob = update(self.get(tag).unwrap_or(EMPTY));
        self.insert(tag, &ob)
    }
}

pub struct KeysIter<'a> {
    object: &'a Object,
    index: usize,
}
impl Iterator for KeysIter<'_> {
    type Item = Tag;
    fn next(&mut self) -> Option<Self::Item> {
        if self.index >= self.object.tags().len() {
            None
        } else {
            self.index += 1;
            Some(Tag(self.object.tags()[self.index - 1]))
        }
    }
}

pub struct EntriesIter<'a, T: ?Sized> {
    object: &'a Object,
    index: usize,
    ty: PhantomData<T>,
}
impl<'a, T: ValueMarker<'a> + ?Sized> Iterator for EntriesIter<'a, T> {
    type Item = (Tag, T::Inner);
    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if self.index >= self.object.tags().len() {
                return None;
            }
            if T::Inner::TYPE != self.object.offset_type(self.index) {
                self.index += 1;
                continue;
            }
            let value = self.object.get_typed::<T>(self.index)?;
            let tag = self.object.tags()[self.index];
            self.index += 1;
            return Some((Tag(tag), value));
        }
    }
}

pub struct FieldIter<'a, T: ?Sized> {
    object: &'a Object,
    index: usize,
    tag: u32,
    ty: PhantomData<T>,
}
impl<'a, T: ValueMarker<'a> + ?Sized> Iterator for FieldIter<'a, T> {
    type Item = T::Inner;
    fn next(&mut self) -> Option<Self::Item> {
        if self.index >= self.object.tags().len() {
            return None;
        }
        if self.object.tags()[self.index] != self.tag {
            return None;
        }
        let val = self.object.get_typed::<T>(self.index);
        self.index += 1;
        val
    }
}