heatslop / src / app.rs

use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Instant;

use winit::application::ApplicationHandler;
use winit::dpi::PhysicalSize;
use winit::event::WindowEvent;
use winit::event_loop::ActiveEventLoop;
use winit::window::{Window, WindowId};

use crate::cpu;
use crate::gpu::GpuState;

pub struct App {
    window: Option<Arc<Window>>,
    gpu: Option<GpuState>,
    running: Arc<AtomicBool>,
    primes_found: Arc<AtomicU64>,
    numbers_tested: Arc<AtomicU64>,
    cpu_threads: Vec<std::thread::JoinHandle<()>>,
    frame_count: u64,
    last_fps_time: Instant,
    fps: f64,
    last_tested_count: u64,
    tests_per_sec: f64,
}

impl App {
    pub fn new() -> Self {
        Self {
            window: None,
            gpu: None,
            running: Arc::new(AtomicBool::new(true)),
            primes_found: Arc::new(AtomicU64::new(0)),
            numbers_tested: Arc::new(AtomicU64::new(0)),
            cpu_threads: Vec::new(),
            frame_count: 0,
            last_fps_time: Instant::now(),
            fps: 0.0,
            last_tested_count: 0,
            tests_per_sec: 0.0,
        }
    }

    fn spawn_cpu_workers(&mut self) {
        let num_threads = num_cpus::get();
        log::info!("Spawning {} CPU threads for prime search", num_threads);

        for i in 0..num_threads {
            let running = self.running.clone();
            let primes_found = self.primes_found.clone();
            let numbers_tested = self.numbers_tested.clone();
            let handle = std::thread::Builder::new()
                .name(format!("prime-worker-{}", i))
                .spawn(move || {
                    cpu::prime_worker(
                        running,
                        primes_found,
                        numbers_tested,
                        i as u64,
                        num_threads as u64,
                    );
                })
                .expect("Failed to spawn CPU worker thread");
            self.cpu_threads.push(handle);
        }
    }

    fn update_title(&mut self) {
        self.frame_count += 1;
        let now = Instant::now();
        let elapsed = now.duration_since(self.last_fps_time).as_secs_f64();

        let current_tested = self.numbers_tested.load(Ordering::Relaxed);
        let current_primes = self.primes_found.load(Ordering::Relaxed);

        if elapsed >= 0.5 {
            self.fps = self.frame_count as f64 / elapsed;
            let delta = current_tested.saturating_sub(self.last_tested_count);
            self.tests_per_sec = delta as f64 / elapsed;
            self.last_tested_count = current_tested;
            self.frame_count = 0;
            self.last_fps_time = now;
        }

        let num_cores = num_cpus::get();

        let title = format!(
            "\u{1F525} HEATSLOP \u{2014} GPU: {:.0} fps | CPU: {} cores \u{00D7} primes \u{2014} {} found, {} tests/s | Stay warm!",
            self.fps,
            num_cores,
            format_number(current_primes),
            format_number(self.tests_per_sec as u64),
        );

        if let Some(window) = &self.window {
            window.set_title(&title);
        }
    }
}

fn format_number(n: u64) -> String {
    if n >= 1_000_000_000 {
        format!("{:.2}B", n as f64 / 1_000_000_000.0)
    } else if n >= 1_000_000 {
        format!("{:.2}M", n as f64 / 1_000_000.0)
    } else if n >= 1_000 {
        format!("{:.1}K", n as f64 / 1_000.0)
    } else {
        format!("{}", n)
    }
}

impl ApplicationHandler for App {
    fn resumed(&mut self, event_loop: &ActiveEventLoop) {
        if self.window.is_some() {
            return;
        }

        let window_attrs = Window::default_attributes()
            .with_title("\u{1F525} HEATSLOP \u{2014} Starting up...")
            .with_inner_size(PhysicalSize::new(1280u32, 720u32));

        let window = Arc::new(event_loop.create_window(window_attrs).unwrap());

        let gpu = pollster::block_on(GpuState::new(window.clone()));

        self.window = Some(window);
        self.gpu = Some(gpu);

        // Start CPU workers
        self.spawn_cpu_workers();
    }

    fn window_event(
        &mut self,
        event_loop: &ActiveEventLoop,
        _window_id: WindowId,
        event: WindowEvent,
    ) {
        match event {
            WindowEvent::CloseRequested => {
                log::info!("Shutting down...");
                self.running.store(false, Ordering::Relaxed);
                event_loop.exit();
            }
            WindowEvent::Resized(new_size) => {
                if let Some(gpu) = &mut self.gpu {
                    gpu.resize(new_size);
                }
            }
            WindowEvent::RedrawRequested => {
                self.update_title();

                if let Some(gpu) = &mut self.gpu {
                    match gpu.render() {
                        Ok(_) => {}
                        Err(wgpu::SurfaceError::Lost) => {
                            let size = gpu.size;
                            gpu.resize(size);
                        }
                        Err(wgpu::SurfaceError::OutOfMemory) => {
                            log::error!("Out of GPU memory!");
                            event_loop.exit();
                        }
                        Err(e) => {
                            log::warn!("Render error: {:?}", e);
                        }
                    }
                }

                // Request another frame immediately (no vsync = max GPU throughput)
                if let Some(window) = &self.window {
                    window.request_redraw();
                }
            }
            WindowEvent::KeyboardInput { event, .. } => {
                if event.physical_key
                    == winit::keyboard::PhysicalKey::Code(winit::keyboard::KeyCode::Escape)
                {
                    self.running.store(false, Ordering::Relaxed);
                    event_loop.exit();
                }
            }
            _ => {}
        }
    }
}

impl Drop for App {
    fn drop(&mut self) {
        self.running.store(false, Ordering::Relaxed);
        for handle in self.cpu_threads.drain(..) {
            let _ = handle.join();
        }
    }
}