//! Registry-aware full-scene capture, for play mode (and future scene files).
//!
//! [`Scene::to_ron`](super::Scene::to_ron) records only the node-baked
//! `Node`/`Transform`/hierarchy. A [`SceneSnapshot`] additionally captures
//! **every reflected component** on each entity through the
//! [`TypeRegistry`], so a scene mutated while *playing* (physics moving bodies,
//! scripts spawning or editing entities) can be restored **bit-for-bit** when
//! play stops — avoiding Unity's classic "edited in play mode, lost it" footgun.
//!
//! Why a separate type from [`SceneData`](super::serialize): the plain RON form
//! is registry-free (it can round-trip without knowing any component types),
//! whereas a snapshot needs the registry to enumerate and serialize arbitrary
//! components. Keeping the two apart means the cheap path stays cheap.
//!
//! Fidelity is bounded by what is *registered*, plus the engine's intrinsic
//! components: every reflected type in the [`TypeRegistry`] is captured, as are
//! the built-in non-reflected components (`Node`/`Transform`, plus the inspector-
//! hidden [`Tags`] and [`DisabledComponents`]). A *module's* component that is
//! neither registered nor one of those is invisible to capture — modules that
//! want play-mode survival register their components, which they do anyway to be
//! editable. Within that set, capture → restore → capture is stable.

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

use hecs::Entity;
use serde::{Deserialize, Serialize};

use super::{DisabledComponents, Node, Scene, SceneError};
use crate::layer::Tags;
use crate::math::Transform;
use crate::reflect::TypeRegistry;

/// Components captured explicitly via [`Scene::spawn`] on restore, so they are
/// excluded from the per-node component map to avoid storing them twice.
/// (`Layer` is *not* here: it is auto-attached on spawn but carries authored
/// data, so it round-trips through the component map like any other component.)
const SPAWN_BAKED: [&str; 2] = ["Node", "Transform"];

/// One entity in a flattened [`SceneSnapshot`]. `children` holds indices into
/// the surrounding [`SceneSnapshot::nodes`] list (entity handles are not stable
/// across a capture/restore, so positions stand in for references).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
struct SnapshotNode {
    name: String,
    enabled: bool,
    transform: Transform,
    children: Vec<usize>,
    /// Every *other* registered component on the node, `type_name` → RON.
    /// A `BTreeMap` so the serialized form is order-stable.
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    components: BTreeMap<String, String>,
    /// The entity's gameplay [`Tags`], if any. Captured directly (not via the
    /// registry) because `Tags` is an engine-intrinsic component edited through
    /// the inspector's dedicated Groups UI, not registered as a generic
    /// reflected type — without this, Play→Stop would wipe group membership.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    tags: Option<Tags>,
    /// The entity's [`DisabledComponents`] set, if any. Captured directly for
    /// the same reason as [`tags`](Self::tags): it's hidden engine metadata, not
    /// a reflected authored component.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    disabled: Option<DisabledComponents>,
}

/// A complete, restorable capture of a [`Scene`]: hierarchy plus every reflected
/// component on each node. See the [module docs](self).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SceneSnapshot {
    nodes: Vec<SnapshotNode>,
    roots: Vec<usize>,
}

impl SceneSnapshot {
    /// Captures `scene` in full, serializing every component the `registry`
    /// knows about on each entity. Built from a deterministic pre-order walk of
    /// the hierarchy, so two captures of equal scenes compare equal.
    pub fn capture(scene: &Scene, registry: &TypeRegistry) -> Self {
        let mut index: HashMap<Entity, usize> = HashMap::with_capacity(scene.len());
        let mut order: Vec<Entity> = Vec::with_capacity(scene.len());
        for &root in scene.roots() {
            assign_indices(scene, root, &mut index, &mut order);
        }

        let nodes = order
            .iter()
            .map(|&entity| {
                // Pull the node-baked fields first, then drop the borrow before
                // the reflection walk reads the same world.
                let (name, enabled) = {
                    let node = scene
                        .get::<Node>(entity)
                        .expect("entity in hierarchy must have a Node");
                    (node.name.clone(), node.enabled)
                };
                let transform = scene
                    .local_transform(entity)
                    .expect("entity in hierarchy must have a Transform");

                let mut components = BTreeMap::new();
                for type_name in registry.components_on(scene.world(), entity) {
                    if SPAWN_BAKED.contains(&type_name) {
                        continue;
                    }
                    if let Ok(ron) = registry.get_ron(scene.world(), entity, type_name) {
                        components.insert(type_name.to_string(), ron);
                    }
                }

                SnapshotNode {
                    name,
                    enabled,
                    transform,
                    children: scene.children(entity).iter().map(|c| index[c]).collect(),
                    components,
                    tags: scene.get::<Tags>(entity).map(|t| (*t).clone()),
                    disabled: scene
                        .get::<DisabledComponents>(entity)
                        .map(|d| (*d).clone()),
                }
            })
            .collect();

        let roots = scene.roots().iter().map(|r| index[r]).collect();
        SceneSnapshot { nodes, roots }
    }

