fieldmasker 0.0.1

use crate::spec::{Kind, Node};
use crate::trie::MaskTrie;
use serde::ser::{
    SerializeMap, SerializeSeq, SerializeStruct, SerializeStructVariant, SerializeTuple,
    SerializeTupleStruct, SerializeTupleVariant,
};
use serde::{Serialize, Serializer};
use std::collections::BTreeMap;

pub(crate) struct Selective<'a, S> {
    pub(crate) ser: S,
    pub(crate) trie: &'a MaskTrie,
    pub(crate) spec: &'a Node,
    pub(crate) include_all: bool,
}

// Wrap any Serialize value (no MaskSpec bound; allow ?Sized).
struct WithMask<'a, T: ?Sized> {
    v: &'a T,
    trie: &'a MaskTrie,
    spec: &'a Node,
    include_all: bool,
}

impl<'a, T: ?Sized + Serialize> Serialize for WithMask<'a, T> {
    fn serialize<S: Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
        let gate = Selective {
            ser: s,
            trie: self.trie,
            spec: self.spec,
            include_all: self.include_all,
        };
        self.v.serialize(gate)
    }
}

impl<'a, S: Serializer> Serializer for Selective<'a, S> {
    type Ok = S::Ok;
    type Error = S::Error;

    type SerializeSeq = GateSeq<'a, S::SerializeSeq>;
    type SerializeTuple = GateSeq<'a, S::SerializeTuple>;
    type SerializeTupleStruct = GateSeq<'a, S::SerializeTupleStruct>;
    type SerializeTupleVariant = GateSeqVar<'a, S::SerializeTupleVariant>;
    type SerializeMap = GateMap<'a, S::SerializeMap>;
    type SerializeStruct = GateStruct<'a, S::SerializeStruct>;
    type SerializeStructVariant = GateStructVar<'a, S::SerializeStructVariant>;

    // Scalars & simple passthroughs
    fn serialize_bool(self, v: bool) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_bool(v)
    }

    fn serialize_i8(self, v: i8) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_i8(v)
    }

    fn serialize_i16(self, v: i16) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_i16(v)
    }

    fn serialize_i32(self, v: i32) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_i32(v)
    }

    fn serialize_i64(self, v: i64) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_i64(v)
    }

    fn serialize_u8(self, v: u8) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_u8(v)
    }

    fn serialize_u16(self, v: u16) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_u16(v)
    }

    fn serialize_u32(self, v: u32) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_u32(v)
    }

    fn serialize_u64(self, v: u64) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_u64(v)
    }

    fn serialize_f32(self, v: f32) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_f32(v)
    }

    fn serialize_f64(self, v: f64) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_f64(v)
    }

    fn serialize_char(self, v: char) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_char(v)
    }

    fn serialize_str(self, v: &str) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_str(v)
    }

    fn serialize_bytes(self, v: &[u8]) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_bytes(v)
    }

    fn serialize_none(self) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_none()
    }

    fn serialize_some<Tx: ?Sized + Serialize>(self, value: &Tx) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_some(&WithMask {
            v: value,
            trie: self.trie,
            spec: self.spec,
            include_all: self.include_all,
        })
    }

    fn serialize_unit(self) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_unit()
    }

    fn serialize_unit_struct(self, name: &'static str) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_unit_struct(name)
    }

    fn serialize_unit_variant(
        self,
        n: &'static str,
        i: u32,
        v: &'static str,
    ) -> Result<Self::Ok, Self::Error> {
        self.ser.serialize_unit_variant(n, i, v)
    }

    // Newtype struct: just serialize inner via our gate (Serde permits identical representation)
    fn serialize_newtype_struct<Tx: ?Sized + Serialize>(
        self,
        _name: &'static str,
        value: &Tx,
    ) -> Result<Self::Ok, Self::Error> {
        value.serialize(Selective {
            ser: self.ser,
            trie: self.trie,
            spec: self.spec,
            include_all: self.include_all,
        })
    }

    fn serialize_newtype_variant<Tx: ?Sized + Serialize>(
        self,
        n: &'static str,
        i: u32,
        v: &'static str,
        value: &Tx,
    ) -> Result<Self::Ok, Self::Error> {
        // Variants are treated as a separate envelope; forward through underlying serializer.
        self.ser.serialize_newtype_variant(n, i, v, value)
    }

    // Seqs / tuples: descend into element spec; same trie applies to each element.
    fn serialize_seq(self, len: Option<usize>) -> Result<Self::SerializeSeq, Self::Error> {
        let inner = self.ser.serialize_seq(len)?;
        let elem = match &self.spec.kind {
            Kind::Array(elem) | Kind::Map(elem) => &**elem,
            _ => self.spec,
        };
        Ok(GateSeq {
            inner,
            trie: self.trie,
            spec: elem,
            include_all: self.include_all || self.trie.terminal,
        })
    }

    fn serialize_tuple(self, len: usize) -> Result<Self::SerializeTuple, Self::Error> {
        let inner = self.ser.serialize_tuple(len)?;
        Ok(GateSeq {
            inner,
            trie: self.trie,
            spec: self.spec,
            include_all: self.include_all,
        })
    }

    fn serialize_tuple_struct(
        self,
        name: &'static str,
        len: usize,
    ) -> Result<Self::SerializeTupleStruct, Self::Error> {
        let inner = self.ser.serialize_tuple_struct(name, len)?;
        Ok(GateSeq {
            inner,
            trie: self.trie,
            spec: self.spec,
            include_all: self.include_all,
        })
    }

    fn serialize_tuple_variant(
        self,
        name: &'static str,
        idx: u32,
        var: &'static str,
        len: usize,
    ) -> Result<Self::SerializeTupleVariant, Self::Error> {
        let inner = self.ser.serialize_tuple_variant(name, idx, var, len)?;
        Ok(GateSeqVar {
            inner,
            trie: self.trie,
            spec: self.spec,
            include_all: self.include_all,
        })
    }

    // Maps: apply same trie to values (drop all entries if no selection at this node).
    fn serialize_map(self, len: Option<usize>) -> Result<Self::SerializeMap, Self::Error> {
        let inner = self.ser.serialize_map(len)?;
        let elem = match &self.spec.kind {
            Kind::Map(elem) => &**elem,
            _ => self.spec,
        };
        let allow = self.include_all || self.trie.terminal || !self.trie.children.is_empty();
        Ok(GateMap {
            inner,
            trie: self.trie,
            spec: elem,
            allow,
        })
    }

    // Structs: gate fields by name using object spec and trie children.
    fn serialize_struct(
        self,
        name: &'static str,
        len: usize,
    ) -> Result<Self::SerializeStruct, Self::Error> {
        let inner = self.ser.serialize_struct(name, len)?;
        let fields = match &self.spec.kind {
            Kind::Object(map) => Some(map),
            _ => None,
        };
        Ok(GateStruct {
            inner,
            trie: self.trie,
            fields,
            include_all: self.include_all || self.trie.terminal,
        })
    }

    fn serialize_struct_variant(
        self,
        name: &'static str,
        idx: u32,
        var: &'static str,
        len: usize,
    ) -> Result<Self::SerializeStructVariant, Self::Error> {
        let inner = self.ser.serialize_struct_variant(name, idx, var, len)?;
        let fields = match &self.spec.kind {
            Kind::Object(map) => Some(map),
            _ => None,
        };
        Ok(GateStructVar {
            inner,
            trie: self.trie,
            fields,
            include_all: self.include_all || self.trie.terminal,
        })
    }
}

