five_five_five/src/main.rs

use std::{
    time::{
        Duration,
        Instant,
    },
};

use eframe::egui;

use tokio::sync::
{
    mpsc,
    watch,
};

mod capture;
mod controller;
mod driver;
mod hand_tracking;

use driver::
{
    Driver,
    MsgToDriver,
    MsgToGui,
};

fn main() -> Result <(), Error>
{
    let mut args = std::env::args ();
    let _exe_name = args.next ();
    match args.next ().as_deref ()
    {
        None => main_egui (),
        Some ("capture") =>
        {
            let mut capture = capture::Capture::new ().unwrap ();
            let mut buf = vec! [0u8; capture.size_image ()];
            
            for _ in 0..10
            {
                capture.wait_for_frame (&mut buf).unwrap ();
            }
            
            println! ("Measuring...");
            let start = Instant::now ();
            for i in 0..30
            {
                dbg! (i);
                capture.wait_for_frame (&mut buf).unwrap ();
            }
            let stop = Instant::now ();
            
            println! ("{} ms", (stop - start).as_millis ());
        },
        Some ("hand-tracking") => hand_tracking::main (),
        Some ("network") => main_network (),
        Some (_) => eprintln! ("Unknown subcommand"),
    }
    
    Ok (())
}



struct App {
    img: egui::ColorImage,
    texture: Option <egui::TextureHandle>,
    gui_frames: u64,
    camera_frames: u64,
    gui_fps: u64,
    camera_fps: u64,
    latency: Duration,
    
    send_to_ctl: mpsc::Sender <MsgToDriver>,
    recv_at_gui: mpsc::Receiver <MsgToGui>,
    
    requesting_frames: bool,
    debug_capture: bool,
    
    next_stat_refresh: Instant,
}

impl App {
    fn new (
        cc: &eframe::CreationContext<'_>,
        send_to_ctl: mpsc::Sender <MsgToDriver>,
        recv_at_gui: mpsc::Receiver <MsgToGui>,
        send_gui_handle: tokio::sync::oneshot::Sender <egui::Context>,
    ) -> Self 
    {
        let gui_ctx = cc.egui_ctx.clone ();
        send_gui_handle.send (gui_ctx).unwrap ();
        
        let img = egui::ColorImage::new ([1280,720], egui::Color32::TEMPORARY_COLOR);
        
        Self 
        {
            img,
            texture: None,
            gui_frames: 0,
            camera_frames: 0,
            gui_fps: 0,
            camera_fps: 0,
            latency: Duration::from_millis(0),
            send_to_ctl,
            recv_at_gui,
            requesting_frames: true,
            debug_capture: false,
            next_stat_refresh: Instant::now (),
        }
    }
}

impl eframe::App for App {
    fn update(&mut self, ctx: &egui::Context, _frame: &mut eframe::Frame) {
        egui::CentralPanel::default().show(ctx, |ui| {
            let texture = self.texture.get_or_insert_with(||
            {
                ui.ctx ().load_texture ("my_image", self.img.clone (), Default::default())
            });
            
            if self.requesting_frames
            {
                self.send_to_ctl.blocking_send (MsgToDriver::GuiNeedsRgbaFrame).unwrap ();
            }
            self.send_to_ctl.blocking_send(MsgToDriver::GuiDebugCapture (self.debug_capture)).unwrap ();
            
            while let Ok (msg) = self.recv_at_gui.try_recv()
            {
                match msg
                {
                    MsgToGui::NewRgbaFrame ((frame, _present_time)) =>
                    {
                        texture.set (egui::ColorImage::from_rgba_premultiplied([1280,720], &frame.data), Default::default ());
                        self.latency = std::cmp::max (self.latency, Instant::now () - frame.capture_time);
                        self.camera_frames += 1;
                        // println! ("     Camera");
                    },
                }
            }
            
            let texture: &egui::TextureHandle = texture;
            
            let now = Instant::now ();
            if now >= self.next_stat_refresh
            {
                self.gui_fps = self.gui_frames;
                self.camera_fps = self.camera_frames;
                self.gui_frames = 0;
                self.camera_frames = 0;
                self.latency = Duration::default ();
                self.next_stat_refresh += Duration::from_secs(1);
            }
            
            ui.heading (format! ("{0} camera FPS, latency {1} ms", self.camera_fps, self.latency.as_millis ()));
            let available = ui.available_size() - egui::Vec2::new (0.0, 20.0);
            let tex_size = egui::Vec2::new (texture.size()[0] as f32, texture.size()[1] as f32);
            let scaled_width = available.y * tex_size.x / tex_size.y;
            let scaled_height = available.x * tex_size.y / tex_size.x;
            
            let size = if scaled_width <= available.x
            {
                egui::Vec2::new (scaled_width, available.y)
            }
            else
            {
                egui::Vec2::new (available.x, scaled_height)
            };
            ui.image (texture, size);
            
            self.gui_frames += 1;
            // println! ("GUI");
            
            ui.checkbox(&mut self.requesting_frames, "Run");
            ui.checkbox(&mut self.debug_capture, "Debug capture");
        });
    }
}

