threecrate-visualization 0.8.0

//! Interactive 3D viewer with UI controls
//! 
//! This module provides a simplified interactive viewer for 3D data

use std::sync::Arc;
use winit::{
    application::ApplicationHandler,
    event::{WindowEvent, ElementState, MouseButton},
    event_loop::{EventLoop, ActiveEventLoop},
    window::{Window, WindowId},
    keyboard::Key,
    dpi::PhysicalPosition,
};

use threecrate_core::{PointCloud, TriangleMesh, Result, Point3f, ColoredPoint3f, Error};
use threecrate_gpu::{
    PointCloudRenderer, RenderConfig, PointVertex,
    MeshRenderer, MeshRenderConfig, ShadingMode, PbrMaterial, MeshLightingParams, mesh_to_gpu_mesh,
};
use threecrate_algorithms::{ICPResult, PlaneSegmentationResult};
use crate::camera::Camera;

use nalgebra::{Vector3, Point3};

/// Types of data that can be displayed
#[derive(Debug, Clone)]
pub enum ViewData {
    Empty,
    PointCloud(PointCloud<Point3f>),
    ColoredPointCloud(PointCloud<ColoredPoint3f>),
    Mesh(TriangleMesh),
}

/// Camera control modes
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum CameraMode {
    Orbit,
    Pan,
    Zoom,
}

/// Pipeline processing type
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum PipelineType {
    Cpu,
    Gpu,
}

/// ICP algorithm parameters
#[derive(Debug, Clone)]
pub struct ICPParams {
    pub max_iterations: usize,
    pub convergence_threshold: f32,
    pub max_correspondence_distance: f32,
}

impl Default for ICPParams {
    fn default() -> Self {
        Self {
            max_iterations: 50,
            convergence_threshold: 0.001,
            max_correspondence_distance: 1.0,
        }
    }
}

/// RANSAC algorithm parameters
#[derive(Debug, Clone)]
pub struct RANSACParams {
    pub max_iterations: usize,
    pub distance_threshold: f32,
}

impl Default for RANSACParams {
    fn default() -> Self {
        Self {
            max_iterations: 1000,
            distance_threshold: 0.1,
        }
    }
}

/// UI state for all panels and controls (kept for future use)
#[derive(Debug)]
pub struct UIState {
    pub render_panel_open: bool,
    pub algorithm_panel_open: bool,
    pub camera_panel_open: bool,
    pub stats_panel_open: bool,
    pub icp_params: ICPParams,
    pub ransac_params: RANSACParams,
    pub source_cloud: Option<PointCloud<Point3f>>,
    pub target_cloud: Option<PointCloud<Point3f>>,
    pub icp_result: Option<ICPResult>,
    pub ransac_result: Option<PlaneSegmentationResult>,
}

impl Default for UIState {
    fn default() -> Self {
        Self {
            render_panel_open: false,
            algorithm_panel_open: false,
            camera_panel_open: false,
            stats_panel_open: false,
            icp_params: ICPParams::default(),
            ransac_params: RANSACParams::default(),
            source_cloud: None,
            target_cloud: None,
            icp_result: None,
            ransac_result: None,
        }
    }
}

/// Interactive 3D viewer with comprehensive UI controls
pub struct InteractiveViewer {
    current_data: ViewData,
    camera: Camera,
    camera_mode: CameraMode,
    last_mouse_pos: Option<PhysicalPosition<f64>>,
    mouse_pressed: bool,
    right_mouse_pressed: bool,
    debug_frame_count: usize,
    vertices_dirty: bool,
    /// Active shading mode for mesh rendering
    pub shading_mode: ShadingMode,
    /// Per-mesh PBR material properties
    pub material: PbrMaterial,
    /// Lighting parameters (ambient, light position/intensity/color, tone-mapping)
    pub lighting_params: MeshLightingParams,
    /// Set to true when lighting_params has changed and must be uploaded to the GPU
    lighting_dirty: bool,
}

impl InteractiveViewer {
    /// Create a new interactive viewer
    pub fn new() -> Result<Self> {
        let camera = Camera::new(
            Point3::new(5.0, 5.0, 5.0),
            Point3::new(0.0, 0.0, 0.0),
            Vector3::new(0.0, 1.0, 0.0),
            45.0,
            1.0,
            0.1,
            100.0,
        );

        Ok(Self {
            current_data: ViewData::Empty,
            camera,
            camera_mode: CameraMode::Orbit,
            last_mouse_pos: None,
            mouse_pressed: false,
            right_mouse_pressed: false,
            debug_frame_count: 0,
            vertices_dirty: true,
            shading_mode: ShadingMode::Flat,
            material: PbrMaterial::default(),
            lighting_params: MeshLightingParams::default(),
            lighting_dirty: false,
        })
    }