// Wrappers

pub(crate) struct GateSeq<'a, I> {
    inner: I,
    trie: &'a MaskTrie,
    spec: &'a Node,
    include_all: bool,
}

impl<'a, I: SerializeSeq> SerializeSeq for GateSeq<'a, I> {
    type Ok = I::Ok;
    type Error = I::Error;

    fn serialize_element<Tx: ?Sized + Serialize>(&mut self, value: &Tx) -> Result<(), Self::Error> {
        self.inner.serialize_element(&WithMask {
            v: value,
            trie: self.trie,
            spec: self.spec,
            include_all: self.include_all,
        })
    }

    fn end(self) -> Result<Self::Ok, Self::Error> {
        self.inner.end()
    }
}

impl<'a, I: SerializeTuple> SerializeTuple for GateSeq<'a, I> {
    type Ok = I::Ok;
    type Error = I::Error;

    fn serialize_element<Tx: ?Sized + Serialize>(&mut self, value: &Tx) -> Result<(), Self::Error> {
        self.inner.serialize_element(&WithMask {
            v: value,
            trie: self.trie,
            spec: self.spec,
            include_all: self.include_all,
        })
    }

    fn end(self) -> Result<Self::Ok, Self::Error> {
        self.inner.end()
    }
}

impl<'a, I: SerializeTupleStruct> SerializeTupleStruct for GateSeq<'a, I> {
    type Ok = I::Ok;
    type Error = I::Error;

    fn serialize_field<Tx: ?Sized + Serialize>(&mut self, value: &Tx) -> Result<(), Self::Error> {
        self.inner.serialize_field(&WithMask {
            v: value,
            trie: self.trie,
            spec: self.spec,
            include_all: self.include_all,
        })
    }

    fn end(self) -> Result<Self::Ok, Self::Error> {
        self.inner.end()
    }
}

pub(crate) struct GateSeqVar<'a, I> {
    inner: I,
    trie: &'a MaskTrie,
    spec: &'a Node,
    include_all: bool,
}

