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

use crate::{Object, Tag, TypedTag, ValueStore};
use bytemuck::{try_cast_slice, try_cast_vec};
use log::trace;
use std::borrow::Cow;

#[derive(PartialEq, Eq, Hash)]
pub struct ObjectBuffer(pub Vec<u32>);

impl ObjectBuffer {
    pub fn empty() -> Self {
        Self(vec![0])
    }
    pub fn as_object<'a>(&'a self) -> Object<'a> {
        Object::load(&self.0).unwrap()
    }
    pub fn to_bytes(self) -> Vec<u8> {
        self.0.into_iter().flat_map(u32::to_le_bytes).collect()
    }
    pub fn new(fields: &mut [(Tag, &dyn ValueStore)]) -> ObjectBuffer {
        let mut tags = Vec::new();
        let mut offsets = Vec::new();
        let mut values = Vec::new();
        fields.sort_by_key(|(t, _)| t.0);
        let mut temp = Vec::new();
        for (tag, val) in fields {
            tags.push(tag.0);
            let ty = val.get_type();
            let off = (values.len() as u32) << 5 | (ty as u32) << 2;
            if ty.is_aligned() {
                offsets.push(off);
                val.store_aligned(&mut values);
            } else {
                temp.clear();
                val.store_unaligned(&mut temp);
                let pad = pad_vec(&mut temp);
                offsets.push(off | pad);
                values.extend(&*slice_u8_to_u32(&temp));
            }
        }
        ObjectBuffer(
            [tags.len() as u32]
                .into_iter()
                .chain(tags)
                .chain(offsets)
                .chain(values)
                .collect(),
        )
    }
}

pub type OBB = ObjectBufferBuilder;

#[derive(Default)]
pub struct ObjectBufferBuilder(Vec<(Tag, u32, Vec<u32>)>);
impl ObjectBufferBuilder {
    pub fn new() -> Self {
        Self::default()
    }
    pub fn push<T: ValueStore>(&mut self, tag: TypedTag<T>, value: T) {
        let ty = value.get_type();
        let tyb = (ty as u32) << 2;
        if ty.is_aligned() {
            let mut buf = Vec::new();
            value.store_aligned(&mut buf);
            self.0.push((tag.0, tyb, buf));
        } else {
            let mut buf = Vec::new();
            value.store_unaligned(&mut buf);
            let pad = pad_vec(&mut buf);
            self.0.push((tag.0, tyb | pad, vec_u8_to_u32(buf)));
        }
    }
    pub fn with<T: ValueStore>(mut self, tag: TypedTag<T>, value: T) -> Self {
        self.push(tag, value);
        self
    }
    pub fn finish(mut self) -> ObjectBuffer {
        let mut tags = Vec::new();
        let mut offsets = Vec::new();
        let mut values = Vec::new();
        self.0.sort_by_key(|(t, _, _)| t.0);
        for (tag, mut off, buf) in self.0 {
            tags.push(tag.0);
            off |= (values.len() as u32) << 5;
            offsets.push(off);
            values.extend(buf);
        }
        ObjectBuffer(
            [tags.len() as u32]
                .into_iter()
                .chain(tags)
                .chain(offsets)
                .chain(values)
                .collect(),
        )
    }
}

impl From<Vec<u8>> for ObjectBuffer {
    fn from(value: Vec<u8>) -> Self {
        ObjectBuffer(vec_u8_to_u32(value))
    }
}

#[inline]
pub(super) fn pad_vec(temp: &mut Vec<u8>) -> u32 {
    let mut pad = 0;
    while temp.len() % 4 != 0 {
        pad += 1;
        temp.push(0);
    }
    pad
}

pub fn vec_u8_to_u32(value: Vec<u8>) -> Vec<u32> {
    try_cast_vec(value).unwrap_or_else(|(_, v)| {
        trace!("encountered unalined vec");
        v.into_iter()
            .array_chunks()
            .map(u32::from_ne_bytes)
            .collect()
    })
}
pub fn slice_u8_to_u32<'a>(value: &'a [u8]) -> Cow<'a, [u32]> {
    try_cast_slice(value)
        .map(Cow::Borrowed)
        .unwrap_or_else(|_| {
            trace!("encountered unalined slice");
            Cow::Owned(
                value
                    .into_iter()
                    .copied()
                    .array_chunks()
                    .map(u32::from_ne_bytes)
                    .collect(),
            )
        })
}