//src/engine/scene.rs

pub use crate::frontend::props::{DrawableProps, SharedProps};

use glam::{Mat4, Quat, Vec2, Vec3, vec2, vec3};
use std::collections::HashMap;
use std::path::Path;
use std::path::PathBuf;

use crate::engine::camera::Camera;
use crate::engine::timeline::Timeline;
use crate::frontend::layout::{Anchor, Bounds, Direction, anchor_for_direction, opposite_anchor};
use crate::frontend::updater::UpdaterManager;
use crate::frontend::{DirtyFlags, IntoTattva, TattvaId, tattva_trait::TattvaTrait};
use crate::resource::texture::TextureImage;

#[derive(Debug, Clone, Default)]
pub struct ScreenshotCapture {
    pub times: Vec<f32>,
    pub names: Option<Vec<Option<PathBuf>>>,
}

impl ScreenshotCapture {
    pub fn new<I>(times: I) -> Self
    where
        I: IntoIterator<Item = f32>,
    {
        let mut times: Vec<f32> = times.into_iter().collect();
        times.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
        Self { times, names: None }
    }

    pub fn with_names<I, P>(mut self, names: I) -> Self
    where
        I: IntoIterator<Item = Option<P>>,
        P: Into<PathBuf>,
    {
        self.names = Some(names.into_iter().map(|name| name.map(Into::into)).collect());
        self
    }
}

#[derive(Debug, Clone)]
pub struct GifCapture {
    pub name: String,
    pub times: Vec<f32>,
}

impl GifCapture {
    pub fn new<I>(name: impl Into<String>, times: I) -> Self
    where
        I: IntoIterator<Item = f32>,
    {
        let mut times: Vec<f32> = times.into_iter().collect();
        times.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
        Self {
            name: name.into(),
            times,
        }
    }
}

/// The Scene represents the authoritative Frontend state.
pub struct Scene {
    /// Authoritative Tattvas (Source of Truth)
    /// We use TattvaTrait to allow different types (Circle, Latex, etc.) in one list.
    pub tattvas: HashMap<TattvaId, Box<dyn TattvaTrait>>,

    /// Time & Animation
    pub scene_time: f32,
    pub timelines: HashMap<String, Timeline>,
    pub screenshot_captures: Vec<ScreenshotCapture>,
    pub gif_captures: Vec<GifCapture>,

    /// Updaters - callbacks that run every frame
    pub updaters: UpdaterManager,

    /// Global State
    pub camera: Camera,
    pub global_model: Mat4,

    /// Identity bookkeeping
    next_tattva_id: TattvaId,
    removed_tattva_ids: Vec<TattvaId>,
}

impl Scene {
    pub fn new() -> Self {
        Self {
            tattvas: HashMap::new(),
            scene_time: 0.0,
            timelines: HashMap::new(),
            screenshot_captures: Vec::new(),
            gif_captures: Vec::new(),
            updaters: UpdaterManager::new(),
            camera: Camera::default(),
            global_model: Mat4::IDENTITY,
            next_tattva_id: 1,
            removed_tattva_ids: Vec::new(),
        }
    }

    /// Adds a Tattva to the scene and returns its stable ID.
    pub fn add<T: TattvaTrait + 'static>(&mut self, mut tattva: T) -> TattvaId {
        let id = self.next_tattva_id;
        self.next_tattva_id += 1;

        // Tattvas need to know their own ID for animation targeting
        // (Assuming a .set_id() helper on the trait)
        tattva.set_id(id);