    /// Rebuilds a fresh [`Scene`] from this snapshot. Spawns every node (which
    /// auto-attaches the node-baked `Node`/`Transform`/`Layer`), re-links the
    /// hierarchy, then applies each captured component over the defaults.
    ///
    /// Entity handles in the new scene differ from the captured ones — callers
    /// holding an [`Entity`] (e.g. an editor selection) must drop or re-resolve
    /// it after a restore.
    ///
    /// # Errors
    /// [`Deserialize`](SceneError::Deserialize) if an index reference is out of
    /// range or a component's stored RON no longer parses for its type.
    pub fn restore(&self, registry: &TypeRegistry) -> Result<Scene, SceneError> {
        let mut scene = Scene::new();
        let n = self.nodes.len();

        // Spawn every entity first (as a root) so all indices resolve before
        // wiring parent/child links.
        let entities: Vec<Entity> = self
            .nodes
            .iter()
            .map(|rec| {
                scene.spawn(
                    Node {
                        name: rec.name.clone(),
                        enabled: rec.enabled,
                    },
                    rec.transform,
                )
            })
            .collect();

        // Re-link the hierarchy.
        for (i, rec) in self.nodes.iter().enumerate() {
            for &child_idx in &rec.children {
                let child = *entities.get(child_idx).ok_or_else(|| {
                    SceneError::Deserialize(format!(
                        "child index {child_idx} out of range (have {n} nodes)"
                    ))
                })?;
                scene
                    .set_parent(child, Some(entities[i]))
                    .map_err(|e| SceneError::Deserialize(e.to_string()))?;
            }
        }

        // Apply captured components over the spawn defaults.
        for (i, rec) in self.nodes.iter().enumerate() {
            for (type_name, ron) in &rec.components {
                registry
                    .set_ron(scene.world_mut(), entities[i], type_name, ron)
                    .map_err(|e| SceneError::Deserialize(e.to_string()))?;
            }
            // Reinstate the engine-intrinsic, non-reflected components.
            if let Some(tags) = &rec.tags {
                let _ = scene.world_mut().insert_one(entities[i], tags.clone());
            }
            if let Some(disabled) = &rec.disabled {
                let _ = scene.world_mut().insert_one(entities[i], disabled.clone());
            }
        }

        Ok(scene)
    }

    /// Serializes the snapshot to a pretty-printed RON string.
    ///
    /// # Errors
    /// [`Serialize`](SceneError::Serialize) if encoding fails.
    pub fn to_ron(&self) -> Result<String, SceneError> {
        ron::ser::to_string_pretty(self, ron::ser::PrettyConfig::default())
            .map_err(|e| SceneError::Serialize(e.to_string()))
    }

    /// Reconstructs a snapshot from a string produced by [`to_ron`](Self::to_ron).
    ///
    /// # Errors
    /// [`Deserialize`](SceneError::Deserialize) if the text is not valid.
    pub fn from_ron(ron: &str) -> Result<Self, SceneError> {
        ron::from_str(ron).map_err(|e| SceneError::Deserialize(e.to_string()))
    }
}

/// Pre-order index assignment: mirrors the plain-RON walk so snapshot and
/// `to_ron` order entities identically.
fn assign_indices(
    scene: &Scene,
    entity: Entity,
    index: &mut HashMap<Entity, usize>,
    order: &mut Vec<Entity>,
) {
    index.insert(entity, order.len());
    order.push(entity);
    for &child in scene.children(entity) {
        assign_indices(scene, child, index, order);
    }
}