    /// Set point cloud data
    pub fn set_point_cloud(&mut self, cloud: &PointCloud<Point3f>) {
        self.current_data = ViewData::PointCloud(cloud.clone());
        self.vertices_dirty = true;
        println!("Set point cloud with {} points", cloud.len());
    }

    /// Set colored point cloud data
    pub fn set_colored_point_cloud(&mut self, cloud: &PointCloud<ColoredPoint3f>) {
        self.current_data = ViewData::ColoredPointCloud(cloud.clone());
        self.vertices_dirty = true;
        println!("Set colored point cloud with {} points", cloud.len());
    }

    /// Set mesh data
    pub fn set_mesh(&mut self, mesh: &TriangleMesh) {
        self.current_data = ViewData::Mesh(mesh.clone());
        self.vertices_dirty = true;
        println!("Set mesh with {} vertices and {} faces", mesh.vertices.len(), mesh.faces.len());
    }

    /// Set the shading mode used when rendering meshes.
    pub fn set_shading_mode(&mut self, mode: ShadingMode) {
        self.shading_mode = mode;
        println!("Shading mode: {:?}", mode);
    }

    /// Set the PBR material applied to the rendered mesh.
    pub fn set_material(&mut self, material: PbrMaterial) {
        self.material = material;
    }

    /// Set full lighting parameters (ambient, light position/intensity/colour, gamma, exposure).
    pub fn set_lighting_params(&mut self, params: MeshLightingParams) {
        self.lighting_params = params;
        self.lighting_dirty = true;
    }

    /// Run the interactive viewer
    pub fn run(self) -> Result<()> {
        println!("Starting threecrate Interactive Viewer...");

        // Create event loop
        let event_loop = EventLoop::new().map_err(|e| Error::Io(std::io::Error::new(std::io::ErrorKind::Other, format!("Failed to create event loop: {}", e))))?;

        // Create application handler
        let mut app = ViewerApp {
            viewer: self,
            window: None,
            point_renderer: None,
            mesh_renderer: None,
            cached_vertices: Vec::new(),
            screenshot_pending: false,
        };

        // Run the event loop
        event_loop.run_app(&mut app).map_err(|e| Error::Io(std::io::Error::new(std::io::ErrorKind::Other, format!("Event loop error: {}", e))))?;

        Ok(())
    }
}

impl Default for InteractiveViewer {
    fn default() -> Self {
        Self::new().expect("Failed to create InteractiveViewer")
    }
}

/// Application handler for the viewer
struct ViewerApp {
    viewer: InteractiveViewer,
    window: Option<Arc<Window>>,
    point_renderer: Option<PointCloudRenderer<'static>>,
    mesh_renderer: Option<MeshRenderer<'static>>,
    cached_vertices: Vec<PointVertex>,
    screenshot_pending: bool,
}

