serde_arrow 0.14.0

use std::collections::{BTreeMap, HashMap};

use marrow::{
    array::{Array, FixedSizeBinaryArray},
    datatypes::FieldMeta,
};
use serde::{Serialize, Serializer};

use crate::internal::{
    error::{fail, set_default, try_, Context, ContextSupport, Error, Result},
    serialization::utils::impl_serializer,
    utils::array_ext::{ArrayExt, CountArray, SeqArrayExt},
};

use super::{array_builder::ArrayBuilder, binary_builder::U8Serializer};

#[derive(Debug, Clone)]

pub struct FixedSizeBinaryBuilder {
    pub name: String,
    pub seq: CountArray,
    pub buffer: Vec<u8>,
    pub current_n: usize,
    pub n: usize,
    metadata: HashMap<String, String>,
}

impl FixedSizeBinaryBuilder {
    pub fn new(
        name: String,
        n: usize,
        is_nullable: bool,
        metadata: HashMap<String, String>,
    ) -> Self {
        Self {
            name,
            seq: CountArray::new(is_nullable),
            buffer: Vec::new(),
            n,
            current_n: 0,
            metadata,
        }
    }

    pub fn take(&mut self) -> ArrayBuilder {
        ArrayBuilder::FixedSizeBinary(Self {
            name: self.name.clone(),
            metadata: self.metadata.clone(),
            seq: self.seq.take(),
            buffer: std::mem::take(&mut self.buffer),
            current_n: std::mem::take(&mut self.current_n),
            n: self.n,
        })
    }

    pub fn is_nullable(&self) -> bool {
        self.seq.validity.is_some()
    }

    pub fn into_array_and_field_meta(self) -> Result<(Array, FieldMeta)> {
        let meta = FieldMeta {
            name: self.name,
            metadata: self.metadata,
            nullable: self.seq.validity.is_some(),
        };
        let array = Array::FixedSizeBinary(FixedSizeBinaryArray {
            n: self.n.try_into()?,
            validity: self.seq.validity,
            data: self.buffer,
        });
        Ok((array, meta))
    }

    pub fn reserve(&mut self, additional: usize) {
        self.seq.reserve(additional);
        self.buffer.reserve(additional * self.n);
    }

    pub fn serialize_default_value(&mut self) -> Result<()> {
        try_(|| {
            self.seq.push_seq_default()?;
            for _ in 0..self.n {
                self.buffer.push(0);
            }
            Ok(())
        })
        .ctx(self)
    }

    pub fn serialize_value<V: Serialize>(&mut self, value: V) -> Result<()> {
        value.serialize(&mut *self).ctx(self)
    }
}

impl FixedSizeBinaryBuilder {
    fn start(&mut self) -> Result<()> {
        self.current_n = 0;
        self.seq.start_seq()
    }

    fn element<V: Serialize + ?Sized>(&mut self, value: &V) -> Result<()> {
        let byte = value.serialize(U8Serializer)?;
        self.buffer.push(byte);
        self.current_n += 1;

        Ok(())
    }

    fn end(&mut self) -> Result<()> {
        if self.current_n != self.n {
            fail!(
                "Invalid number of elements for fixed size binary: got {actual}, expected {expected}",
                actual = self.current_n,
                expected = self.n,
            );
        }
        self.seq.end_seq()
    }
}

impl Context for FixedSizeBinaryBuilder {
    fn annotate(&self, annotations: &mut BTreeMap<String, String>) {
        set_default(annotations, "field", &self.name);
        set_default(
            annotations,
            "data_type",
            format!("FixedSizeBinary({n})", n = self.n),
        );
    }
}

impl<'a> Serializer for &'a mut FixedSizeBinaryBuilder {
    impl_serializer!(
        'a, FixedSizeBinaryBuilder;
        override serialize_none,
        override serialize_seq,
        override serialize_tuple,
        override serialize_bytes,
        override serialize_str,
    );

    fn serialize_none(self) -> Result<()> {
        self.seq.push_seq_none()?;
        for _ in 0..self.n {
            self.buffer.push(0);
        }
        Ok(())
    }

    fn serialize_seq(self, _len: Option<usize>) -> Result<Self::SerializeSeq> {
        // TOOD: fix reservation
        self.start()?;
        Ok(super::utils::SerializeSeq::FixedSizeBinary(self))
    }

    fn serialize_tuple(self, _len: usize) -> Result<Self::SerializeTuple> {
        // TOOD: fix reservation
        self.start()?;
        Ok(Self::SerializeTuple::FixedSizeBinary(self))
    }

    fn serialize_bytes(self, v: &[u8]) -> Result<()> {
        if v.len() != self.n {
            fail!(
                in self,
                "Invalid number of elements for fixed size binary: got {actual}, expected {expected}",
                actual = v.len(),
                expected = self.n,
            );
        }

        self.seq.start_seq()?;
        self.buffer.extend(v);
        self.seq.end_seq()
    }

    fn serialize_str(self, v: &str) -> Result<()> {
        serde::Serializer::serialize_bytes(self, v.as_bytes())
    }
}

impl serde::ser::SerializeSeq for &mut FixedSizeBinaryBuilder {
    type Ok = ();
    type Error = Error;

    fn serialize_element<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<()> {
        self.element(value)
    }

    fn end(self) -> Result<()> {
        FixedSizeBinaryBuilder::end(&mut *self)
    }
}

impl serde::ser::SerializeTuple for &mut FixedSizeBinaryBuilder {
    type Ok = ();
    type Error = Error;

    fn serialize_element<T: ?Sized + Serialize>(&mut self, value: &T) -> Result<()> {
        self.element(value)
    }

    fn end(self) -> Result<()> {
        FixedSizeBinaryBuilder::end(&mut *self)
    }
}