impl Scene {
    /// Captures this scene in full (hierarchy + every reflected component) into
    /// a restorable [`SceneSnapshot`]. See that type for why it differs from
    /// [`to_ron`](Self::to_ron).
    pub fn snapshot(&self, registry: &TypeRegistry) -> SceneSnapshot {
        SceneSnapshot::capture(self, registry)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::math::Vec3;
    use crate::render::{MeshRenderer, PrimitiveShape};

    /// A registry seeded like the editor's: node-baked types plus a couple of
    /// modular components, so snapshots exercise the component map.
    fn registry() -> TypeRegistry {
        let mut r = TypeRegistry::new();
        r.register_reflected::<Transform>("Transform");
        r.register_reflected::<Node>("Node");
        r.register_reflected::<crate::layer::Layer>("Layer");
        r.register_reflected::<MeshRenderer>("MeshRenderer");
        r
    }

    fn sample_scene() -> Scene {
        let mut scene = Scene::new();
        let parent = scene.spawn("parent", Transform::IDENTITY);
        let child = scene.spawn(
            "child",
            Transform::from_translation(Vec3::new(1.0, 2.0, 3.0)),
        );
        scene.set_parent(child, Some(parent)).unwrap();
        scene
            .world_mut()
            .insert_one(
                child,
                MeshRenderer {
                    shape: PrimitiveShape::Sphere,
                    ..MeshRenderer::default()
                },
            )
            .unwrap();
        scene
    }

    #[test]
    fn capture_restore_round_trips_components_and_hierarchy() {
        let reg = registry();
        let scene = sample_scene();
        let snap = scene.snapshot(&reg);

        let restored = snap.restore(&reg).unwrap();
        // Hierarchy (plain RON) matches.
        assert_eq!(scene.to_ron().unwrap(), restored.to_ron().unwrap());
        // And the full snapshot (incl. components) matches.
        assert_eq!(snap, restored.snapshot(&reg));
    }

    #[test]
    fn restore_reinstates_a_modular_component() {
        let reg = registry();
        let scene = sample_scene();
        let snap = scene.snapshot(&reg);
        let restored = snap.restore(&reg).unwrap();

        // The child's MeshRenderer (Sphere) survived the round-trip.
        let child = restored
            .entities()
            .find(|&e| {
                restored
                    .get::<Node>(e)
                    .map(|n| n.name == "child")
                    .unwrap_or(false)
            })
            .unwrap();
        let mesh = restored.get::<MeshRenderer>(child).unwrap();
        assert_eq!(mesh.shape, PrimitiveShape::Sphere);
    }

    #[test]
    fn stop_reverts_a_play_mode_mutation_bit_for_bit() {
        // The play-mode contract: snapshot, mutate (as a tick would), restore,
        // and the scene returns to its captured form.
        let reg = registry();
        let mut scene = sample_scene();
        let snap = scene.snapshot(&reg);

        // Mutate every transform, as a physics/script tick might.
        let entities: Vec<_> = scene.entities().collect();
        for e in entities {
            let t = scene.local_transform(e).unwrap();
            scene.set_local_transform(e, Transform::from_translation(t.translation + Vec3::X));
        }
        assert_ne!(snap, scene.snapshot(&reg));

        let reverted = snap.restore(&reg).unwrap();
        assert_eq!(snap, reverted.snapshot(&reg));
    }

    #[test]
    fn captures_intrinsic_tags_and_disabled() {
        use crate::layer::Tags;
        use crate::scene::DisabledComponents;

        let reg = registry();
        let mut scene = sample_scene();
        let parent = scene
            .entities()
            .find(|&e| {
                scene
                    .get::<Node>(e)
                    .map(|n| n.name == "parent")
                    .unwrap_or(false)
            })
            .unwrap();
        scene
            .world_mut()
            .insert_one(parent, Tags::single("Enemy"))
            .unwrap();
        let mut dc = DisabledComponents::new();
        dc.set_disabled("MeshRenderer", true);
        scene.world_mut().insert_one(parent, dc).unwrap();

        let snap = scene.snapshot(&reg);
        let restored = snap.restore(&reg).unwrap();
        // Both engine-intrinsic, non-reflected components survive the round-trip
        // even though neither is in the registry.
        assert_eq!(snap, restored.snapshot(&reg));
        let rparent = restored
            .entities()
            .find(|&e| {
                restored
                    .get::<Node>(e)
                    .map(|n| n.name == "parent")
                    .unwrap_or(false)
            })
            .unwrap();
        assert!(restored.get::<Tags>(rparent).unwrap().contains("Enemy"));
        assert!(restored
            .get::<DisabledComponents>(rparent)
            .unwrap()
            .is_disabled("MeshRenderer"));
    }

    #[test]
    fn ron_round_trips() {
        let reg = registry();
        let snap = sample_scene().snapshot(&reg);
        let ron = snap.to_ron().unwrap();
        assert_eq!(snap, SceneSnapshot::from_ron(&ron).unwrap());
    }
}
