flax 0.7.1

use core::marker::PhantomData;

use alloc::{collections::BTreeMap, format, string::String, vec::Vec};
use serde::{
    de::{self, DeserializeSeed, SeqAccess, VariantAccess, Visitor},
    Deserialize, Deserializer,
};

use crate::{
    archetype::{BatchSpawn, Storage},
    component::{ComponentDesc, ComponentValue},
    Component, Entity, EntityBuilder, World,
};

use super::{RowFields, SerializeFormat, WorldFields};

#[derive(Clone)]
struct Slot {
    /// Takes a whole column and returns a serializer for it
    deser_col: fn(
        deserializer: &mut dyn erased_serde::Deserializer,
        len: usize,
        component: ComponentDesc,
    ) -> erased_serde::Result<Storage>,
    deser_one: fn(
        deserializer: &mut dyn erased_serde::Deserializer,
        component: ComponentDesc,
        builder: &mut EntityBuilder,
    ) -> erased_serde::Result<()>,
    desc: ComponentDesc,
}

/// [ T, T, T ]
struct DeserializeStorage<'a> {
    slot: &'a Slot,
    len: usize,
}

impl<'a, 'de> DeserializeSeed<'de> for DeserializeStorage<'a> {
    type Value = Storage;

    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
    where
        D: Deserializer<'de>,
    {
        let mut deserializer = <dyn erased_serde::Deserializer>::erase(deserializer);
        let storage = (self.slot.deser_col)(&mut deserializer, self.len, self.slot.desc)
            .map_err(de::Error::custom)?;

        Ok(storage)
    }
}

#[derive(Clone, Default)]
/// Incrementally construct a [crate::serialize::DeserializeContext]
pub struct DeserializeBuilder {
    slots: BTreeMap<String, Slot>,
}

impl DeserializeBuilder {
    /// Creates a new DeserializeBuilder
    pub fn new() -> Self {
        Default::default()
    }

    /// Register a component using the component's name
    ///
    /// See [`Self::with_name`]
    pub fn with<T>(&mut self, component: Component<T>) -> &mut Self
    where
        T: ComponentValue + for<'x> Deserialize<'x>,
    {
        self.with_name(component.name(), component)
    }

    /// Register a new component to be deserialized
    pub fn with_name<T>(&mut self, key: impl Into<String>, component: Component<T>) -> &mut Self
    where
        T: ComponentValue + for<'x> Deserialize<'x>,
    {
        fn deser_col<T: ComponentValue + for<'x> Deserialize<'x>>(
            deserializer: &mut dyn erased_serde::Deserializer,
            len: usize,
            desc: ComponentDesc,
        ) -> erased_serde::Result<Storage> {
            deserializer.deserialize_seq(StorageVisitor::<T> {
                desc,
                cap: len,
                _marker: PhantomData,
            })
        }

        fn deser_one<T: ComponentValue + for<'x> Deserialize<'x>>(
            deserializer: &mut dyn erased_serde::Deserializer,
            desc: ComponentDesc,
            builder: &mut EntityBuilder,
        ) -> erased_serde::Result<()> {
            let value = T::deserialize(deserializer)?;
            builder.set(desc.downcast(), value);
            Ok(())
        }

        let key = key.into();

        self.slots.insert(
            key,
            Slot {
                deser_col: deser_col::<T>,
                deser_one: deser_one::<T>,
                desc: component.desc(),
            },
        );
        self
    }

    /// Finish constructing the deserialization context
    pub fn build(&mut self) -> DeserializeContext {
        DeserializeContext {
            slots: self.slots.clone(),
        }
    }
}

/// Describes how to deserialize the world from the described components.
pub struct DeserializeContext {
    slots: BTreeMap<String, Slot>,
}

impl DeserializeContext {
    /// Deserializes the world from the supplied deserializer.
    /// Automatically uses the row or column major format depending on the
    /// underlying data.
    pub fn deserialize<'de, D>(&self, deserializer: D) -> core::result::Result<World, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_enum("World", &["row", "col"], WorldVisitor { context: self })
    }

    fn get(&self, key: &str) -> Result<&Slot, String> {
        self.slots
            .get(key)
            .ok_or_else(|| format!("Unknown component key: {key:?}"))
    }
}

struct WorldVisitor<'a> {
    context: &'a DeserializeContext,
}

impl<'a, 'de> Visitor<'de> for WorldVisitor<'a> {
    type Value = World;

    fn expecting(&self, formatter: &mut core::fmt::Formatter) -> core::fmt::Result {
        write!(formatter, "A map like structure containing the world")
    }

    fn visit_enum<A>(self, data: A) -> Result<Self::Value, A::Error>
    where
        A: de::EnumAccess<'de>,
    {
        let (format, variant) = data.variant::<SerializeFormat>()?;
        let world = match format {
            SerializeFormat::ColumnMajor => variant.struct_variant(
                &["archetypes"],
                WorldColumnVisitor {
                    context: self.context,
                },
            )?,
            SerializeFormat::RowMajor => variant.struct_variant(
                &["entities"],
                WorldRowVisitor {
                    context: self.context,
                },
            )?,
        };
        Ok(world)
    }
}

