bevy_simple_screenshot 0.1.2

//! Background disk I/O and image encoding for saving screenshots.

use bevy::prelude::*;
use bevy::render::render_resource::TextureFormat;
use chrono::{DateTime, Local};
use crossbeam_channel::{Receiver, Sender};
use image::DynamicImage;
use std::io::Write;
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;
use uuid::Uuid;

use crate::buffer::CapturedScreenshot;
use crate::burn_in::apply_burn_in;
use crate::config::{ImageFormat, ScreenshotConfig};

/// Resource for sending screenshots to the save thread.
#[derive(Resource)]
pub struct SaveTaskSender(pub Sender<CapturedScreenshot>);

/// Resource that holds the save thread handle for graceful shutdown.
#[derive(Resource)]
pub struct SaveThreadHandle {
    handle: Option<std::thread::JoinHandle<()>>,
    shutdown_signal: Arc<AtomicBool>,
}

impl SaveThreadHandle {
    /// Signal the save thread to finish and wait for completion.
    pub fn shutdown(&mut self) {
        self.shutdown_signal.store(true, Ordering::SeqCst);
        if let Some(handle) = self.handle.take() {
            // Wait up to 5 seconds for graceful shutdown
            let _ = handle.join();
        }
    }
}

impl Drop for SaveThreadHandle {
    fn drop(&mut self) {
        self.shutdown();
    }
}

/// Spawn the background I/O thread for saving screenshots.
pub fn spawn_save_thread(
    receiver: Receiver<CapturedScreenshot>,
    config: ScreenshotConfig,
) -> SaveThreadHandle {
    let shutdown_signal = Arc::new(AtomicBool::new(false));
    let shutdown_clone = shutdown_signal.clone();

    let handle = std::thread::spawn(move || {
        save_loop(receiver, config, shutdown_clone);
    });

    SaveThreadHandle {
        handle: Some(handle),
        shutdown_signal,
    }
}

/// Clean up orphaned temporary files from previous runs that were interrupted.
/// These files have the pattern `.tmp_<uuid>.<ext>` in the output directory.
fn cleanup_orphaned_temp_files(output_dir: &Path) {
    let Ok(entries) = std::fs::read_dir(output_dir) else {
        return;
    };

    let mut cleaned = 0;
    for entry in entries.flatten() {
        let path = entry.path();
        if let Some(filename) = path.file_name().and_then(|n| n.to_str()) {
            if filename.starts_with(".tmp_") {
                if let Err(e) = std::fs::remove_file(&path) {
                    bevy::log::warn!("Failed to clean up orphaned temp file {:?}: {}", path, e);
                } else {
                    cleaned += 1;
                }
            }
        }
    }

    if cleaned > 0 {
        bevy::log::info!(
            "Cleaned up {} orphaned temporary screenshot file(s) from previous run",
            cleaned
        );
    }
}

fn save_loop(
    receiver: Receiver<CapturedScreenshot>,
    config: ScreenshotConfig,
    shutdown_signal: Arc<AtomicBool>,
) {
    // Ensure output directory exists
    if let Err(e) = std::fs::create_dir_all(&config.output_dir) {
        bevy::log::error!("Failed to create screenshot directory: {}", e);
        return;
    }

    // Clean up any orphaned temp files from previous interrupted runs
    cleanup_orphaned_temp_files(&config.output_dir);

    loop {
        // Check for shutdown signal with a timeout
        match receiver.recv_timeout(Duration::from_millis(100)) {
            Ok(screenshot) => {
                if let Err(e) = save_screenshot(&screenshot, &config) {
                    bevy::log::error!("Failed to save screenshot: {}", e);
                }
            }
            Err(crossbeam_channel::RecvTimeoutError::Timeout) => {
                // Check if we should shut down
                if shutdown_signal.load(Ordering::SeqCst) {
                    // Drain any remaining screenshots before exiting
                    let mut remaining = 0;
                    while let Ok(screenshot) = receiver.try_recv() {
                        if let Err(e) = save_screenshot(&screenshot, &config) {
                            bevy::log::error!("Failed to save screenshot: {}", e);
                        }
                        remaining += 1;
                    }
                    if remaining > 0 {
                        bevy::log::info!(
                            "Screenshot save thread: saved {} remaining screenshot(s) before shutdown",
                            remaining
                        );
                    }
                    bevy::log::info!("Screenshot save thread: graceful shutdown complete");
                    break;
                }
            }
            Err(crossbeam_channel::RecvTimeoutError::Disconnected) => {
                // Channel closed, drain remaining and exit
                let mut remaining = 0;
                while let Ok(screenshot) = receiver.try_recv() {
                    if let Err(e) = save_screenshot(&screenshot, &config) {
                        bevy::log::error!("Failed to save screenshot: {}", e);
                    }
                    remaining += 1;
                }
                if remaining > 0 {
                    bevy::log::info!(
                        "Screenshot save thread: saved {} remaining screenshot(s) before exit",
                        remaining
                    );
                }
                break;
            }
        }
    }
}

