📄 sink/src/model.rs
//! The viewer UI: virtual-camera viewfinder and steering controls.
//!
//! [`ViewerModel`] owns the whole pipeline for the UI β€” the frame source,
//! the Brush renderer, the latest frame's splats β€” and does its work
//! inside [`ViewerModel::ui`], which is called once per egui frame (by
//! the eframe shell in [`crate::app`] or the kittest harness in tests).
//! Everything is headless-testable: no windowing is involved here.

use std::time::Instant;

use anyhow::Result;
use egui::{ColorImage, TextureHandle, TextureOptions};

use render::{RenderedImage, SplatRenderer};
use rig::{Camera, FrameSource, Rig, SplatBatch};

/// The viewer's interactive state, independent of egui internals so it
/// can be asserted on directly in tests.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ViewerState {
    /// Virtual camera position between the side cameras:
    /// `0.0` = left camera, `1.0` = right camera.
    pub t: f32,
    /// Whether the frame source is being played back.
    pub playing: bool,
}

/// What the last pipeline pass achieved, for the stats readout and tests.
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct FrameStats {
    /// Frames pulled from the source so far.
    pub frame_index: u64,
    /// Splats in the latest frame's batch.
    pub splats: usize,
    /// Splats visible in the latest virtual-camera render.
    pub visible_splats: u32,
    /// Wall-clock time of the latest render, milliseconds.
    pub render_ms: f32,
    /// The latest frame's source timestamp, seconds.
    pub timestamp_secs: f64,
}

/// Drives the pipeline and builds the viewer UI.
pub struct ViewerModel {
    rig: Rig,
    source: Box<dyn FrameSource>,
    renderer: SplatRenderer,
    state: ViewerState,
    /// The latest frame's splats, re-rendered whenever `t` changes.
    current: Option<SplatBatch>,
    last_image: Option<RenderedImage>,
    texture: Option<TextureHandle>,
    stats: FrameStats,
}

impl ViewerModel {
    pub fn new(rig: Rig, source: Box<dyn FrameSource>, renderer: SplatRenderer) -> Self {
        Self {
            rig,
            source,
            renderer,
            state: ViewerState {
                t: 0.5,
                playing: true,
            },
            current: None,
            last_image: None,
            texture: None,
            stats: FrameStats::default(),
        }
    }

    /// The interactive state.
    pub fn state(&self) -> ViewerState {
        self.state
    }

    /// The latest pipeline statistics.
    pub fn stats(&self) -> FrameStats {
        self.stats
    }

    /// The latest rendered virtual-camera image.
    pub fn last_image(&self) -> Option<&RenderedImage> {
        self.last_image.as_ref()
    }

    /// Steer the virtual camera; `t` is clamped to `[0, 1]`.
    pub fn set_t(&mut self, t: f32) {
        self.state.t = t.clamp(0.0, 1.0);
    }

    /// The virtual camera the model currently renders from: a webcam
    /// (the rig's left camera's intrinsics and resolution) posed at
    /// `t` between the side cameras.
    pub fn virtual_camera(&self) -> Camera {
        self.rig.virtual_camera(
            self.state.t,
            self.rig.left.intrinsics,
            self.rig.left.resolution,
        )
    }

    fn toggle_playing(&mut self) {
        self.state.playing = !self.state.playing;
    }

    /// One pipeline pass: advance the source if playing, then (re)render
    /// the latest splats from the current virtual camera. Re-rendering
    /// while paused is what makes `t` scrubbing live.
    fn update_pipeline(&mut self) -> Result<()> {
        if self.state.playing {
            let frame = self.source.next_frame()?;
            self.stats.frame_index += 1;
            self.stats.timestamp_secs = frame.timestamp_secs;
            self.current = Some(frame.splats);
        }
        let Some(batch) = self.current.as_ref() else {
            return Ok(()); // no frame yet (e.g. paused before first pass)
        };
        let camera = self.virtual_camera();
        let start = Instant::now();
        let (image, render_stats) = self.renderer.render(batch, &camera)?;
        self.stats.render_ms = start.elapsed().as_secs_f32() * 1e3;
        self.stats.splats = render_stats.total_splats;
        self.stats.visible_splats = render_stats.visible_splats;
        self.last_image = Some(image);
        Ok(())
    }

    /// Run one pipeline pass and build the UI. Called once per egui
    /// frame with the root [`egui::Ui`] (by the eframe shell in
    /// [`crate::app`] or the kittest harness in tests).
    pub fn ui(&mut self, ui: &mut egui::Ui) {
        if let Err(error) = self.update_pipeline() {
            // Keep the last good frame on screen; the source may recover.
            log::error!("pipeline error: {error:#}");
        }

        egui::Panel::bottom("controls").show(ui, |ui| {
            ui.horizontal(|ui| {
                let play_label = if self.state.playing { "Pause" } else { "Play" };
                if ui.button(play_label).clicked() {
                    self.toggle_playing();
                }
                ui.add(
                    egui::Slider::new(&mut self.state.t, 0.0..=1.0).text("Virtual camera position"),
                );
                ui.label(format!("{:.0}%", self.state.t * 100.0));
                ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
                    ui.label(format!(
                        "frame {} Β· {} splats Β· {} visible Β· {:.1} ms",
                        self.stats.frame_index,
                        self.stats.splats,
                        self.stats.visible_splats,
                        self.stats.render_ms,
                    ));
                });
            });
        });

        egui::CentralPanel::default_margins().show(ui, |ui| {
            let Some(image) = self.last_image.as_ref() else {
                ui.label("Waiting for the first frame…");
                return;
            };
            // Brush's packed output is premultiplied over the background.
            let color_image = ColorImage::from_rgba_premultiplied(
                [image.width as usize, image.height as usize],
                &image.rgba,
            );
            if let Some(texture) = self.texture.as_mut() {
                texture.set(color_image, TextureOptions::LINEAR);
            } else {
                self.texture = Some(ui.ctx().load_texture(
                    "virtual_camera",
                    color_image,
                    TextureOptions::LINEAR,
                ));
            }
            if let Some(texture) = &self.texture {
                ui.add(
                    egui::Image::from_texture(texture)
                        .maintain_aspect_ratio(true)
                        .max_size(ui.available_size()),
                );
            }
        });
    }
}