        self.tattvas.insert(id, Box::new(tattva));
        id
    }

    /// Adds a concrete state object to the scene at a given position.
    pub fn add_tattva<T>(&mut self, state: T, position: Vec3) -> TattvaId
    where
        T: crate::projection::Project + crate::frontend::layout::Bounded + Send + Sync + 'static,
    {
        let tattva = state.into_tattva();
        {
            let mut props = DrawableProps::write(&tattva.props);
            props.position = position;
        }
        self.add(tattva)
    }

    pub fn add_vector_latex(
        &mut self,
        equation: crate::frontend::collection::math::equation::VectorLatexEquation,
    ) -> crate::frontend::collection::math::equation::VectorEquationHandle {
        equation.add_to_scene(self)
    }

    pub fn add_vector_typst(
        &mut self,
        equation: crate::frontend::collection::math::equation::VectorTypstEquation,
    ) -> crate::frontend::collection::math::equation::VectorEquationHandle {
        equation.add_to_scene(self)
    }

    pub fn add_vector_formula_latex(
        &mut self,
        equation: crate::frontend::collection::math::equation::VectorLatexEquation,
    ) -> crate::frontend::collection::math::equation::VectorEquationHandle {
        self.add_vector_latex(equation)
    }

    pub fn add_vector_formula_typst(
        &mut self,
        equation: crate::frontend::collection::math::equation::VectorTypstEquation,
    ) -> crate::frontend::collection::math::equation::VectorEquationHandle {
        self.add_vector_typst(equation)
    }

    /// Hides a tattva completely from rendering until another animation or mutation reveals it.
    pub fn hide_tattva(&mut self, id: TattvaId) {
        self.hide(id);
    }

    pub fn add_screenshot_capture(&mut self, capture: ScreenshotCapture) {
        self.screenshot_captures.push(capture);
    }

    pub fn capture_screenshots<I>(&mut self, times: I)
    where
        I: IntoIterator<Item = f32>,
    {
        self.add_screenshot_capture(ScreenshotCapture::new(times));
    }

    pub fn capture_screenshots_named<I, P>(&mut self, names: I)
    where
        I: IntoIterator<Item = (f32, Option<P>)>,
        P: Into<PathBuf>,
    {
        let entries: Vec<(f32, Option<PathBuf>)> = names
            .into_iter()
            .map(|(time, name)| (time, name.map(Into::into)))
            .collect();
        let times = entries.iter().map(|(time, _)| *time).collect::<Vec<_>>();
        let names = entries
            .into_iter()
            .map(|(_, name)| name)
            .collect::<Vec<_>>();
        self.add_screenshot_capture(ScreenshotCapture::new(times).with_names(names));
    }

    pub fn add_gif_capture(&mut self, capture: GifCapture) {
        self.gif_captures.push(capture);
    }

    pub fn capture_gif<I>(&mut self, name: impl Into<String>, times: I)
    where
        I: IntoIterator<Item = f32>,
    {
        self.add_gif_capture(GifCapture::new(name, times));
    }

    /// Installs or replaces a named timeline in the scene.
    pub fn set_timeline(&mut self, name: impl Into<String>, timeline: Timeline) {
        self.timelines.insert(name.into(), timeline);
    }

    /// Preferred authored API for the common single-timeline case.
    pub fn play(&mut self, timeline: Timeline) {
        self.set_timeline("main", timeline);
    }

    pub fn play_named(&mut self, name: impl Into<String>, timeline: Timeline) {
        self.set_timeline(name, timeline);
    }

    pub fn add_textured_surface<F>(
        &mut self,
        u_range: (f32, f32),
        v_range: (f32, f32),
        f: F,
        texture: TextureImage,
        position: Vec3,
    ) -> TattvaId
    where
        F: Fn(f32, f32) -> Vec3 + Send + Sync + 'static,
    {
        let surface =
            crate::frontend::collection::graph::parametric_surface::ParametricSurface::new(
                u_range, v_range, f,
            )
            .with_texture(texture);
        self.add_tattva(surface, position)
    }

    pub fn add_textured_surface_with_path(
        &mut self,
        surface: crate::frontend::collection::graph::parametric_surface::ParametricSurface,
        texture_path: impl AsRef<Path>,
        position: Vec3,
    ) -> anyhow::Result<TattvaId> {
        let surface = surface.with_texture_path(texture_path)?;
        Ok(self.add_tattva(surface, position))
    }

    /// Retrieves a Tattva for inspection or mutation.
    pub fn get_tattva_any_mut(&mut self, id: TattvaId) -> Option<&mut (dyn TattvaTrait + '_)> {
        match self.tattvas.get_mut(&id) {
            Some(b) => Some(b.as_mut()),
            None => None,
        }
    }

    pub fn get_tattva_any(&self, id: TattvaId) -> Option<&(dyn TattvaTrait + '_)> {
        match self.tattvas.get(&id) {
            Some(b) => Some(b.as_ref()),
            None => None,
        }
    }

    pub fn get_tattva_typed<T: 'static>(
        &self,
        id: TattvaId,
    ) -> Option<&crate::frontend::Tattva<T>> {
        self.get_tattva_any(id)?
            .as_any()
            .downcast_ref::<crate::frontend::Tattva<T>>()
    }

    pub fn get_tattva_typed_mut<T: 'static>(
        &mut self,
        id: TattvaId,
    ) -> Option<&mut crate::frontend::Tattva<T>> {
        self.get_tattva_any_mut(id)?
            .as_any_mut()
            .downcast_mut::<crate::frontend::Tattva<T>>()
    }

    /// Primary update loop for the frontend.
    pub fn update(&mut self, dt: f32) {
        self.scene_time += dt;

        // 1. Tick all timelines (this mutates Tattva props/state)
        // We temporarily move them out to avoid borrow checker conflicts with 'self'
        let mut timelines = std::mem::take(&mut self.timelines);
        for (_, tl) in timelines.iter_mut() {
            tl.update(self.scene_time, self);
        }
        self.timelines = timelines;

        // 2. Update traced paths
        self.update_traced_paths();

        // 3. Run all updaters
        let updaters = std::mem::take(&mut self.updaters);
        updaters.update_all(self, dt);
        self.updaters = updaters;
    }

    /// Update all traced paths in the scene
    fn update_traced_paths(&mut self) {
        use crate::frontend::Tattva;
        use crate::frontend::collection::utility::TracedPath;

        // Collect IDs of traced paths and their tracked objects
        let mut traced_paths: Vec<(TattvaId, TattvaId)> = Vec::new();

        for (id, tattva) in self.tattvas.iter() {
            // TracedPath is wrapped in Tattva<TracedPath>
            if let Some(wrapped) = tattva.as_any().downcast_ref::<Tattva<TracedPath>>() {
                traced_paths.push((*id, wrapped.state.tracked_object_id));
            }
        }

        // Update each traced path
        for (traced_path_id, tracked_object_id) in traced_paths {
            // Get the tracked object's position and rotation
            if let Some(tracked_obj) = self.get_tattva_any(tracked_object_id) {
                let props = DrawableProps::read(tracked_obj.props());
                let obj_pos = props.position;
                let obj_rot = props.rotation;
                drop(props);

                // Get the traced path and record the point
                if let Some(tattva) = self.tattvas.get_mut(&traced_path_id) {
                    if let Some(wrapped) = tattva.as_any_mut().downcast_mut::<Tattva<TracedPath>>()
                    {
                        // Compute the traced point position using the point function
                        let traced_point = (wrapped.state.point_fn)(obj_pos, obj_rot);
                        let old_count = wrapped.state.path_points.len();
                        wrapped.state.record_point(traced_point);
                        let new_count = wrapped.state.path_points.len();

                        // If a new point was added, mark the tattva as dirty
                        if new_count > old_count {
                            wrapped.mark_dirty(DirtyFlags::GEOMETRY);
                        }
                    }
                }
            }
        }
    }

    /// Returns an iterator over all Tattvas for the Sync Boundary to process.
    pub fn tattvas_iter_mut(
        &mut self,
    ) -> impl Iterator<Item = (&TattvaId, &mut Box<dyn TattvaTrait>)> {
        self.tattvas.iter_mut()
    }

    pub fn local_bounds(&self, id: TattvaId) -> Option<Bounds> {
        self.get_tattva_any(id).map(|t| t.local_bounds())
    }

    pub fn world_bounds(&self, id: TattvaId) -> Option<Bounds> {
        let tattva = self.get_tattva_any(id)?;
        let local = tattva.local_bounds();
        let props = DrawableProps::read(tattva.props());
        let scaled = local.scale(vec2(props.scale.x, props.scale.y));
        Some(scaled.translate(vec2(props.position.x, props.position.y)))
    }

    pub fn anchor_position(&self, id: TattvaId, anchor: Anchor) -> Option<Vec2> {
        self.world_bounds(id).map(|b| b.anchor(anchor))
    }

    pub fn set_position_2d(&mut self, id: TattvaId, position: Vec2) {
        if let Some(tattva) = self.get_tattva_any_mut(id) {
            let mut props = DrawableProps::write(tattva.props());
            props.position = vec3(position.x, position.y, props.position.z);
            drop(props);
            tattva.mark_dirty(DirtyFlags::TRANSFORM);
        }
    }

    pub fn set_position_3d(&mut self, id: TattvaId, position: Vec3) {
        if let Some(tattva) = self.get_tattva_any_mut(id) {
            let mut props = DrawableProps::write(tattva.props());
            props.position = position;
            drop(props);
            tattva.mark_dirty(DirtyFlags::TRANSFORM);
        }
    }

    pub fn set_scale(&mut self, id: TattvaId, scale: Vec3) {
        if let Some(tattva) = self.get_tattva_any_mut(id) {
            let mut props = DrawableProps::write(tattva.props());
            props.scale = scale;
            drop(props);
            tattva.mark_dirty(DirtyFlags::TRANSFORM);
        }
    }

    pub fn set_rotation(&mut self, id: TattvaId, rotation: Quat) {
        if let Some(tattva) = self.get_tattva_any_mut(id) {
            let mut props = DrawableProps::write(tattva.props());
            props.rotation = rotation;
            drop(props);
            tattva.mark_dirty(DirtyFlags::TRANSFORM);
        }
    }

    pub fn set_opacity(&mut self, id: TattvaId, opacity: f32) {
        if let Some(tattva) = self.get_tattva_any_mut(id) {
            let mut props = DrawableProps::write(tattva.props());
            props.opacity = opacity.clamp(0.0, 1.0);
            props.visible = props.opacity > 0.001;
            drop(props);
            tattva.mark_dirty(DirtyFlags::STYLE | DirtyFlags::VISIBILITY);
        }
    }

    pub fn show(&mut self, id: TattvaId) {
        self.set_opacity(id, 1.0);
    }

    pub fn hide(&mut self, id: TattvaId) {
        self.set_opacity(id, 0.0);
    }

    pub fn align_to(&mut self, moving: TattvaId, target: TattvaId, anchor: Anchor) {
        let Some(moving_anchor) = self.anchor_position(moving, anchor) else {
            return;
        };
        let Some(target_anchor) = self.anchor_position(target, anchor) else {
            return;
        };
        let delta = target_anchor - moving_anchor;
        if let Some(bounds) = self.world_bounds(moving) {
            let current_pos = bounds.center();
            let new_pos = match anchor {
                Anchor::Up | Anchor::Down => glam::vec2(current_pos.x, current_pos.y + delta.y),
                Anchor::Left | Anchor::Right => glam::vec2(current_pos.x + delta.x, current_pos.y),
                _ => current_pos + delta,
            };
            self.set_position_2d(moving, new_pos);
        }
    }

    pub fn next_to(
        &mut self,
        moving: TattvaId,
        target: TattvaId,
        direction: Direction,
        padding: f32,
    ) {
        let Some(target_bounds) = self.world_bounds(target) else {
            return;
        };
        let Some(local_bounds) = self.local_bounds(moving) else {
            return;
        };

        let moving_anchor = opposite_anchor(direction);
        let target_anchor = anchor_for_direction(direction);
        let target_point = target_bounds.anchor(target_anchor);
        let local_anchor = local_bounds.anchor(moving_anchor);

        let offset = match direction {
            Direction::Up => vec2(0.0, padding),
            Direction::Down => vec2(0.0, -padding),
            Direction::Left => vec2(-padding, 0.0),
            Direction::Right => vec2(padding, 0.0),
        };

        self.set_position_2d(moving, target_point + offset - local_anchor);
    }

    pub fn to_edge(&mut self, id: TattvaId, direction: Direction, margin: f32) {
        let Some(local_bounds) = self.local_bounds(id) else {
            return;
        };
        let frame = self.frame_bounds();
        let target_anchor = anchor_for_direction(direction);
        let moving_anchor = opposite_anchor(direction);
        let edge_point = frame.anchor(target_anchor);
        let margin_offset = match direction {
            Direction::Up => vec2(0.0, -margin),
            Direction::Down => vec2(0.0, margin),
            Direction::Left => vec2(margin, 0.0),
            Direction::Right => vec2(-margin, 0.0),
        };
        self.set_position_2d(
            id,
            edge_point + margin_offset - local_bounds.anchor(moving_anchor),
        );
    }

    pub fn frame_bounds(&self) -> Bounds {
        let view_width = self.camera.view_width();
        let half_width = view_width / 2.0;
        let half_height = half_width / crate::engine::camera::ASPECT_RATIO;
        Bounds::new(
            vec2(-half_width, -half_height),
            vec2(half_width, half_height),
        )
    }

    pub fn clear(&mut self) {
        self.tattvas.clear();
        self.timelines.clear();
        self.screenshot_captures.clear();
        self.gif_captures.clear();
        self.updaters.clear();
        self.scene_time = 0.0;
        self.next_tattva_id = 1;
        self.removed_tattva_ids.clear();
        self.camera = Camera::default();
    }

    // camera
    pub fn camera(&self) -> &Camera {
        &self.camera
    }

    pub fn camera_mut(&mut self) -> &mut Camera {
        &mut self.camera
    }

    /// Replaces the implementation of an existing Tattva.
    /// This is used for shape morphing where we swap types (e.g., Circle -> Path).
    pub fn replace_tattva(&mut self, id: TattvaId, tattva: Box<dyn TattvaTrait>) {
        self.tattvas.insert(id, tattva);
    }

    /// Removes a tattva from the scene and records its ID so the backend can
    /// despawn any cached render entities tied to it.
    pub fn remove_tattva(&mut self, id: TattvaId) -> Option<Box<dyn TattvaTrait>> {
        let removed = self.tattvas.remove(&id);
        if removed.is_some() {
            self.removed_tattva_ids.push(id);
        }
        removed
    }

    pub fn take_removed_tattva_ids(&mut self) -> Vec<TattvaId> {
        std::mem::take(&mut self.removed_tattva_ids)
    }

    /// Add an updater callback for a tattva
    /// The callback will be called every frame with (scene, tattva_id, dt)
    /// Returns an index that can be used to remove the updater later
    pub fn add_updater<F>(&mut self, tattva_id: TattvaId, callback: F) -> usize
    where
        F: Fn(&mut Scene, TattvaId, f32) + Send + Sync + 'static,
    {
        self.updaters.add_updater(tattva_id, callback)
    }

    /// Remove an updater by its index
    pub fn remove_updater(&mut self, index: usize) {
        self.updaters.remove_updater(index);
    }

    /// Remove all updaters for a specific tattva
    pub fn remove_updaters_for(&mut self, tattva_id: TattvaId) {
        self.updaters.remove_updaters_for_tattva(tattva_id);
    }
}

impl Default for Scene {
    fn default() -> Self {
        Self::new()
    }
}