fn save_screenshot(
    screenshot: &CapturedScreenshot,
    config: &ScreenshotConfig,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    // Validate image data before saving
    let width = screenshot.image.width();
    let height = screenshot.image.height();

    if width == 0 || height == 0 {
        return Err(format!(
            "Screenshot has invalid dimensions ({}x{}). This may indicate the screenshot was \
             captured before the render system was ready or during shutdown.",
            width, height
        )
        .into());
    }

    let data = screenshot.image.data.as_ref().ok_or(
        "Screenshot has no image data. This may indicate the screenshot was \
         captured before the render system was ready or during shutdown.",
    )?;

    if data.is_empty() {
        return Err(
            "Screenshot has empty image data. This may indicate the screenshot was \
             captured before the render system was ready or during shutdown."
                .into(),
        );
    }

    // Validate data size matches expected dimensions
    let expected_size = (width * height * 4) as usize; // RGBA = 4 bytes per pixel
    let actual_size = data.len();
    if actual_size != expected_size {
        return Err(format!(
            "Screenshot data size mismatch: expected {} bytes for {}x{} RGBA, got {} bytes. \
             Image may be corrupted.",
            expected_size, width, height, actual_size
        )
        .into());
    }

    // Get format settings for this key
    let format = config
        .keys
        .get(&screenshot.key)
        .and_then(|k| k.format)
        .unwrap_or(config.format);

    let jpeg_quality = config
        .keys
        .get(&screenshot.key)
        .and_then(|k| k.jpeg_quality)
        .unwrap_or(config.jpeg_quality);

    // Generate filename
    let filename = generate_filename(
        &config.filename_pattern,
        &screenshot.key,
        &screenshot.description,
        screenshot.timestamp,
        format,
    );

    let path = config.output_dir.join(&filename);

    // Convert Bevy Image to DynamicImage
    let mut dynamic_image = bevy_image_to_dynamic(&screenshot.image)?;

    // Apply crop if specified
    if let Some(crop) = &screenshot.crop_region {
        // Clamp crop region to image bounds
        if let Some(clamped) = crop.clone().clamp_to_bounds(width, height) {
            dynamic_image = dynamic_image.crop_imm(
                clamped.x,
                clamped.y,
                clamped.width,
                clamped.height,
            );
        }
        // If clamp_to_bounds returns None, the crop region is completely outside
        // the image - we just skip cropping and save the full image
    }

    // Apply burn-in if enabled
    let burn_in_config = config
        .keys
        .get(&screenshot.key)
        .and_then(|k| k.burn_in.as_ref())
        .unwrap_or(&config.burn_in);

    apply_burn_in(&mut dynamic_image, screenshot, burn_in_config);

    // Save atomically using temp file + rename pattern
    atomic_save(&path, &dynamic_image, format, jpeg_quality)?;

    bevy::log::info!("Screenshot saved: {}", path.display());
    Ok(())
}