struct DeserializeEntities<'a> {
    context: &'a DeserializeContext,
    world: &'a mut World,
}

impl<'de, 'a> DeserializeSeed<'de> for DeserializeEntities<'a> {
    type Value = ();

    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_seq(self)
    }
}

impl<'de, 'a> Visitor<'de> for DeserializeEntities<'a> {
    type Value = ();

    fn expecting(&self, formatter: &mut core::fmt::Formatter) -> core::fmt::Result {
        write!(formatter, "an entity id followed by a map of components")
    }

    fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
    where
        A: SeqAccess<'de>,
    {
        let mut builder = EntityBuilder::new();
        while let Some(id) = seq.next_element_seed(DeserializeEntity {
            context: self.context,
            builder: &mut builder,
        })? {
            // The world that is serialized into is empty
            builder.spawn_at(self.world, id).map_err(|e| {
                de::Error::custom(format!("Duplicate entities in deserialized world: {e}"))
            })?;
        }

        Ok(())
    }
}

/// (id, components)
struct DeserializeEntity<'a> {
    context: &'a DeserializeContext,
    builder: &'a mut EntityBuilder,
}

impl<'de, 'a> DeserializeSeed<'de> for DeserializeEntity<'a> {
    type Value = Entity;

    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_tuple_struct("Entity", 2, self)
    }
}

impl<'de, 'a> Visitor<'de> for DeserializeEntity<'a> {
    type Value = Entity;

    fn expecting(&self, formatter: &mut core::fmt::Formatter) -> core::fmt::Result {
        write!(formatter, "an entity id followed by a map of components")
    }

    fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
    where
        A: SeqAccess<'de>,
    {
        let id = seq
            .next_element()?
            .ok_or_else(|| de::Error::invalid_length(1, &self))?;

        seq.next_element_seed(DeserializeEntityData {
            context: self.context,
            builder: self.builder,
        })?
        .ok_or_else(|| de::Error::invalid_length(0, &self))?;

        Ok(id)
    }
}

struct DeserializeEntityData<'a> {
    context: &'a DeserializeContext,
    builder: &'a mut EntityBuilder,
}

impl<'de, 'a> DeserializeSeed<'de> for DeserializeEntityData<'a> {
    type Value = ();

    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_map(self)
    }
}

/// { component: value }
impl<'de, 'a> Visitor<'de> for DeserializeEntityData<'a> {
    type Value = ();

    fn expecting(&self, formatter: &mut core::fmt::Formatter) -> core::fmt::Result {
        write!(formatter, "a map of component values")
    }

    fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
    where
        A: de::MapAccess<'de>,
    {
        while let Some(key) = map.next_key::<&str>()? {
            let slot = self.context.get(key).map_err(de::Error::custom)?;
            map.next_value_seed(DeserializeComponent {
                slot,
                builder: self.builder,
            })?;
        }

        Ok(())
    }
}

/// A single component value
struct DeserializeComponent<'a> {
    slot: &'a Slot,
    builder: &'a mut EntityBuilder,
}

impl<'de, 'a> DeserializeSeed<'de> for DeserializeComponent<'a> {
    type Value = ();

    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
    where
        D: Deserializer<'de>,
    {
        let mut deserializer = <dyn erased_serde::Deserializer>::erase(deserializer);
        (self.slot.deser_one)(&mut deserializer, self.slot.desc, self.builder)
            .map_err(de::Error::custom)?;

        Ok(())
    }
}

/// Deserializes a list of archetypes
struct WorldRowVisitor<'a> {
    context: &'a DeserializeContext,
}

impl<'de, 'a> Visitor<'de> for WorldRowVisitor<'a> {
    type Value = World;

    fn expecting(&self, formatter: &mut core::fmt::Formatter) -> core::fmt::Result {
        write!(formatter, "a struct containing a sequence of entities")
    }

    fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
    where
        A: de::SeqAccess<'de>,
    {
        let mut world = World::new();

        seq.next_element_seed(DeserializeEntities {
            context: self.context,
            world: &mut world,
        })?
        .ok_or_else(|| de::Error::invalid_length(1, &self))?;

        Ok(world)
    }

    fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
    where
        A: de::MapAccess<'de>,
    {
        let mut world = World::new();

        while let Some(key) = map.next_key()? {
            match key {
                RowFields::Entities => map.next_value_seed(DeserializeEntities {
                    context: self.context,
                    world: &mut world,
                })?,
            }
        }

        Ok(world)
    }
}

/// Deserializes a list of archetypes
struct WorldColumnVisitor<'a> {
    context: &'a DeserializeContext,
}

impl<'de, 'a> Visitor<'de> for WorldColumnVisitor<'a> {
    type Value = World;

    fn expecting(&self, formatter: &mut core::fmt::Formatter) -> core::fmt::Result {
        write!(formatter, "a struct containing a sequence of archetypes")
    }

    fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
    where
        A: de::MapAccess<'de>,
    {
        let mut world = World::new();
        let mut has_archetypes = false;

        while let Some(key) = map.next_key()? {
            match key {
                WorldFields::Archetypes => {
                    if has_archetypes {
                        return Err(de::Error::duplicate_field("archetypes"));
                    }

                    map.next_value_seed(DeserializeArchetypes {
                        context: self.context,
                        world: &mut world,
                    })?;

                    has_archetypes = true;
                }
            }
        }

        Ok(world)
    }

    fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
    where
        A: SeqAccess<'de>,
    {
        let mut world = World::new();

        seq.next_element_seed(DeserializeArchetypes {
            context: self.context,
            world: &mut world,
        })?
        .ok_or_else(|| de::Error::invalid_length(0, &self))?;

        Ok(world)
    }
}

/// Deserializes a list of archetypes
struct DeserializeArchetypes<'a> {
    context: &'a DeserializeContext,
    world: &'a mut World,
}

impl<'a, 'de> DeserializeSeed<'de> for DeserializeArchetypes<'a> {
    type Value = ();

    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_seq(ArchetypesVisitor {
            context: self.context,
            world: self.world,
        })
    }
}

struct ArchetypesVisitor<'a> {
    context: &'a DeserializeContext,
    world: &'a mut World,
}

impl<'a, 'de> Visitor<'de> for ArchetypesVisitor<'a> {
    type Value = ();

    fn expecting(&self, formatter: &mut core::fmt::Formatter) -> core::fmt::Result {
        write!(formatter, "expected a sequence of archetypes")
    }

    fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
    where
        A: SeqAccess<'de>,
    {
        let world = self.world;
        while let Some((ids, mut batch)) = seq.next_element_seed(DeserializeArchetype {
            context: self.context,
        })? {
            world
                .spawn_batch_at(&ids, &mut batch)
                .expect("Entity ids are not duplicated");
        }

        Ok(())
    }
}

struct DeserializeArchetype<'a> {
    context: &'a DeserializeContext,
}

impl<'de, 'a> DeserializeSeed<'de> for DeserializeArchetype<'a> {
    type Value = (Vec<Entity>, BatchSpawn);

    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_tuple_struct(
            "Arch",
            2,
            ArchetypeVisitor {
                context: self.context,
            },
        )
    }
}

struct ArchetypeVisitor<'a> {
    context: &'a DeserializeContext,
}

impl<'a, 'de> Visitor<'de> for ArchetypeVisitor<'a> {
    type Value = (Vec<Entity>, BatchSpawn);

    fn expecting(&self, formatter: &mut core::fmt::Formatter) -> core::fmt::Result {
        write!(formatter, "an archetype of entities and components")
    }

    fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
    where
        A: SeqAccess<'de>,
    {
        let entities: Vec<Entity> = seq
            .next_element()?
            .ok_or_else(|| de::Error::invalid_length(0, &self))?;

        let components = seq
            .next_element_seed(DeserializeStorages {
                len: entities.len(),
                context: self.context,
            })?
            .ok_or_else(|| de::Error::invalid_length(1, &self))?;

        Ok((entities, components))
    }
}

struct DeserializeStorages<'a> {
    len: usize,
    context: &'a DeserializeContext,
}

impl<'de, 'a> DeserializeSeed<'de> for DeserializeStorages<'a> {
    type Value = BatchSpawn;

    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_map(StoragesVisitor {
            len: self.len,
            context: self.context,
        })
    }
}

struct StoragesVisitor<'a> {
    len: usize,
    context: &'a DeserializeContext,
}

impl<'de, 'a> Visitor<'de> for StoragesVisitor<'a> {
    type Value = BatchSpawn;

    fn expecting(&self, formatter: &mut core::fmt::Formatter) -> core::fmt::Result {
        write!(formatter, "a map of component values")
    }

    fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
    where
        A: de::MapAccess<'de>,
    {
        let mut batch = BatchSpawn::new(self.len);
        while let Some(key) = map.next_key::<&'de str>()? {
            let slot = self.context.get(key).map_err(de::Error::custom)?;

            let storage = map.next_value_seed(DeserializeStorage {
                slot,
                len: self.len,
            })?;

            batch.append(storage).map_err(de::Error::custom)?;
        }

        Ok(batch)
    }
}

/// Visit a single column of component values
struct StorageVisitor<T: ComponentValue> {
    desc: ComponentDesc,
    cap: usize,
    _marker: PhantomData<T>,
}

impl<'de, T: ComponentValue + de::Deserialize<'de>> Visitor<'de> for StorageVisitor<T> {
    type Value = Storage;

    fn expecting(&self, formatter: &mut core::fmt::Formatter) -> core::fmt::Result {
        write!(formatter, "A sequence of component values of the same type")
    }

    fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
    where
        A: SeqAccess<'de>,
    {
        let mut storage = Storage::with_capacity(self.desc, self.cap);

        while let Some(item) = seq.next_element::<T>()? {
            unsafe { storage.push(item) }
        }

        Ok(storage)
    }
}