impl<'a, I: SerializeTupleVariant> SerializeTupleVariant for GateSeqVar<'a, I> {
    type Ok = I::Ok;
    type Error = I::Error;

    fn serialize_field<Tx: ?Sized + Serialize>(&mut self, value: &Tx) -> Result<(), Self::Error> {
        self.inner.serialize_field(&WithMask {
            v: value,
            trie: self.trie,
            spec: self.spec,
            include_all: self.include_all,
        })
    }

    fn end(self) -> Result<Self::Ok, Self::Error> {
        self.inner.end()
    }
}

pub(crate) struct GateMap<'a, I> {
    inner: I,
    trie: &'a MaskTrie,
    spec: &'a Node,
    allow: bool,
}

impl<'a, I: SerializeMap> SerializeMap for GateMap<'a, I> {
    type Ok = I::Ok;
    type Error = I::Error;
    fn serialize_key<Tx: ?Sized + Serialize>(&mut self, key: &Tx) -> Result<(), Self::Error> {
        if self.allow {
            self.inner.serialize_key(key)
        } else {
            Ok(())
        }
    }
    fn serialize_value<Tx: ?Sized + Serialize>(&mut self, value: &Tx) -> Result<(), Self::Error> {
        if self.allow {
            self.inner.serialize_value(&WithMask {
                v: value,
                trie: self.trie,
                spec: self.spec,
                include_all: self.trie.terminal,
            })
        } else {
            Ok(())
        }
    }
    fn end(self) -> Result<Self::Ok, Self::Error> {
        self.inner.end()
    }
    fn serialize_entry<K: ?Sized + Serialize, V: ?Sized + Serialize>(
        &mut self,
        key: &K,
        value: &V,
    ) -> Result<(), Self::Error> {
        if !self.allow {
            return Ok(());
        }
        self.inner.serialize_entry(
            key,
            &WithMask {
                v: value,
                trie: self.trie,
                spec: self.spec,
                include_all: self.trie.terminal,
            },
        )
    }
}

pub(crate) struct GateStruct<'a, I> {
    inner: I,
    trie: &'a MaskTrie,
    fields: Option<&'a BTreeMap<String, Node>>,
    include_all: bool,
}

impl<'a, I: SerializeStruct> SerializeStruct for GateStruct<'a, I> {
    type Ok = I::Ok;
    type Error = I::Error;

    fn serialize_field<Tx: ?Sized + Serialize>(
        &mut self,
        key: &'static str,
        value: &Tx,
    ) -> Result<(), Self::Error> {
        if self.include_all {
            return self.inner.serialize_field(key, value);
        }
        let Some(map) = self.fields else {
            return Ok(());
        };
        let Some(child_spec) = map.get(key) else {
            return Ok(());
        };
        if let Some(child_trie) = self.trie.children.get(key) {
            self.inner.serialize_field(
                key,
                &WithMask {
                    v: value,
                    trie: child_trie,
                    spec: child_spec,
                    include_all: child_trie.terminal,
                },
            )
        } else if let Some(wild) = self.trie.children.get("*") {
            self.inner.serialize_field(
                key,
                &WithMask {
                    v: value,
                    trie: wild,
                    spec: child_spec,
                    include_all: wild.terminal,
                },
            )
        } else {
            Ok(())
        }
    }

    fn end(self) -> Result<Self::Ok, Self::Error> {
        self.inner.end()
    }
}

pub(crate) struct GateStructVar<'a, I> {
    inner: I,
    trie: &'a MaskTrie,
    fields: Option<&'a BTreeMap<String, Node>>,
    include_all: bool,
}
impl<'a, I: SerializeStructVariant> SerializeStructVariant for GateStructVar<'a, I> {
    type Ok = I::Ok;
    type Error = I::Error;

    fn serialize_field<Tx: ?Sized + Serialize>(
        &mut self,
        key: &'static str,
        value: &Tx,
    ) -> Result<(), Self::Error> {
        if self.include_all {
            return self.inner.serialize_field(key, value);
        }
        let Some(map) = self.fields else {
            return Ok(());
        };
        let Some(child_spec) = map.get(key) else {
            return Ok(());
        };
        if let Some(child_trie) = self.trie.children.get(key) {
            self.inner.serialize_field(
                key,
                &WithMask {
                    v: value,
                    trie: child_trie,
                    spec: child_spec,
                    include_all: child_trie.terminal,
                },
            )
        } else if let Some(wild) = self.trie.children.get("*") {
            self.inner.serialize_field(
                key,
                &WithMask {
                    v: value,
                    trie: wild,
                    spec: child_spec,
                    include_all: wild.terminal,
                },
            )
        } else {
            Ok(())
        }
    }

    fn end(self) -> Result<Self::Ok, Self::Error> {
        self.inner.end()
    }
}