fn main_egui ()
{
    let rt = tokio::runtime::Runtime::new ().unwrap ();
    
    let (send_to_ctl, recv_at_ctl) = mpsc::channel (8);
    let (send_to_gui, recv_at_gui) = mpsc::channel (8);
    let (send_gui_handle, recv_gui_handle) = tokio::sync::oneshot::channel ();
    
    let send_to_ctl_2 = send_to_ctl.clone ();
    rt.spawn (async move { 
        let gui_ctx = recv_gui_handle.await.unwrap ();
        
        let mut driver = Driver::new 
        (
            send_to_ctl_2,
            recv_at_ctl,
            gui_ctx,
            send_to_gui,
        );
        driver.run ().await 
    });
    
    // GUIs are the most needy, picky, finicky piece of any program
    // so they get spawned last, given the main thread, and babied.
    // 
    // Not to mention that `run_native` is callback heck. I'm sure it's
    // for a good reason, but still.
    
    let native_options = eframe::NativeOptions::default();
    eframe::run_native("Five Five Five", native_options, Box::new(|cc| Box::new(App::new(cc, send_to_ctl, recv_at_gui, send_gui_handle)))).unwrap ();
}

fn main_network ()
{
    use tokio::io::
    {
        AsyncReadExt,
        AsyncWriteExt,
    };
    
    enum MsgToController
    {
        AcceptedClient ((u64, mpsc::Sender <MsgToServer>)),
        DisconnectedClient (u64),
        GeneratedData (u64),
        Request ((u64, u8)),
    }
    
    struct MsgToServer (Vec <u8>);
    
    struct Client
    {
        sender: mpsc::Sender <MsgToServer>,
        pause_time: Instant,
    }
    
    let rt = tokio::runtime::Runtime::new ().unwrap ();
    
    rt.block_on (async 
    {
        let (send_to_ctl, mut recv_at_ctl) = mpsc::channel (8);
        let (send_to_gen, mut recv_at_gen) = watch::channel (Instant::now ());
        
        // Control task
        tokio::spawn (async move
        {
            let mut clients = std::collections::BTreeMap::default ();
            
            while let Some (msg) = recv_at_ctl.recv ().await
            {
                match msg
                {
                    MsgToController::AcceptedClient ((id, sender)) =>
                    {
                        let client = Client
                        {
                            sender,
                            pause_time: Instant::now (),
                        };
                        clients.insert (id, client);
                    },
                    MsgToController::DisconnectedClient(id) => 
                    {
                        clients.remove (&id);
                    },
                    MsgToController::GeneratedData (data) =>
                    {
                        let now = Instant::now ();
                        for (_id, client) in clients.iter_mut()
                        .filter(|(_id, client)| client.pause_time > now)
                        {
                            client.sender.send (MsgToServer (format! ("Data {data}\n").into_bytes())).await.unwrap ();
                        }
                    },
                    MsgToController::Request ((id, _req)) =>
                    {
                        let client = clients.get_mut (&id).unwrap ();
                        
                        client.pause_time = Instant::now () + Duration::from_secs (2);
                        client.sender.send (MsgToServer (format! (":V {id}\n").into_bytes())).await.unwrap ();
                        send_to_gen.send (client.pause_time).unwrap ();
                    },
                }
            }
        });
        
        // Bogus data generator task
        
        let send_to_ctl_2 = send_to_ctl.clone ();
        tokio::spawn (async move
        {
            let mut i = 0;
            
            while recv_at_gen.changed ().await.is_ok ()
            {
                let pause_time = *recv_at_gen.borrow_and_update ();
                while Instant::now () < pause_time
                {
                    println! ("Generating...");
                    tokio::time::sleep (Duration::from_millis(200)).await;
                    
                    send_to_ctl_2.send (MsgToController::GeneratedData (i)).await.unwrap ();
                    i += 1;
                }
            };
        });
        
        let server = tokio::net::TcpListener::bind ("127.0.0.1:9001").await.unwrap ();
        
        // Accept clients
        
        tokio::spawn (async move
        {
            let mut client_id = 0u64;
            while let Ok ((stream, _remote_addr)) = server.accept ().await
            {
                let (send_to_server, mut recv_at_server) = mpsc::channel (8);
                send_to_ctl.send (MsgToController::AcceptedClient ((client_id, send_to_server))).await.unwrap ();
                
                let (mut net_read, mut net_write) = stream.into_split ();
                
                // Read from clients
                
                let send_to_ctl_2 = send_to_ctl.clone ();
                tokio::spawn (async move
                {
                    loop 
                    {
                        let mut buf = vec! [0u8, 0u8];
                        if let Err (_) = net_read.read_exact (&mut buf).await
                        {
                            break;
                        }
                        send_to_ctl_2.send (MsgToController::Request ((client_id, buf [0]))).await.unwrap ();
                    }
                    
                    send_to_ctl_2.send (MsgToController::DisconnectedClient (client_id)).await.ok ();
                });
                
                // Write to clients
                
                let send_to_ctl_2 = send_to_ctl.clone ();
                tokio::spawn (async move
                {
                    while let Some (msg) = recv_at_server.recv ().await
                    {
                        net_write.write_all (&msg.0).await.unwrap ();
                    }
                    
                    send_to_ctl_2.send (MsgToController::DisconnectedClient (client_id)).await.ok ();
                });
                
                client_id += 1;
            }
        }).await.unwrap ();
    });
}

#[derive (Debug, thiserror::Error)]
enum Error
{
    #[error ("capture")]
    Capture (#[from] capture::Error),
}