//! Live DINO feature view in a desktop window.
//!
//! The same pipeline the headset runs — capture, encode, PCA colour,
//! display — pointed at a monitor instead of a passthrough camera. Two
//! reasons it exists:
//!
//! * It is the demo you can actually show someone. Point it at a video, a
//!   photo, a game, and watch objects resolve into stable colour regions.
//! * It de-risks the Quest camera. The capture loop, the downscale, and the
//!   "new frame arrives asynchronously while the window redraws" structure
//!   are all the same; the passthrough camera becomes another
//!   `FrameSource` and nothing around it changes.
//!
//! ```text
//! cargo run --release --features capture --example viewer -- [model.safetensors] [monitor]
//! ```
//!
//! Without a checkpoint it runs on synthetic weights, which still shows the
//! scene's structure but colours it arbitrarily.

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

use blade_graphics as gpu;
use dinovision::dinov3::Config;
use dinovision::inference::{self, Weights, Worker};
use dinovision::render::GridView;
use dinovision::source::FrameSource;

fn surface_config(size: winit::dpi::PhysicalSize<u32>) -> gpu::SurfaceConfig {
    gpu::SurfaceConfig {
        size: gpu::Extent {
            width: size.width.max(1),
            height: size.height.max(1),
            depth: 1,
        },
        usage: gpu::TextureUsage::TARGET,
        display_sync: gpu::DisplaySync::Recent,
        ..Default::default()
    }
}

struct Viewer {
    context: Arc<gpu::Context>,
    surface: gpu::Surface,
    encoder: gpu::CommandEncoder,
    view: GridView,
    worker: Worker,
    source: Box<dyn FrameSource>,
    config: Config,
    last_shown: u64,
    frames: u64,
    encodes: u64,
    last_report: Instant,
    window: winit::window::Window,
}

impl Viewer {
    fn redraw(&mut self) {
        // Feed the encoder whenever it is free. Unlike the headset there is
        // no compositor to starve here, so it runs flat out.
        if self.worker.is_ready()
            && let Some(rgb) = self.source.next_frame()
        {
            let patches = dinovision::preprocess::patches_from_rgb8(rgb, &self.config);
            self.worker.submit(patches);
        }

        let generation = self.worker.generation();
        if generation != self.last_shown
            && let Some(grid) = self.worker.latest()
        {
            self.view.upload(&self.context, grid.patch_rgb());
            self.last_shown = generation;
            self.encodes += 1;
        }

        let frame = self.surface.acquire_frame();
        self.encoder.start();
        self.encoder.init_texture(frame.texture());
        {
            let mut pass = self.encoder.render("view", gpu::RenderTargetSet {
                colors: &[gpu::RenderTarget {
                    view: frame.texture_view(),
                    init_op: gpu::InitOp::DontCare,
                    finish_op: gpu::FinishOp::Store,
                }],
                depth_stencil: None,
            });
            self.view.draw(&mut pass, 0);
        }
        self.encoder.present(frame);
        let _sp = self.context.submit(&mut self.encoder);

        self.frames += 1;
        if self.last_report.elapsed() >= Duration::from_secs(3) {
            let secs = self.last_report.elapsed().as_secs_f64();
            let latency = self
                .worker
                .latest()
                .map(|g| g.latency_ms)
                .unwrap_or(f64::NAN);
            self.window.set_title(&format!(
                "dinovision — {:.0} fps display, {:.1} Hz inference ({:.0} ms)",
                self.frames as f64 / secs,
                self.encodes as f64 / secs,
                latency
            ));
            self.frames = 0;
            self.encodes = 0;
            self.last_report = Instant::now();
        }
    }
}

#[derive(Default)]
struct App {
    viewer: Option<Viewer>,
    weights: Option<std::path::PathBuf>,
    decoder: Option<std::path::PathBuf>,
    monitor: usize,
}