impl ApplicationHandler for ViewerApp {
    fn resumed(&mut self, event_loop: &ActiveEventLoop) {
        if self.window.is_none() {
            println!("Creating window and initializing renderers...");

            // Create window with attributes
            let window_attrs = Window::default_attributes()
                .with_title("threecrate Interactive Viewer")
                .with_inner_size(winit::dpi::LogicalSize::new(1200.0, 800.0));

            let window = match event_loop.create_window(window_attrs) {
                Ok(w) => Arc::new(w),
                Err(e) => {
                    eprintln!("Failed to create window: {}", e);
                    event_loop.exit();
                    return;
                }
            };

            // Update camera aspect ratio
            let size = window.inner_size();
            self.viewer.camera.aspect_ratio = size.width as f32 / size.height as f32;

            // Leak the Arc to get a 'static reference
            // This is safe because the window will live for the duration of the program
            let window_ref: &'static Window = unsafe {
                std::mem::transmute::<&Window, &'static Window>(window.as_ref())
            };

            // Initialize renderers using the static reference
            let pc_config = RenderConfig::default();
            let point_renderer = match pollster::block_on(PointCloudRenderer::new(window_ref, pc_config)) {
                Ok(r) => r,
                Err(e) => {
                    eprintln!("Failed to create point cloud renderer: {}", e);
                    event_loop.exit();
                    return;
                }
            };

            let mesh_config = MeshRenderConfig::default();
            let mesh_renderer = match pollster::block_on(MeshRenderer::new(window_ref, mesh_config)) {
                Ok(r) => r,
                Err(e) => {
                    eprintln!("Failed to create mesh renderer: {}", e);
                    event_loop.exit();
                    return;
                }
            };

            self.window = Some(window);
            self.point_renderer = Some(point_renderer);
            self.mesh_renderer = Some(mesh_renderer);

            println!("Viewer initialized successfully. Window should now be visible.");
            println!();
            println!("Controls:");
            println!("  O / P / Z  — Orbit / Pan / Zoom camera mode");
            println!("  R          — Reset camera");
            println!("  M          — Toggle Flat / PBR shading");
            println!("  S          — Save screenshot (screenshot_<timestamp>.png)");
            println!("  [ / ]      — Decrease / Increase ambient light strength");
            println!("  - / =      — Decrease / Increase light intensity");
            println!();
        }
    }

    fn window_event(&mut self, event_loop: &ActiveEventLoop, _id: WindowId, event: WindowEvent) {
        let Some(window) = &self.window else { return; };
        let Some(point_renderer) = &mut self.point_renderer else { return; };
        let Some(mesh_renderer) = &mut self.mesh_renderer else { return; };

        match event {
            WindowEvent::CloseRequested => {
                event_loop.exit();
            }
            WindowEvent::Resized(new_size) => {
                point_renderer.resize(new_size);
                mesh_renderer.resize(new_size);
                self.viewer.camera.aspect_ratio = new_size.width as f32 / new_size.height as f32;
            }
            WindowEvent::MouseInput { state, button, .. } => {
                match button {
                    MouseButton::Left => {
                        self.viewer.mouse_pressed = state == ElementState::Pressed;
                    }
                    MouseButton::Right => {
                        self.viewer.right_mouse_pressed = state == ElementState::Pressed;
                    }
                    _ => {}
                }
            }
            WindowEvent::CursorMoved { position, .. } => {
                if let Some(last_pos) = self.viewer.last_mouse_pos {
                    let delta_x = position.x - last_pos.x;
                    let delta_y = position.y - last_pos.y;

                    if self.viewer.mouse_pressed {
                        match self.viewer.camera_mode {
                            CameraMode::Orbit => {
                                self.viewer.camera.orbit(delta_x as f32 * 0.01, delta_y as f32 * 0.01);
                            }
                            CameraMode::Pan => {
                                self.viewer.camera.pan(delta_x as f32 * 0.01, delta_y as f32 * 0.01);
                            }
                            _ => {}
                        }
                    }
                }
                self.viewer.last_mouse_pos = Some(position);
            }
            WindowEvent::MouseWheel { delta, .. } => {
                let scroll_delta = match delta {
                    winit::event::MouseScrollDelta::LineDelta(_, y) => y,
                    winit::event::MouseScrollDelta::PixelDelta(pos) => pos.y as f32 / 100.0,
                };
                self.viewer.camera.zoom(scroll_delta * 0.1);
            }
            WindowEvent::KeyboardInput { event, .. } => {
                if event.state == ElementState::Pressed {
                    match &event.logical_key {
                        Key::Character(c) => {
                            match c.as_str() {
                                "o" | "O" => {
                                    self.viewer.camera_mode = CameraMode::Orbit;
                                    println!("Switched to Orbit mode");
                                }
                                "p" | "P" => {
                                    self.viewer.camera_mode = CameraMode::Pan;
                                    println!("Switched to Pan mode");
                                }
                                "z" | "Z" => {
                                    self.viewer.camera_mode = CameraMode::Zoom;
                                    println!("Switched to Zoom mode");
                                }
                                "r" | "R" => {
                                    self.viewer.camera.reset();
                                    println!("Reset camera");
                                }
                                // Toggle between Flat and PBR shading (M key)
                                "m" | "M" => {
                                    self.viewer.shading_mode = match self.viewer.shading_mode {
                                        ShadingMode::Flat => ShadingMode::Pbr,
                                        ShadingMode::Pbr => ShadingMode::Flat,
                                    };
                                    println!("Shading mode: {:?}", self.viewer.shading_mode);
                                }
                                // Screenshot (S key)
                                "s" | "S" => {
                                    self.screenshot_pending = true;
                                    println!("Screenshot requested…");
                                }
                                // Ambient strength  [ decrease  ] increase
                                "[" => {
                                    self.viewer.lighting_params.ambient_strength =
                                        (self.viewer.lighting_params.ambient_strength - 0.01).max(0.0);
                                    self.viewer.lighting_dirty = true;
                                    println!("Ambient strength: {:.3}", self.viewer.lighting_params.ambient_strength);
                                }
                                "]" => {
                                    self.viewer.lighting_params.ambient_strength =
                                        (self.viewer.lighting_params.ambient_strength + 0.01).min(1.0);
                                    self.viewer.lighting_dirty = true;
                                    println!("Ambient strength: {:.3}", self.viewer.lighting_params.ambient_strength);
                                }
                                // Light intensity  - decrease  = increase
                                "-" => {
                                    self.viewer.lighting_params.light_intensity =
                                        (self.viewer.lighting_params.light_intensity - 0.1).max(0.0);
                                    self.viewer.lighting_dirty = true;
                                    println!("Light intensity: {:.2}", self.viewer.lighting_params.light_intensity);
                                }
                                "=" => {
                                    self.viewer.lighting_params.light_intensity =
                                        (self.viewer.lighting_params.light_intensity + 0.1).min(10.0);
                                    self.viewer.lighting_dirty = true;
                                    println!("Light intensity: {:.2}", self.viewer.lighting_params.light_intensity);
                                }
                                _ => {}
                            }
                        }
                        _ => {}
                    }
                }
            }
            WindowEvent::RedrawRequested => {
                // Sync lighting params to GPU when changed
                if self.viewer.lighting_dirty {
                    mesh_renderer.update_lighting(self.viewer.lighting_params);
                    self.viewer.lighting_dirty = false;
                }

                // Update camera matrices
                let view_matrix = self.viewer.camera.view_matrix();
                let proj_matrix = self.viewer.camera.projection_matrix();
                let camera_pos = self.viewer.camera.position.coords;
                point_renderer.update_camera(view_matrix, proj_matrix, camera_pos);
                mesh_renderer.update_camera(view_matrix, proj_matrix, camera_pos);

                // Rebuild quad vertices only when data has changed
                if self.viewer.vertices_dirty {
                    self.cached_vertices = match &self.viewer.current_data {
                        ViewData::PointCloud(cloud) => {
                            let mut vertices = Vec::with_capacity(cloud.len() * 6);
                            for point in cloud.iter() {
                                let size = 0.02;
                                let pos = [point.x, point.y, point.z];
                                let color = [1.0, 1.0, 1.0];
                                let normal = [0.0, 0.0, 1.0];

                                let v1 = PointVertex::from_point(&Point3f::new(pos[0] - size, pos[1] - size, pos[2]), color, 16.0, normal);
                                let v2 = PointVertex::from_point(&Point3f::new(pos[0] + size, pos[1] - size, pos[2]), color, 16.0, normal);
                                let v3 = PointVertex::from_point(&Point3f::new(pos[0] + size, pos[1] + size, pos[2]), color, 16.0, normal);
                                let v4 = PointVertex::from_point(&Point3f::new(pos[0] - size, pos[1] + size, pos[2]), color, 16.0, normal);

                                vertices.push(v1);
                                vertices.push(v2);
                                vertices.push(v3);
                                vertices.push(v1);
                                vertices.push(v3);
                                vertices.push(v4);
                            }
                            vertices
                        }
                        ViewData::ColoredPointCloud(cloud) => {
                            let mut vertices = Vec::with_capacity(cloud.len() * 6);
                            for point in cloud.iter() {
                                let size = 0.02;
                                let pos = [point.position.x, point.position.y, point.position.z];
                                let color = [
                                    point.color[0] as f32 / 255.0,
                                    point.color[1] as f32 / 255.0,
                                    point.color[2] as f32 / 255.0,
                                ];
                                let normal = [0.0, 0.0, 1.0];

                                let v1 = PointVertex::from_point(&Point3f::new(pos[0] - size, pos[1] - size, pos[2]), color, 16.0, normal);
                                let v2 = PointVertex::from_point(&Point3f::new(pos[0] + size, pos[1] - size, pos[2]), color, 16.0, normal);
                                let v3 = PointVertex::from_point(&Point3f::new(pos[0] + size, pos[1] + size, pos[2]), color, 16.0, normal);
                                let v4 = PointVertex::from_point(&Point3f::new(pos[0] - size, pos[1] + size, pos[2]), color, 16.0, normal);

                                vertices.push(v1);
                                vertices.push(v2);
                                vertices.push(v3);
                                vertices.push(v1);
                                vertices.push(v3);
                                vertices.push(v4);
                            }
                            vertices
                        }
                        ViewData::Mesh(_) | ViewData::Empty => vec![],
                    };
                    self.viewer.vertices_dirty = false;
                }

                // Debug: Print vertex count periodically
                if !self.cached_vertices.is_empty() {
                    if self.viewer.debug_frame_count % 60 == 0 {
                        println!("Rendering {} vertices", self.cached_vertices.len());
                    }
                }
                self.viewer.debug_frame_count += 1;

                match &self.viewer.current_data {
                    ViewData::PointCloud(_) | ViewData::ColoredPointCloud(_) => {
                        if self.screenshot_pending {
                            self.screenshot_pending = false;
                            println!("Screenshot not yet supported for point clouds.");
                        }
                        if !self.cached_vertices.is_empty() {
                            if let Err(e) = point_renderer.render(&self.cached_vertices) {
                                eprintln!("Render error: {}", e);
                            }
                        }
                    }
                    ViewData::Mesh(mesh) => {
                        if !mesh.vertices.is_empty() && !mesh.faces.is_empty() {
                            let indices: Vec<u32> = mesh
                                .faces
                                .iter()
                                .flat_map(|f| [f[0] as u32, f[1] as u32, f[2] as u32])
                                .collect();

                            let normals_opt = mesh.normals.as_ref().map(|n| n.as_slice());

                            let colors_f32: Option<Vec<[f32; 3]>> = mesh.colors.as_ref().map(|cols| {
                                cols.iter()
                                    .map(|c| [c[0] as f32 / 255.0, c[1] as f32 / 255.0, c[2] as f32 / 255.0])
                                    .collect()
                            });
                            let colors_opt = colors_f32.as_ref().map(|c| c.as_slice());

                            // Use the viewer's material and shading mode
                            let gpu_mesh = mesh_to_gpu_mesh(
                                &mesh.vertices,
                                &indices,
                                normals_opt,
                                colors_opt,
                                Some(self.viewer.material),
                            );

                            if let Err(e) = mesh_renderer.render(&gpu_mesh, self.viewer.shading_mode) {
                                eprintln!("Mesh render error: {}", e);
                            }

                            // Handle screenshot request
                            if self.screenshot_pending {
                                self.screenshot_pending = false;
                                match mesh_renderer.render_to_texture(&gpu_mesh, self.viewer.shading_mode) {
                                    Ok((pixels, format, w, h)) => {
                                        save_screenshot(&pixels, format, w, h);
                                    }
                                    Err(e) => eprintln!("Screenshot render error: {}", e),
                                }
                            }
                        }
                    }
                    ViewData::Empty => {}
                }

                // Request next frame
                window.request_redraw();
            }
            _ => {}
        }
    }

    fn about_to_wait(&mut self, _event_loop: &ActiveEventLoop) {
        if let Some(window) = &self.window {
            window.request_redraw();
        }
    }
}

