/*
    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)]

mod buffer;
pub mod inspect;
mod registry;
#[cfg(test)]
mod tests;
mod value;
pub use buffer::*;
pub use registry::*;
pub use value::*;

use std::marker::PhantomData;

#[repr(transparent)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct Tag(pub u32);
pub struct TypedTag<T>(pub Tag, pub PhantomData<T>);

#[derive(Debug, Clone, Copy)]
pub struct Object<'a> {
    tags: &'a [u32],
    offsets: &'a [u32],
    values: &'a [u32],
}
impl<'a> Object<'a> {
    pub fn load(buf: &'a [u32]) -> Option<Self> {
        let nf = *buf.get(0)? as usize;
        if buf.len() < 1 + nf * 2 {
            return None;
        }
        Some(Self {
            tags: &buf[1..1 + nf],
            offsets: &buf[1 + nf..1 + nf + nf],
            values: &buf[1 + nf + nf..],
        })
    }
    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)
    }
    #[inline]
    fn offset(&self, i: usize) -> usize {
        self.offsets
            .get(i)
            .map(|&v| v >> 2)
            .unwrap_or(self.values.len() as u32) as usize
    }
    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) << 2);

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

        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) << 2);

        let start = (start_raw >> 2) * 4;
        let padding = start_raw & 0b11;
        let end = (end_raw >> 2) * 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<'b: 'a, T: Value<'b>>(&'b self, index: usize) -> Option<T> {
        if T::ALIGNED {
            T::load_aligned(self.get_aligned(index)?)
        } else {
            T::load_unaligned(self.get_unaligned(index)?)
        }
    }
    pub fn get<'b: 'a, T: Value<'b>>(&'b self, tag: TypedTag<T>) -> Option<T> {
        self.get_typed(self.find_field(tag.0)?)
    }
    pub fn keys<'b: 'a>(&'b self) -> ObjectIter<'b> {
        ObjectIter {
            object: self,
            index: 0,
        }
    }
    pub fn iter<'b: 'a, T>(&'b self, tag: TypedTag<T>) -> FieldIter<'b, T> {
        FieldIter {
            object: self,
            index: self.tags.partition_point(|&x| x < tag.0.0),
            tag: tag.0.0,
            ty: PhantomData,
        }
    }
    pub fn insert<T: ValueStore>(&self, tag: TypedTag<T>, value: T) -> ObjectBuffer {
        self.insert_multi(tag, &[value])
    }
    pub fn insert_multi<T: ValueStore>(&self, tag: TypedTag<T>, values: &[T]) -> ObjectBuffer {
        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 + 1);
        let values_suffix = self.offset(suffix + 1);

        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) << 2;
            if val.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(bytemuck::cast_slice(&temp));
            }
        }
        let insert_size = (buf.len() - values_new) as u32;
        buf[suffix_offs..suffix_end]
            .iter_mut()
            .for_each(|e| *e += insert_size << 2);
        buf.extend(&self.values[values_suffix..]);

        ObjectBuffer(buf)
    }
}

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

pub struct FieldIter<'a, T> {
    object: &'a Object<'a>,
    index: usize,
    tag: u32,
    ty: PhantomData<T>,
}
impl<'a, T: Value<'a>> Iterator for FieldIter<'a, T> {
    type Item = T;
    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(self.index);
        self.index += 1;
        val
    }
}