impl winit::application::ApplicationHandler for App {
    fn resumed(&mut self, event_loop: &winit::event_loop::ActiveEventLoop) {
        if self.viewer.is_some() {
            return;
        }
        let window = event_loop
            .create_window(
                winit::window::Window::default_attributes()
                    .with_title("dinovision — starting…")
                    .with_inner_size(winit::dpi::LogicalSize::new(720, 720)),
            )
            .expect("failed to create window");

        let context = Arc::new(unsafe {
            gpu::Context::init(gpu::ContextDesc {
                presentation: true,
                validation: false,
                ..Default::default()
            })
            .expect("failed to initialize GPU context")
        });

        let surface = context
            .create_surface_configured(&window, surface_config(window.inner_size()))
            .expect("failed to create surface");

        let config = Config::vits16().at_resolution(224);
        // A reconstruction fills the whole image, so the display grid is the
        // image size; PCA colouring only has one value per patch.
        let display = match self.decoder.clone() {
            Some(path) => {
                log::info!("decoder: {}", path.display());
                inference::Display::Reconstruction(path)
            }
            None => inference::Display::PcaColour,
        };
        let view_grid = match &display {
            inference::Display::Reconstruction(_) => config.image_size,
            inference::Display::PcaColour => config.grid(),
        };
        let view = GridView::new(&context, surface.info().format, view_grid);

        let weights = match self.weights.clone() {
            Some(p) => {
                log::info!("weights: {}", p.display());
                Weights::SafeTensors(p)
            }
            None => {
                log::warn!("no checkpoint given — synthetic weights, colours are arbitrary");
                Weights::Synthetic
            }
        };
        // One submission: nothing else is competing for this GPU, so the
        // chunking that matters on a headset would only cost overhead.
        let worker = inference::spawn(Arc::clone(&context), config.clone(), weights, None, 1, display);

        let source: Box<dyn FrameSource> = match dinovision::source::ScreenCapture::new(
            config.image_size,
            self.monitor,
        ) {
            Ok(s) => Box::new(s),
            Err(e) => {
                log::warn!("screen capture unavailable ({e}); falling back to the test pattern");
                Box::new(dinovision::source::TestPattern::new(config.image_size))
            }
        };

        let encoder = context.create_command_encoder(gpu::CommandEncoderDesc {
            name: "viewer",
            buffer_count: 2,
            manual_barriers: false,
        });

        self.viewer = Some(Viewer {
            context,
            surface,
            encoder,
            view,
            worker,
            source,
            config,
            last_shown: 0,
            frames: 0,
            encodes: 0,
            last_report: Instant::now(),
            window,
        });
    }

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

    fn window_event(
        &mut self,
        event_loop: &winit::event_loop::ActiveEventLoop,
        _id: winit::window::WindowId,
        event: winit::event::WindowEvent,
    ) {
        let Some(viewer) = self.viewer.as_mut() else {
            return;
        };
        match event {
            winit::event::WindowEvent::CloseRequested => event_loop.exit(),
            winit::event::WindowEvent::KeyboardInput {
                event:
                    winit::event::KeyEvent {
                        physical_key: winit::keyboard::PhysicalKey::Code(code),
                        state: winit::event::ElementState::Pressed,
                        ..
                    },
                ..
            } => match code {
                winit::keyboard::KeyCode::Escape => event_loop.exit(),
                // Refit the colour basis on demand: point the capture at
                // something new and the old principal components will be a
                // poor fit for it.
                winit::keyboard::KeyCode::KeyR => {
                    log::info!("refitting colour basis");
                    viewer.worker.request_refit();
                }
                _ => {}
            },
            winit::event::WindowEvent::Resized(size) => {
                viewer
                    .context
                    .reconfigure_surface(&mut viewer.surface, surface_config(size));
            }
            winit::event::WindowEvent::RedrawRequested => viewer.redraw(),
            _ => {}
        }
    }
}

fn main() {
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();

    let mut args = std::env::args().skip(1);
    let weights = args.next().map(std::path::PathBuf::from).filter(|p| {
        let ok = p.exists();
        if !ok {
            log::warn!("{} does not exist; using synthetic weights", p.display());
        }
        ok
    });
    // A decoder file switches the view from PCA colour to a real RGB
    // reconstruction.
    let decoder = args.next().map(std::path::PathBuf::from).filter(|p| p.exists());
    let monitor = args.next().and_then(|s| s.parse().ok()).unwrap_or(0);

    println!("R refits the colour basis · Esc quits");

    let event_loop = winit::event_loop::EventLoop::new().expect("failed to create event loop");
    event_loop.set_control_flow(winit::event_loop::ControlFlow::Poll);
    event_loop
        .run_app(&mut App {
            viewer: None,
            weights,
            decoder,
            monitor,
        })
        .expect("event loop failed");
}