/// Encode raw GPU pixel bytes as a PNG file.
///
/// Handles both RGBA and BGRA surface formats transparently.
fn save_screenshot(pixels: &[u8], format: wgpu::TextureFormat, width: u32, height: u32) {
    use wgpu::TextureFormat;

    // Convert BGRA → RGBA if the surface uses a BGRA format (common on macOS/Metal)
    let rgba: Vec<u8> = match format {
        TextureFormat::Bgra8Unorm | TextureFormat::Bgra8UnormSrgb => {
            let mut buf = Vec::with_capacity(pixels.len());
            for chunk in pixels.chunks_exact(4) {
                buf.push(chunk[2]); // R ← B
                buf.push(chunk[1]); // G
                buf.push(chunk[0]); // B ← R
                buf.push(chunk[3]); // A
            }
            buf
        }
        _ => pixels.to_vec(),
    };

    let timestamp = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0);
    let path = format!("screenshot_{}.png", timestamp);

    match image::RgbaImage::from_raw(width, height, rgba) {
        Some(img) => match img.save(&path) {
            Ok(()) => println!("Screenshot saved: {}", path),
            Err(e) => eprintln!("Failed to save screenshot: {}", e),
        },
        None => eprintln!("Failed to build image buffer for screenshot"),
    }
}