/// Save an image atomically using temp file + rename pattern.
///
/// This ensures that if the process is killed during save, we either have:
/// - The complete file (if rename succeeded)
/// - An orphaned temp file (if killed before rename)
///
/// Orphaned temp files are cleaned up on next startup.
fn atomic_save(
    final_path: &Path,
    image: &DynamicImage,
    format: ImageFormat,
    jpeg_quality: crate::config::JpegQuality,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    let parent_dir = final_path
        .parent()
        .ok_or("Invalid path: no parent directory")?;

    // Generate temp file path in the same directory (for atomic rename)
    let extension = match format {
        ImageFormat::Png => "png",
        ImageFormat::Jpeg => "jpg",
    };
    let temp_filename = format!(".tmp_{}.{}", Uuid::new_v4(), extension);
    let temp_path = parent_dir.join(&temp_filename);

    // Write to temp file
    {
        let file = std::fs::File::create(&temp_path)?;
        let mut writer = std::io::BufWriter::new(file);

        match format {
            ImageFormat::Png => {
                let encoder = image::codecs::png::PngEncoder::new(&mut writer);
                image.write_with_encoder(encoder)?;
            }
            ImageFormat::Jpeg => {
                let encoder =
                    image::codecs::jpeg::JpegEncoder::new_with_quality(&mut writer, jpeg_quality.as_u8());
                image.write_with_encoder(encoder)?;
            }
        }

        // Flush the buffer
        writer.flush()?;

        // Sync to disk to ensure data is persisted before rename
        writer.into_inner()?.sync_all()?;
    }

    // Atomic rename (on POSIX systems, rename is atomic if src and dst are on same filesystem)
    std::fs::rename(&temp_path, final_path)?;

    Ok(())
}

/// Convert a Bevy Image to an image crate DynamicImage.
fn bevy_image_to_dynamic(
    image: &Image,
) -> Result<DynamicImage, Box<dyn std::error::Error + Send + Sync>> {
    let width = image.width();
    let height = image.height();

    let data = image
        .data
        .as_ref()
        .ok_or("Image has no data")?;

    // Handle different texture formats
    match image.texture_descriptor.format {
        TextureFormat::Rgba8UnormSrgb | TextureFormat::Rgba8Unorm => {
            let rgba_image = image::RgbaImage::from_raw(width, height, data.clone())
                .ok_or("Failed to create RGBA image from raw data")?;
            Ok(DynamicImage::ImageRgba8(rgba_image))
        }
        TextureFormat::Bgra8UnormSrgb | TextureFormat::Bgra8Unorm => {
            // Convert BGRA to RGBA
            let mut rgba_data = data.clone();
            for chunk in rgba_data.chunks_exact_mut(4) {
                chunk.swap(0, 2); // Swap B and R
            }
            let rgba_image = image::RgbaImage::from_raw(width, height, rgba_data)
                .ok_or("Failed to create RGBA image from BGRA data")?;
            Ok(DynamicImage::ImageRgba8(rgba_image))
        }
        other => Err(format!("Unsupported texture format: {:?}", other).into()),
    }
}

fn generate_filename(
    pattern: &str,
    key: &str,
    description: &str,
    timestamp: DateTime<Local>,
    format: ImageFormat,
) -> String {
    let extension = match format {
        ImageFormat::Png => "png",
        ImageFormat::Jpeg => "jpg",
    };

    // Sanitize description for filename
    let safe_description: String = description
        .chars()
        .map(|c| {
            if c.is_alphanumeric() || c == '_' || c == '-' {
                c
            } else {
                '_'
            }
        })
        .collect();

    // Sanitize key for filename
    let safe_key: String = key
        .chars()
        .map(|c| {
            if c.is_alphanumeric() || c == '_' || c == '-' {
                c
            } else {
                '_'
            }
        })
        .collect();

    let timestamp_str = timestamp.format("%Y%m%d_%H%M%S_%3f").to_string();

    let filename = pattern
        .replace("{timestamp}", &timestamp_str)
        .replace("{key}", &safe_key)
        .replace("{description}", &safe_description);

    format!("{}.{}", filename, extension)
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::TimeZone;

    #[test]
    fn test_generate_filename_basic() {
        let timestamp = Local.with_ymd_and_hms(2024, 1, 15, 10, 30, 45).unwrap();
        let filename = generate_filename(
            "{timestamp}_{key}_{description}",
            "combat",
            "boss_defeated",
            timestamp,
            ImageFormat::Png,
        );
        assert!(filename.starts_with("20240115_103045_"));
        assert!(filename.contains("combat"));
        assert!(filename.contains("boss_defeated"));
        assert!(filename.ends_with(".png"));
    }

    #[test]
    fn test_generate_filename_sanitizes_description() {
        let timestamp = Local.with_ymd_and_hms(2024, 1, 15, 10, 30, 45).unwrap();
        let filename = generate_filename(
            "{timestamp}_{key}_{description}",
            "test",
            "hello world/bad:chars",
            timestamp,
            ImageFormat::Jpeg,
        );
        assert!(!filename.contains(' '));
        assert!(!filename.contains('/'));
        assert!(!filename.contains(':'));
        assert!(filename.ends_with(".jpg"));
    }
}