I'm working on a rust program that takes a directory full of images as input, and outputs a video file which chaotically flashes, strobes, and displays images at random times, following intricate and chaotic patterns. The transitions are fast and chaotic, creating an almost epileptic effect. It has a modular way of implementing interesting effects. I'll share the source code for each source file below. Please update it to work not just with static images, but also animateds gifs.


////// src/config.rs //////


#[derive(Debug, Clone)]
pub struct Config {
    pub input_dirs: Vec<String>,
    pub output_file: String,
    pub width: usize,
    pub height: usize,
    pub framerate: u32,
    pub duration_seconds: u32,
    pub total_frames: usize,
    pub timebase_num: u32,
    pub timebase_den: u32,
    pub quantizer: usize,
    pub speed: usize,
    pub bit_depth: usize,
    // Add other configuration parameters as needed
}

impl Default for Config
{
    fn default() -> Self {
        let framerate = 60;
        let duration_seconds = crate::DEFAULT_DURATION_S;
        let total_frames = framerate * duration_seconds;
        Self 
        { 
            input_dirs: crate::INPUT_DIRS.iter().map(|d| d.to_string()).collect(), 
            output_file: crate::DEFAULT_OUTPUT_FILE.to_string(), 
            width: crate::DEFAULT_WIDTH, 
            height: crate::DEFAULT_HEIGHT, 
            framerate, 
            duration_seconds, 
            total_frames: total_frames as usize,
            timebase_num: 1u32, 
            timebase_den: framerate,
            quantizer: 100,
            speed: 6,
            bit_depth: 8,
        }
    }
}


////// src/effects.rs //////

use image::{ImageBuffer, Pixel, Rgba, RgbaImage};
use rand::prelude::*;
use rusttype::Font;

use crate::types;

// Define the Effect enum with all possible effects
#[derive(Clone)]
pub enum Effect {
    InvertColors,
    ColorFlash(image::Rgba<u8>),
    Sparkle,
    Spiral,
    TransparencyStrobe(f32), // The parameter represents strobe speed.
    OverlayText(String),
    ZoomIn(f32), // This parameter represents the zoom speed.
    Kaleidoscope,
    BitFocus,
    Glow { max_intensity: f32 },
    GrowShrink { speed: f32 },     // New effect for growing and shrinking
    FadeOut,                       // New effect for fading out
    SplitApart,                    // New effect for splitting apart
    LowOpacityFlash {
        flash_duration: usize,
        flash_interval: usize,
    },
    // Add more effects as needed
}


// EffectFactory to generate random effects
pub struct EffectFactory {
    pub rng: ThreadRng,
    texts: Vec<String>,
    font: rusttype::Font<'static>,
}

impl EffectFactory {
    pub fn new() -> Self {
        // Random number generator
        let font_data = include_bytes!(r"C:\Project\Workspace\Javascript\gpt4o-site\media-cycler-rs\Atop-R99O3.ttf") as &[u8];
        Self
        {
            rng: rand::thread_rng(),
            texts: crate::TEXTS.iter().map(|t| t.to_string()).collect(),
            font: Font::try_from_bytes(font_data).expect("Error loading font"),
        }
    }

    pub fn generate_random_effects(&mut self) -> Vec<Effect> {
        let mut effects = Vec::new();

        if self.rng.gen_bool(0.3) {
            effects.push(Effect::InvertColors);
        }
        if self.rng.gen_bool(0.2) {
            let color = image::Rgba([
                self.rng.gen_range(0..=255),
                self.rng.gen_range(0..=255),
                self.rng.gen_range(0..=255),
                128,
            ]);
            effects.push(Effect::ColorFlash(color));
        }
        if self.rng.gen_bool(0.2) {
            effects.push(Effect::Sparkle);
        }
        if self.rng.gen_bool(0.2) {
            effects.push(Effect::Spiral);
        }
        if self.rng.gen_bool(0.1) {
            effects.push(Effect::TransparencyStrobe(self.rng.gen_range(0.2..=10.0)));
        }
        if self.rng.gen_bool(0.2) {
            let text = self.texts.choose(&mut self.rng).unwrap().to_string();
            effects.push(Effect::OverlayText(text));
        }
        if self.rng.gen_bool(0.1) {
            let zoom_speed = self.rng.gen_range(0.001..0.01); // Adjust zoom speed as desired
            effects.push(Effect::ZoomIn(zoom_speed));
        }
        if self.rng.gen_bool(0.02) {
            effects.push(Effect::Kaleidoscope);
        }
        if self.rng.gen_bool(0.3) {
            effects.push(Effect::BitFocus);
        }
        if self.rng.gen_bool(0.2) {
            effects.push(Effect::Glow{max_intensity: self.rng.gen_range(1.5..3.0)});
        }
        effects
    }

    pub fn generate_overlay_effects(&mut self) -> Vec<Effect> {
        let mut effects = Vec::new();
        if self.rng.gen_bool(0.3) {
            let speed = self.rng.gen_range(0.01..0.05);
            effects.push(Effect::GrowShrink { speed });
        }
        if self.rng.gen_bool(0.3) {
            effects.push(Effect::FadeOut);
        }
        if self.rng.gen_bool(0.2) {
            effects.push(Effect::SplitApart);
        }
        if self.rng.gen_bool(0.2) {
            effects.push(Effect::Sparkle);
        }
        if self.rng.gen_bool(0.2) {
            effects.push(Effect::Spiral);
        }
        if self.rng.gen_bool(0.2) {
            effects.push(Effect::Glow{max_intensity: self.rng.gen_range(1.5..3.0)});
        }
        if self.rng.gen_bool(0.3) {
            effects.push(Effect::BitFocus);
        }
        if self.rng.gen_bool(0.2) {
            let color = image::Rgba([
                self.rng.gen_range(0..=255),
                self.rng.gen_range(0..=255),
                self.rng.gen_range(0..=255),
                128,
            ]);
            effects.push(Effect::ColorFlash(color));
        }
        if self.rng.gen_bool(0.01) {
            effects.push(Effect::InvertColors);
        }

        effects
    }

    pub fn generate_random_global_effects(&mut self) -> Vec<Effect> {
        let mut effects = Vec::new();

        if self.rng.gen_bool(0.3) {
            effects.push(Effect::InvertColors);
        }
        if self.rng.gen_bool(0.2) {
            let color = image::Rgba([
                self.rng.gen_range(0..=255),
                self.rng.gen_range(0..=255),
                self.rng.gen_range(0..=255),
                128,
            ]);
            effects.push(Effect::ColorFlash(color));
        }
        if self.rng.gen_bool(0.2) {
            effects.push(Effect::Sparkle);
        }
        if self.rng.gen_bool(0.2) {
            effects.push(Effect::Spiral);
        }
        if self.rng.gen_bool(0.2) {
            let text = self.texts.choose(&mut self.rng).unwrap().to_string();
            effects.push(Effect::OverlayText(text));
        }
        effects
    }

    // Implement effect functions
    /// Returns new image scale
    pub fn apply_effects(
        &mut self,
        image: &mut image::RgbaImage,
        frame_number: usize,
        state: &types::ImageState,
    ) -> ImageBuffer<Rgba<u8>, Vec<u8>> {
        let mut current_scale = state.scale;
        let mut alpha = 1.0;
        for effect in &state.effects {
            match effect {
                Effect::InvertColors => Self::invert_colors(image),
                Effect::ColorFlash(color) => Self::add_color_flash(image, &color),
                Effect::Sparkle => self.add_sparkle(image),
                Effect::Spiral => self.add_spiral(image),
                Effect::TransparencyStrobe(speed) => {
                    Self::apply_transparency_strobe(image, frame_number, state.start_frame, *speed)
                },
                Effect::OverlayText(text) => {
                    self.overlay_text(image, &text)
                },
                Effect::ZoomIn(zoom_speed) => {
                    current_scale = Self::calculate_zoom_scale(
                        state.scale,
                        frame_number,
                        state.start_frame,
                        *zoom_speed,
                    );
                }
                Effect::Kaleidoscope => Self::apply_kaleidoscope(image),
                Effect::BitFocus => Self::apply_bit_focus(
                    image,
                    frame_number,
                    state.start_frame,
                    state.end_frame,
                ),
                Effect::Glow { max_intensity } => {
                    self.apply_glow(image, frame_number, state, *max_intensity);
                },
                Effect::GrowShrink { speed } => {
                    let scale_factor = self.calculate_grow_shrink_scale(
                        frame_number,
                        state.start_frame,
                        state.end_frame,
                        *speed,
                    );
                    current_scale *= scale_factor;
                },
                Effect::FadeOut => {
                    alpha = self.calculate_fade_out_alpha(
                        frame_number,
                        state.start_frame,
                        state.end_frame,
                    );
                },
                Effect::SplitApart => {
                    *image = self.apply_split_apart(image, frame_number, state);
                },
                Effect::LowOpacityFlash {
                    flash_duration,
                    flash_interval,
                } => {
                    alpha = self.calculate_low_opacity_flash_alpha(
                        frame_number,
                        state.start_frame,
                        *flash_duration,
                        *flash_interval,
                    );
                },
            }
        }
        // Apply scaling
        let scaled_image = image::imageops::resize(
            image,
            (image.width() as f32 * current_scale) as u32,
            (image.height() as f32 * current_scale) as u32,
            image::imageops::FilterType::Lanczos3,
        );

        // Apply alpha (transparency)
        let final_image = self.apply_alpha(&scaled_image, alpha);

        final_image
    }

    // Function to calculate alpha for LowOpacityFlash effect
    fn calculate_low_opacity_flash_alpha(
        &self,
        frame_number: usize,
        start_frame: usize,
        flash_duration: usize,
        flash_interval: usize,
    ) -> f32 {
        let elapsed_frames = frame_number - start_frame;
        let cycle = flash_duration + flash_interval;

        // Determine where we are in the flash cycle
        let position_in_cycle = elapsed_frames % cycle;

        if position_in_cycle < flash_duration {
            1.0 // Full opacity during the flash
        } else {
            0.2 // Low opacity otherwise
        }
    }

    // Function to calculate scale for GrowShrink effect
    fn calculate_grow_shrink_scale(
        &self,
        frame_number: usize,
        start_frame: usize,
        end_frame: usize,
        speed: f32,
    ) -> f32 {
        let total_frames = (end_frame - start_frame) as f32;
        let elapsed_frames = (frame_number - start_frame) as f32;

        let progress = elapsed_frames / total_frames;
        let scale = (progress * speed * std::f32::consts::PI * 2.0).sin() * 0.5 + 1.0; // Oscillate between 0.5 and 1.5

        scale
    }

    // Function to calculate alpha for FadeOut effect
    fn calculate_fade_out_alpha(
        &self,
        frame_number: usize,
        start_frame: usize,
        end_frame: usize,
    ) -> f32 {
        let total_frames = (end_frame - start_frame) as f32;
        let elapsed_frames = (frame_number - start_frame) as f32;

        let progress = (elapsed_frames / total_frames).clamp(0.0, 1.0);
        let alpha = 1.0 - progress; // Fade out from 1.0 to 0.0

        alpha
    }

    // Function to apply alpha to an image
    fn apply_alpha(
        &self,
        image: &RgbaImage,
        alpha: f32,
    ) -> RgbaImage {
        let mut final_image = image.clone();

        for pixel in final_image.pixels_mut() {
            pixel[3] = (pixel[3] as f32 * alpha).clamp(0.0, 255.0) as u8;
        }

        final_image
    }

    // Function to apply the SplitApart effect
    fn apply_split_apart(
        &self,
        image: &mut RgbaImage,
        frame_number: usize,
        state: &types::ImageState,
    ) -> RgbaImage {
        // For simplicity, we'll split the image into two halves moving apart

        let total_frames = (state.end_frame - state.start_frame) as f32;
        let elapsed_frames = (frame_number - state.start_frame) as f32;
        let progress = (elapsed_frames / total_frames).clamp(0.0, 1.0);

        let width = image.width();
        let height = image.height();

        let mut new_image = RgbaImage::new(width, height);

        // Left half moves left
        let offset = (progress * width as f32 / 2.0) as i64;
        let left_half = image::imageops::crop_imm(image, 0, 0, width / 2, height).to_image();
        image::imageops::overlay(&mut new_image, &left_half, -offset, 0);

        // Right half moves right
        let right_half = image::imageops::crop_imm(image, width / 2, 0, width / 2, height).to_image();
        image::imageops::overlay(&mut new_image, &right_half, offset, 0);

        new_image
    }

    // Implement the apply_glow function
    fn apply_glow(
        &self,
        image: &mut RgbaImage,
        frame_number: usize,
        state: &types::ImageState,
        max_intensity: f32
    ) {
        // Calculate the progress of the effect between 0.0 and 1.0
        let total_frames = (state.end_frame - state.start_frame) as f32;
        let elapsed_frames = (frame_number - state.start_frame) as f32;
        let progress = (elapsed_frames / total_frames).clamp(0.0, 1.0);

        let brightness_factor = 1.0 + progress * (max_intensity - 1.0);

        for pixel in image.pixels_mut() {
            for i in 0..3 {
                let value = pixel[i] as f32 * brightness_factor;
                pixel[i] = value.clamp(0.0, 255.0) as u8;
            }
        }
    }

    // Function to invert colors in an image
    fn invert_colors(image: &mut RgbaImage) {
        for pixel in image.pixels_mut() {
            pixel[0] = 255 - pixel[0];
            pixel[1] = 255 - pixel[1];
            pixel[2] = 255 - pixel[2];
        }
    }

    // Function to add a color flash effect
    fn add_color_flash(image: &mut RgbaImage, color: &image::Rgba<u8>) {
        for pixel in image.pixels_mut() {
            pixel.blend(&color);
        }
    }

    fn apply_transparency_strobe(
        image: &mut RgbaImage,
        frame_number: usize,
        start_frame: usize,
        strobe_speed: f32,
    ) {
        // Calculate the alpha value based on a sine wave for smooth strobing
        let t = (frame_number - start_frame) as f32;
        let alpha = ((strobe_speed * t).sin() * 0.5 + 0.5) * 255.0;
        let alpha = alpha as u8;

        for pixel in image.pixels_mut() {
            pixel[3] = alpha;
        }
    }

    // Function to add sparkle effects
    fn add_sparkle(&mut self, image: &mut RgbaImage) {
        let num_sparkles = self.rng.gen_range(100..500);
        for _ in 0..num_sparkles {
            let x = self.rng.gen_range(0..image.width());
            let y = self.rng.gen_range(0..image.height());
            let pixel = image.get_pixel_mut(x, y);
            pixel[0] = 255;
            pixel[1] = 255;
            pixel[2] = 255;
            pixel[3] = 255;
        }
    }

    // Function to add a spiral effect
    fn add_spiral(&mut self, image: &mut RgbaImage) {
        let center_x = image.width() as i32 as f32 / 2.0;
        let center_y = image.height() as i32 as f32 / 2.0;
        let max_radius = center_x.min(center_y);
        let num_turns = self.rng.gen_range(5.0..10.0);
        let color = image::Rgba([self.rng.gen(), self.rng.gen(), self.rng.gen(), 255]);

        for t in 0..1000 {
            let theta = num_turns * t as f32 * 0.01 * std::f32::consts::PI;
            let radius = max_radius * t as f32 * 0.001;
            let x = center_x + radius * theta.cos();
            let y = center_y + radius * theta.sin();

            if x >= 0.0 && x < image.width() as f32 && y >= 0.0 && y < image.height() as f32 {
                let pixel = image.get_pixel_mut(x as u32, y as u32);
                *pixel = color;
            }
        }
    }

    fn overlay_text(&mut self,
        image: &mut RgbaImage,
        text: &str
    ) {
        use imageproc::drawing::draw_text_mut;
        use rusttype::Scale;

        let scale = Scale {
            x: self.rng.gen_range(20.0..100.0), // Random font size
            y: self.rng.gen_range(20.0..100.0),
        };

        // Random position
        let x = self.rng.gen_range(0..(image.width() as i32));
        let y = self.rng.gen_range(0..(image.height() as i32));

        // Random color
        let color = image::Rgba([
            self.rng.gen(),
            self.rng.gen(),
            self.rng.gen(),
            255,
        ]);

        draw_text_mut(
            image,
            color,
            x,
            y,
            scale,
            &self.font,
            text,
        );
    }

    fn calculate_zoom_scale(
        original_scale: f32,
        frame_number: usize,
        start_frame: usize,
        zoom_speed: f32,
    ) -> f32 {
        let t = (frame_number - start_frame) as f32;
        let scale = original_scale * (1.0 + zoom_speed * t);
        scale
    }

    fn apply_kaleidoscope(image: &mut RgbaImage) {
        let width = image.width();
        let height = image.height();
    
        // Create a new image buffer to store the kaleidoscope result
        let mut output = RgbaImage::new(width, height);
    
        // Number of segments in the kaleidoscope (adjust for different patterns)
        let num_segments = 8;
        let angle_per_segment = 360.0 / num_segments as f32;
        let center_x = width as f32 / 2.0;
        let center_y = height as f32 / 2.0;
    
        // Iterate over each pixel in the output image
        for y in 0..height {
            for x in 0..width {
                // Calculate coordinates relative to the center
                let dx = x as f32 - center_x;
                let dy = y as f32 - center_y;
    
                // Calculate the angle and radius from the center
                let mut angle = dy.atan2(dx).to_degrees();
                let radius = (dx * dx + dy * dy).sqrt();
    
                // Normalize the angle within [0, 360)
                if angle < 0.0 {
                    angle += 360.0;
                }
                // Map the angle into the first segment
                angle = angle % angle_per_segment;
    
                // Reflect the angle to create symmetry
                if angle > angle_per_segment / 2.0 {
                    angle = angle_per_segment - angle;
                }
    
                // Convert back to radians
                let angle_rad = angle.to_radians();
    
                // Calculate the source pixel coordinates
                let src_x = center_x + radius * angle_rad.cos();
                let src_y = center_y + radius * angle_rad.sin();
    
                // Check bounds and sample the pixel from the original image
                if src_x >= 0.0 && src_x < width as f32 && src_y >= 0.0 && src_y < height as f32 {
                    let src_pixel = image.get_pixel(src_x as u32, src_y as u32);
                    output.put_pixel(x, y, *src_pixel);
                } else {
                    // If out of bounds, set pixel to black or transparent
                    output.put_pixel(x, y, image::Rgba([0, 0, 0, 255]));
                }
            }
        }
    
        // Replace the original image data with the kaleidoscope result
        *image = output;
    }

    fn apply_bit_focus(
        image: &mut RgbaImage,
        frame_number: usize,
        start_frame: usize,
        end_frame: usize,
    ) {
        let total_duration = end_frame - start_frame;
        let elapsed = frame_number - start_frame;
    
        // Calculate the progression ratio [0.0, 1.0]
        let progress = elapsed as f32 / total_duration as f32;
    
        // Determine the pixelation level based on progress
        // Start with a small size (high pixelation) and increase to full size
        let min_pixels = 8; // Minimum number of pixels for pixelation (more pixelated)
        let max_pixels = 1; // Maximum number of pixels (full resolution)
    
        // Use an easing function for smoother animation (optional)
        let eased_progress = Self::ease_out_quad(progress);
    
        // Calculate current pixelation level
        let pixel_size = ((min_pixels as f32 * (1.0 - eased_progress)) + (max_pixels as f32 * eased_progress)).max(1.0);
    
        // Resize down to create pixelation, then upscale back to original size
        let (width, height) = image.dimensions();
        let new_width = (width as f32 / pixel_size).max(1.0) as u32;
        let new_height = (height as f32 / pixel_size).max(1.0) as u32;
    
        // Downscale the image to the new dimensions
        let downscaled = image::imageops::resize(
            image,
            new_width,
            new_height,
            image::imageops::FilterType::Nearest,
        );
    
        // Upscale back to original size
        let upscaled = image::imageops::resize(
            &downscaled,
            width,
            height,
            image::imageops::FilterType::Nearest,
        );
    
        // Replace the original image with the upscaled, pixelated image
        *image = upscaled;
    }
    
    // Optional easing function for smoother progression
    fn ease_out_quad(t: f32) -> f32 {
        1.0 - (1.0 - t).powi(2)
    }
}

////// src/frame_renderer.rs //////

use crate::transitions::{apply_reassembly_transition, TransitionEffect};
use crate::{config::Config, transitions::Transition};
use crate::effects::EffectFactory;
use crate::types::ImageState;
use image::{imageops, DynamicImage, GenericImageView, Pixel, RgbaImage};
use rand::{rngs::ThreadRng, Rng};

pub fn render_frame(
    frame_number: usize,
    active_primary_images: &[ImageState],
    active_overlay_images: &[ImageState],
    active_global_effects: &[ImageState],
    config: &Config,
    effect_factory: &mut EffectFactory,
    transitions: &mut [&mut Transition],
) -> RgbaImage {
    // Create a new blank image (frame) with the specified dimensions
    let mut frame = DynamicImage::new_rgba8(config.width as u32, config.height as u32).to_rgba8();

    // Overlay images
    for state in active_primary_images  {
        let mut img = state.image.clone();

        // Apply effects
        let img = effect_factory.apply_effects(
            &mut img,
            frame_number,
            &state
        );

        // Overlay the image onto the frame
        imageops::overlay(&mut frame, &img, state.position.0, state.position.1);
    }

    // Render overlay images (smaller, fast images)
    for state in active_overlay_images {
        let mut img = state.image.clone();

        // Apply effects
        let img = effect_factory.apply_effects(&mut img, frame_number, state);

        // Overlay the image onto the frame
        imageops::overlay(&mut frame, &img, state.position.0, state.position.1);
    }

    // Apply global effects to the frame
    for global_state in active_global_effects {
        effect_factory.apply_effects(&mut frame, 0, global_state);
    }

    // Render transitions
    for transition in transitions {
        if frame_number >= transition.start_frame && frame_number < transition.end_frame {
            match transition.effect {
                TransitionEffect::Reassembly => {
                    let transition_frame = apply_reassembly_transition(
                        frame_number,
                        transition,
                        config,
                    );
                    // Overlay the transition_frame onto the main frame
                    image::imageops::overlay(&mut frame, &transition_frame, 0, 0);
                },
                // Handle other transition effects...
            }
        }
    }

    frame
}

////// src/image_loader.rs //////

use anyhow::Result;
use image::{self, DynamicImage};
use std::fs;

pub fn load_images(input_dirs: &[&str]) -> Result<Vec<DynamicImage>> {
    let mut media = Vec::new();

    for input_dir in input_dirs
    {
        for entry in fs::read_dir(input_dir)? {
            let entry = entry?;
            let path = entry.path();
            if path.is_file() {
                let extension = path.extension().and_then(|e| e.to_str()).unwrap_or("").to_lowercase();

                match extension.as_str() {
                    "jpg" | "jpeg" | "png" | "bmp" | "tiff" | "gif" => {
                        // Load static image
                        if let Ok(img) = image::open(&path) {
                            media.push(img.to_rgba8().into());
                        }
                    },
                    _ => {
                        println!("Unsupported file type: {}", path.display());
                    }
                }
            }
        }
    }

    if media.is_empty() {
        Err(anyhow::anyhow!("No media files found in the specified directory."))
    } else {
        Ok(media)
    }
}

////// src/main.rs //////

use indicatif::{ProgressBar, ProgressStyle};

mod image_loader;
mod effects;
mod frame_renderer;
mod config;
mod scheduler;
mod types;
mod video_encoder;
mod transitions;

const INPUT_DIRS: &[&str] = &[
    r"\\192.168.1.99\eoluSUSE_home\ns\Images", 
    r"\\192.168.1.99\eoluSUSE_home\ns\Images\Set0", 
    r"\\192.168.1.99\eoluSUSE_home\ns\Images\Set1", 
    r"\\192.168.1.99\eoluSUSE_home\ns\Images\Set2", 
    r"\\192.168.1.99\eoluSUSE_home\ns\Images\Set3",
    r"\\192.168.1.99\eoluSUSE_home\ns\Images\Animated", 
    r"\\192.168.1.99\eoluSUSE_home\ns\Images\Animated2", 
    r"\\192.168.1.99\eoluSUSE_home\ns\Images\Animated3", 
    r"\\192.168.1.99\eoluSUSE_home\ns\Images\Animated4", 
];
pub(crate) const TEXTS: &[&str] = &[
    "SUBMIT",
    "OBEY",
    "HORNY",
    "CRAVE",
    "DEEP",
];
// const INPUT_DIRS: &[&str] = &[r"E:\Pictures\iPhone", r"E:\Pictures\Old Backgrounds"];
// pub(crate) const TEXTS: &[&str] = &[
//     "Wow",
//     "Cool",
//     "AI",
//     "Future",
//     "Overdrews",
// ];
pub const DEFAULT_DURATION_S: u32 = 60;
pub const DEFAULT_WIDTH: usize = 1920;
pub const DEFAULT_HEIGHT: usize = 1080;
pub const DEFAULT_OUTPUT_FILE: &str = "output2.ivf";

fn main() -> anyhow::Result<()> {
    let config = config::Config::default();

    // Initialize progress bar
    let pb = ProgressBar::new(config.total_frames as u64);
    pb.set_style(
        ProgressStyle::default_bar()
            .template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {pos}/{len} frames ({eta} remaining)")?
            .progress_chars("#>-")
    );

    // Read all images from the input directory
    let images = image_loader::load_images(INPUT_DIRS).unwrap();

    // Initialize effect factory
    let mut effect_factory = effects::EffectFactory::new();

    // Initialize image scheduler
    let mut scheduler = scheduler::ImageScheduler::new(
        &images,
        &config,
        &mut effect_factory,
    );

    // Initialize video encoder
    let mut encoder = video_encoder::VideoEncoder::new(&config)?;

    // Main frame generation loop
    for frame_number in 0..config.total_frames {
        // Get the list of active image states for the current frame
        let active_primary_images = scheduler.get_active_primary_images(frame_number);
        let active_overlay_images = scheduler.get_active_overlay_images(frame_number);
        let active_globals = scheduler.get_active_global_effects(frame_number);
        let mut active_transitions = scheduler.get_active_transitions(frame_number);

        // Render the frame
        let frame = frame_renderer::render_frame(
            frame_number,
            &active_primary_images,
            &active_overlay_images,
            &active_globals,
            &config,
            &mut effect_factory,
            &mut active_transitions,
        );

        // Encode the frame
        encoder.encode_frame(&frame)?;

        // Update scheduler (if needed)
        scheduler.update(frame_number);

        // Update progress bar
        pb.inc(1);
    }

    // Finalize encoding
    encoder.finish()?;

    // Finish progress bar
    pb.finish_with_message(format!("Video has been successfully encoded to {}\n Convert the IVF file to MP4 using FFmpeg for playback.", config.output_file));

    Ok(())
}


////// src/scheduler.rs //////

use crate::config::Config;
use crate::effects::{self, Effect};
use crate::transitions::{Transition, TransitionEffect};
use crate::types::ImageState;
use rand::prelude::*;
use rand::seq::SliceRandom;

pub struct ImageScheduler<'a> {
    primary_images: Vec<ImageState>,
    overlay_images: Vec<ImageState>, // New vector for overlay images
    rng: ThreadRng,
    config: &'a Config,
    global_effects: Vec<ImageState>,
    transitions: Vec<Transition>,
    overlay_schedule: Vec<usize>,
}

impl<'a> ImageScheduler<'a> {
    pub fn new(
        images: &'a [image::DynamicImage],
        config: &'a crate::config::Config,
        effect_factory: &mut effects::EffectFactory,
    ) -> Self {
        let mut rng = rand::thread_rng();
        let mut scheduled_images = Vec::new();
        let mut overlay_images = Vec::new();

        // Randomly schedule primary images
        for _ in 0..(config.total_frames / 5) {
            // Randomly select an image
            let img = images.choose(&mut rng).unwrap().to_rgba8();

            // Random start frame
            let start_frame = rng.gen_range(0..config.total_frames) as u32;
            // Maximum duration in frames (2 seconds at 30 fps)
            let max_duration = (config.framerate * 2).min(config.total_frames as u32 - start_frame) as usize; // Up to 2 seconds
            // Generate duration favoring shorter durations
            let duration = Self::generate_weighted_duration(&mut rng, max_duration);
            // Calculate end frame
            let end_frame = start_frame as usize + duration;

            // Random scaling factor
            let scale = {
                let max_scale_x = config.width as f32 / img.width() as f32;
                let max_scale_y = config.height as f32 / img.height() as f32;
                let max_scale = max_scale_x.min(max_scale_y).min(1.5); // Allow up to 1.5x scaling
                let min_scale = 0.5_f32.min(max_scale);
                if max_scale > min_scale {
                    rng.gen_range(min_scale..=max_scale)
                } else {
                    min_scale
                }
            };

            // Apply scaling (pre-scaling to improve performance)
            let img = image::imageops::resize(
                &img,
                (img.width() as f32 * scale) as u32,
                (img.height() as f32 * scale) as u32,
                image::imageops::FilterType::Lanczos3,
            );

            // Random position
            let max_x = config.width.saturating_sub(img.width() as usize) as i64;
            let max_y = config.height.saturating_sub(img.height() as usize) as i64;
            let x = if max_x > 0 {
                rng.gen_range(0..=max_x)
            } else {
                0
            };
            let y = if max_y > 0 {
                rng.gen_range(0..=max_y)
            } else {
                0
            };

            // Random effects
            let random_effects = effect_factory.generate_random_effects();

            // Create the image state
            let image_state = ImageState {
                image: img,
                position: (x, y),
                scale: 1.0, // Already scaled
                start_frame: start_frame as usize,
                end_frame,
                effects: random_effects,
            };

            scheduled_images.push(image_state);
        }

        // Generate a schedule for overlay images
        let overlay_schedule = Self::generate_overlay_schedule(config.total_frames, &mut rng);

        // Schedule overlay images
        for frame_number in &overlay_schedule {
            // Randomly select an image
            let img = images.choose(&mut rng).unwrap().to_rgba8();

            // Short durations for overlay images
            let max_duration = 10.min(config.total_frames - frame_number) as usize;
            let duration = Self::generate_short_duration(&mut rng, max_duration);
            let end_frame = frame_number + duration;

            // Small scaling factor
            let scale = rng.gen_range(0.05..=0.2); // Images scaled down to 5% - 20% of original size

            // Apply scaling
            let img = image::imageops::resize(
                &img,
                (img.width() as f32 * scale) as u32,
                (img.height() as f32 * scale) as u32,
                image::imageops::FilterType::Lanczos3,
            );

            // Random position
            let max_x = config.width.saturating_sub(img.width() as usize) as i64;
            let max_y = config.height.saturating_sub(img.height() as usize) as i64;
            let x = if max_x > 0 {
                rng.gen_range(0..=max_x)
            } else {
                0
            };
            let y = if max_y > 0 {
                rng.gen_range(0..=max_y)
            } else {
                0
            };

            // Random effects
            let random_effects = effect_factory.generate_overlay_effects();

            // Create the image state
            let image_state = ImageState {
                image: img,
                position: (x, y),
                scale: 1.0, // Already scaled
                start_frame: *frame_number,
                end_frame,
                effects: random_effects,
            };

            overlay_images.push(image_state);
        }
        
        // Schedule larger, slower overlay images
        for _ in 0..(config.total_frames / 100) { // Adjust the count as needed
            // Randomly select an image
            let img = images.choose(&mut rng).unwrap().to_rgba8();

            // Random start frame
            let start_frame = rng.gen_range(0..config.total_frames) as usize;

            let min_duration = config.framerate * 2;
            let max_possible_duration = config.total_frames - start_frame;

            // Longer duration
            let max_duration = (config.framerate * 5).min(max_possible_duration as u32);
            let duration = if max_duration >= min_duration {
                rng.gen_range(min_duration..=max_duration)
            } else if max_duration > 0 {
                // If max_duration is less than min_duration but greater than zero
                rng.gen_range(1..=max_duration)
            } else {
                // If there's no room for duration, skip this iteration
                continue;
            };
            let end_frame = start_frame + duration as usize;

            // Larger scaling factor
            let scale = rng.gen_range(0.5..=1.5); // Images scaled between 50% and 150% of original size

            // Apply scaling
            let img = image::imageops::resize(
                &img,
                (img.width() as f32 * scale) as u32,
                (img.height() as f32 * scale) as u32,
                image::imageops::FilterType::Lanczos3,
            );

            // Random position
            let max_x = config.width.saturating_sub(img.width() as usize) as i64;
            let max_y = config.height.saturating_sub(img.height() as usize) as i64;
            let x = if max_x > 0 {
                rng.gen_range(0..=max_x)
            } else {
                0
            };
            let y = if max_y > 0 {
                rng.gen_range(0..=max_y)
            } else {
                0
            };

            // Apply low opacity with occasional flashes
            let effects = vec![Effect::LowOpacityFlash {
                flash_duration: rng.gen_range(1..=3), // Flash lasts 1 to 3 frames
                flash_interval: rng.gen_range(30..=100), // Flash every 30 to 100 frames
            }];

            // Create the image state
            let image_state = ImageState {
                image: img,
                position: (x, y),
                scale: 1.0, // Already scaled
                start_frame,
                end_frame,
                effects,
            };

            overlay_images.push(image_state);
        }

        // Generate global effects
        let mut global_effects = Vec::new();
        for _ in 0..(config.total_frames / 50) { // Adjust the frequency as needed
            // Random start frame
            let start_frame = rng.gen_range(0..config.total_frames);

            // Maximum duration in frames
            let max_duration = ((config.framerate as f32 * 2.5)
                .min(config.total_frames as f32 - start_frame as f32)) as usize;
            let duration = Self::generate_weighted_duration(&mut rng, max_duration);
            let end_frame = start_frame + duration;

            // Generate global effects
            let effects = effect_factory.generate_random_global_effects();

            // Create a dummy ImageState for global effects
            let image_state = ImageState {
                image: image::RgbaImage::new(config.width as u32, config.height as u32),
                position: (0, 0),
                scale: 1.0,
                start_frame,
                end_frame,
                effects,
            };

            global_effects.push(image_state);
        }

        // Sort images by start frame
        // TODO: Is this line a hallucination?
        scheduled_images.sort_by_key(|state| state.start_frame);

        let mut scheduler = Self {
            primary_images: scheduled_images,
            overlay_images,
            transitions: Vec::new(),
            rng,
            config,
            global_effects,
            overlay_schedule,
        };

        // Schedule transitions between images
        scheduler.schedule_transitions();

        scheduler
    }

    // New function to generate overlay image schedule
    fn generate_overlay_schedule(total_frames: usize, rng: &mut ThreadRng) -> Vec<usize> {
        let mut overlay_schedule = Vec::new();
        let mut frame_number = 0;
    
        while frame_number < total_frames {
            // Randomly decide whether to schedule an overlay image
            if rng.gen_bool(0.5) { // 50% chance to schedule an image
                overlay_schedule.push(frame_number);
            }

            // Calculate the maximum possible gap without exceeding total_frames
            let max_gap = (total_frames - frame_number).min(30);
            if max_gap == 0 {
                break; // No more frames left
            }
    
            // Random gap before next potential overlay image
            let gap = rng.gen_range(1..=max_gap); // Adjust range as needed
            frame_number += gap;
        }
    
        overlay_schedule
    }

    // Helper function to generate short durations for overlay images
    fn generate_short_duration(rng: &mut ThreadRng, max_duration: usize) -> usize {
        if max_duration == 0 {
            1 // Default to a duration of 1 frame if max_duration is zero
        } else {
            let duration = rng.gen_range(1..=5).min(max_duration);
            duration
        }
    }

    pub fn get_active_primary_images(&self, frame_number: usize) -> Vec<ImageState> {
        self.primary_images
            .iter()
            .filter(|state| state.start_frame <= frame_number && state.end_frame > frame_number)
            .cloned()
            .collect()
    }

    pub fn get_active_overlay_images(&self, frame_number: usize) -> Vec<ImageState> {
        self.overlay_images
            .iter()
            .filter(|state| state.start_frame <= frame_number && state.end_frame > frame_number)
            .cloned()
            .collect()
    }

    pub fn get_active_global_effects(&self, frame_number: usize) -> Vec<ImageState> {
        self.global_effects
            .iter()
            .filter(|state| state.start_frame <= frame_number && state.end_frame > frame_number)
            .cloned()
            .collect()
    }

    pub fn get_active_transitions(&mut self, frame_number: usize) -> Vec<&mut Transition> {
        self.transitions
            .iter_mut()
            .filter(|transition| transition.start_frame <= frame_number && transition.end_frame > frame_number)
            .collect()
    }

    pub fn update(&mut self, _frame_number: usize) {
        // Optionally update the scheduler state if images need to be rescheduled
        // For example, remove images that have ended, add new ones, etc.
        // For now, we'll keep the initial set of images
    }

    fn generate_weighted_duration(rng: &mut ThreadRng, max_duration: usize) -> usize {
        let mut duration_options = Vec::new();

        // Always include duration of 1 frame if max_duration >= 1
        if max_duration >= 1 {
            duration_options.push((1..=1, 70)); // 70% chance of duration being 1 frame
        }

        // Include duration of 2 frames if max_duration >= 2
        if max_duration >= 2 {
            duration_options.push((2..=2, 15)); // 15% chance of duration being 2 frames
        }

        // Include durations between 3 and up to max_duration (min of 5)
        if max_duration >= 3 {
            let upper = max_duration.min(5);
            if upper >= 3 {
                duration_options.push((3..=upper, 10)); // 10% chance
            }
        }

        // Include longer durations if max_duration >= 6
        if max_duration >= 6 {
            duration_options.push((6..=max_duration, 5)); // 5% chance
        }

        // Recalculate total weight
        let total_weight: u32 = duration_options.iter().map(|(_, weight)| *weight).sum();

        // If no duration options are available (should not happen), default to 1 frame
        if total_weight == 0 {
            return 1;
        }
        
        // Generate a random number within the total weight
        let mut choice = rng.gen_range(0..total_weight);
        
        // Select a duration based on weighted probabilities
        for (duration_range, weight) in duration_options {
            if choice < weight {
                return rng.gen_range(duration_range);
            } else {
                choice -= weight;
            }
        }
        
        // Fallback in case of error
        1
    }

    fn generate_weighted_duration_continuous(rng: &mut ThreadRng, max_duration: usize) -> usize {
        // Use exponential distribution to favor shorter durations
        let lambda = 0.1; // Adjust lambda to control decay rate (higher lambda = more short durations)
        let duration = (-rng.gen::<f64>().ln() / lambda).ceil() as usize;
        
        // Ensure duration is at least 1 and does not exceed max_duration
        duration.clamp(1, max_duration)
    }

    pub fn schedule_transitions(&mut self) {
        let rng = &mut self.rng;

        // We need at least two images to create a transition
        if self.primary_images.len() < 2 {
            return;
        }

        // Iterate over pairs of images to create transitions
        for i in 0..(self.primary_images.len() - 1) {
            let (current_state, next_state)= self.primary_images.split_at_mut(i+1);
            let current_state = &mut current_state[current_state.len() - 1];
            let next_state = &mut next_state[0];

            // Decide whether to create a transition between these images
            if rng.gen_bool(0.3) { // 30% chance to create a transition
                // Determine the duration of the transition (e.g., between 1 and 3 seconds)
                let min_transition_duration = self.config.framerate; // 1 second
                let max_transition_duration = self.config.framerate * 3; // 3 seconds
                let transition_duration = rng.gen_range(min_transition_duration..=max_transition_duration);

                // Adjust start and end frames to accommodate the transition
                let transition_midpoint = (current_state.end_frame + next_state.start_frame) / 2;
                let transition_start_frame = transition_midpoint.saturating_sub(transition_duration as usize / 2);
                let transition_end_frame = transition_midpoint.saturating_add(transition_duration as usize / 2);

                // Ensure the frames are within bounds
                let transition_start_frame = transition_start_frame.max(current_state.start_frame);
                let transition_end_frame = transition_end_frame.min(next_state.end_frame);

                // Adjust current and next image durations if necessary
                if current_state.end_frame < transition_end_frame {
                    current_state.end_frame = transition_end_frame;
                }
                if next_state.start_frame > transition_start_frame {
                    next_state.start_frame = transition_start_frame;
                }

                // Resize images to frame dimensions
                let frame_width = self.config.width;
                let frame_height = self.config.height;

                // Resize current image
                let from_image_resized = image::imageops::resize(
                    &current_state.image,
                    frame_width as u32,
                    frame_height as u32,
                    image::imageops::FilterType::Lanczos3,
                );

                // Resize next image
                let to_image_resized = image::imageops::resize(
                    &next_state.image,
                    frame_width as u32,
                    frame_height as u32,
                    image::imageops::FilterType::Lanczos3,
                );

                // Create the transition
                let transition = Transition {
                    from_image: from_image_resized,
                    to_image: to_image_resized,
                    start_frame: transition_start_frame,
                    end_frame: transition_end_frame,
                    duration: transition_end_frame - transition_start_frame,
                    effect: TransitionEffect::Reassembly,
                    particles: Vec::new(), // Particles will be initialized during rendering
                    initialized: false,
                };

                self.transitions.push(transition);
            }
        }
    }
}

////// src/transitions.rs //////

use image::{imageops, Rgba, RgbaImage};
use imageproc::geometric_transformations::{rotate_about_center, Interpolation};
use rand::prelude::*;

use crate::config::Config;

pub struct Transition {
    pub from_image: RgbaImage,
    pub to_image: RgbaImage,
    pub start_frame: usize,
    pub end_frame: usize,
    pub duration: usize,
    pub effect: TransitionEffect,
    pub particles: Vec<Particle>,
    pub initialized: bool, 
}

pub enum TransitionEffect {
    Reassembly,
    // You can add other transition effects in the future
}

// Define the Particle struct
pub struct Particle {
    pub src_x: u32,
    pub src_y: u32,
    pub dst_x: u32,
    pub dst_y: u32,
    pub current_x: f32,
    pub current_y: f32,
    pub block_size: u32,
    pub src_image_data: RgbaImage,
    pub dst_image_data: RgbaImage,
    // Transformations
    pub rotation: f32,
    pub rotation_speed: f32,
    pub scale: f32,
    pub alpha: f32,
}

impl Transition
{
    fn initialize_particles(&mut self, block_size: u32) {
        let from_image = &self.from_image;
        let to_image = &self.to_image;

        let width = from_image.width();
        let height = from_image.height();

        let mut particles = Vec::new();
        let mut rng = rand::thread_rng();

        for y in (0..height).step_by(block_size as usize) {
            for x in (0..width).step_by(block_size as usize) {
                // Calculate block dimensions, ensuring they are at least 1
                let block_width = block_size.min(width - x).max(1);
                let block_height = block_size.min(height - y).max(1);

                // Skip blocks with zero dimensions
                if block_width == 0 || block_height == 0 {
                    continue;
                }

                // Extract block from source image
                let src_block = image::imageops::crop_imm(
                    from_image,
                    x,
                    y,
                    block_size.min(width - x),
                    block_size.min(height - y),
                ).to_image();

                // Extract corresponding block from destination image
                let dst_block = image::imageops::crop_imm(
                    to_image,
                    x,
                    y,
                    block_size.min(width - x),
                    block_size.min(height - y),
                ).to_image();

                // Randomize rotation speed
                let rotation_speed = rng.gen_range(-5.0..5.0);

                // Create the particle
                particles.push(Particle {
                    src_x: x,
                    src_y: y,
                    dst_x: x,
                    dst_y: y,
                    current_x: x as f32,
                    current_y: y as f32,
                    block_size,
                    src_image_data: src_block,
                    dst_image_data: dst_block,
                    rotation: 0.0,
                    rotation_speed,
                    scale: 1.0,
                    alpha: 1.0,
                });
            }
        }

        self.particles = particles;
        self.initialized = true;
    }
}

pub fn apply_reassembly_transition(
    frame_number: usize,
    transition: &mut Transition,
    config: &Config,
) -> RgbaImage {
    let progress = (frame_number - transition.start_frame) as f32 / transition.duration as f32;

    let width = config.width;
    let height = config.height;

    // Initialize particles if not already done
    if !transition.initialized {
        let block_size = 16; // Adjust as needed
        transition.initialize_particles(block_size);
    }

    let mut frame = RgbaImage::new(width as u32, height as u32);

    for particle in &mut transition.particles {
        // Update transformations
        particle.rotation += particle.rotation_speed;
        particle.scale = (1.0 + 0.5 * (1.0 - progress)).max(0.1); // Ensure scale isn't too small

        // Interpolate position
        particle.current_x = particle.src_x as f32 + ((particle.dst_x as f32 - particle.src_x as f32) * progress);
        particle.current_y = particle.src_y as f32 + ((particle.dst_y as f32 - particle.src_y as f32) * progress);

        // Interpolate alpha for blending
        particle.alpha = 1.0 - ((progress * 2.0).min(1.0));

        // Blend source and destination blocks based on progress
        let blended_block = blend_blocks(
            &particle.src_image_data,
            &particle.dst_image_data,
            progress,
            particle.rotation,
            particle.scale,
            particle.alpha,
        );

        // // Use destination block after halfway point
        // if progress > 0.5 && particle.alpha < 0.5 {
        //     // Replace image data with the destination block
        //     particle.image_data = image::imageops::crop_imm(
        //         &transition.to_image,
        //         particle.dst_x,
        //         particle.dst_y,
        //         particle.block_size.min(width as u32 - particle.dst_x),
        //         particle.block_size.min(height as u32 - particle.dst_y),
        //     ).to_image();
        //     particle.alpha = progress; // Increase opacity
        // }

        // // Apply transformations to the block
        // let transformed_block = transform_block(
        //     &particle.image_data,
        //     particle.rotation,
        //     particle.scale,
        //     particle.alpha,
        // );

        // Overlay the transformed block onto the frame
        imageops::overlay(
            &mut frame,
            &blended_block,
            particle.current_x as i64,
            particle.current_y as i64,
        );
    }

    frame
}

fn blend_blocks(
    src_block: &RgbaImage,
    dst_block: &RgbaImage,
    progress: f32,
    rotation: f32,
    scale: f32,
    alpha: f32,
) -> RgbaImage {
    // Ensure block dimensions match
    let (width, height) = src_block.dimensions();

    // Blend the pixels
    let mut blended_block = RgbaImage::new(width, height);
    for (x, y, pixel) in blended_block.enumerate_pixels_mut() {
        let src_pixel = src_block.get_pixel(x, y);
        let dst_pixel = dst_block.get_pixel(x, y);

        // Interpolate between src and dst pixels based on progress
        let blended_pixel = Rgba([
            (src_pixel[0] as f32 * (1.0 - progress) + dst_pixel[0] as f32 * progress) as u8,
            (src_pixel[1] as f32 * (1.0 - progress) + dst_pixel[1] as f32 * progress) as u8,
            (src_pixel[2] as f32 * (1.0 - progress) + dst_pixel[2] as f32 * progress) as u8,
            (src_pixel[3] as f32 * (1.0 - progress) + dst_pixel[3] as f32 * progress) as u8,
        ]);

        *pixel = blended_pixel;
    }

    // Apply transformations
    let transformed_block = transform_block(&blended_block, rotation, scale, alpha);

    transformed_block
}

fn transform_block(
    block: &RgbaImage,
    rotation: f32,
    scale: f32,
    alpha: f32,
) -> RgbaImage {
    let (width, height) = block.dimensions();

    // Early exit if width or height is zero
    if width == 0 || height == 0 {
        // Can't rotate or scale zero-dimension image
        return RgbaImage::new(1, 1); // Return a minimal image
    }

    // Apply rotation
    let rotated = rotate_about_center(
        block,
        rotation.to_radians(),
        Interpolation::Nearest,
        Rgba([0, 0, 0, 0]),
    );

    // Apply scaling
    let new_width = (width as f32 * scale).max(1.0) as u32;
    let new_height = (height as f32 * scale).max(1.0) as u32;

    // Skip if new dimensions are zero
    if new_width == 0 || new_height == 0 {
        return RgbaImage::new(1, 1); // Return a minimal image
    }

    let scaled = imageops::resize(
        &rotated,
        new_width,
        new_height,
        imageops::FilterType::Nearest, // TODO: Try nearest too
    );

    // Apply alpha
    let mut final_block = scaled.clone();
    for pixel in final_block.pixels_mut() {
        let new_alpha = (pixel[3] as f32 * alpha).clamp(0.0, 255.0) as u8;
        pixel[3] = new_alpha;
    }

    final_block
}

fn blend_pixels(
    base: &image::Rgba<u8>,
    overlay: &image::Rgba<u8>,
    alpha: f32,
) -> image::Rgba<u8> {
    let inv_alpha = 1.0 - alpha;
    image::Rgba([
        (base[0] as f32 * inv_alpha + overlay[0] as f32 * alpha) as u8,
        (base[1] as f32 * inv_alpha + overlay[1] as f32 * alpha) as u8,
        (base[2] as f32 * inv_alpha + overlay[2] as f32 * alpha) as u8,
        255,
    ])
}

// Helper function to get pixel safely, returns black if out of bounds
pub fn get_pixel_safe(image: &RgbaImage, x: u32, y: u32) -> image::Rgba<u8> {
    if x < image.width() && y < image.height() {
        *image.get_pixel(x, y)
    } else {
        image::Rgba([0, 0, 0, 255])
    }
}

////// src/types.rs //////

use image::RgbaImage;

#[derive(Clone)]
pub struct ImageState {
    pub image: RgbaImage,
    pub position: (i64, i64),
    pub scale: f32,
    pub start_frame: usize,
    pub end_frame: usize,
    pub effects: Vec<crate::effects::Effect>,
    // Add any additional fields as needed
}

////// src/video_encoder.rs //////

use crate::config::Config;
use anyhow::Result;
use image::RgbaImage;
use rav1e::prelude::*;
use std::fs::File;
use std::io::Write;

pub struct VideoEncoder {
    ctx: Context<u8>,
    output: File,
}

impl VideoEncoder {
    pub fn new(config: &Config) -> Result<Self> {
        // Initialize rav1e encoder
        let mut encoder_cfg = EncoderConfig::default();
        encoder_cfg.width = config.width as usize;
        encoder_cfg.height = config.height as usize;
        encoder_cfg.time_base = Rational::new(config.timebase_num as u64, config.timebase_den as u64);
        encoder_cfg.quantizer = config.quantizer; // Adjust quality (0-255)
        encoder_cfg.speed_settings = SpeedSettings::from_preset(config.speed as u8);
        encoder_cfg.bit_depth = config.bit_depth;
        let cfg = rav1e::Config::default().with_encoder_config(encoder_cfg);
        let ctx: Context<u8> = cfg.new_context()?;

        // Open the output file and write IVF header
        let mut output = File::create(&config.output_file)?;
        write_ivf_header(
            &mut output,
            config.width as u32,
            config.height as u32,
            config.timebase_num,
            config.timebase_den,
        )?;

        Ok(Self { ctx, output })
    }

    pub fn encode_frame(&mut self, img: &RgbaImage) -> Result<()> {
        // Convert image to rav1e frame
        let rav1e_frame = {
            let mut frame = self.ctx.new_frame();
            image_to_frame(&img, &mut frame)?;
            frame
        };

        // Send the frame to the encoder
        self.ctx.send_frame(rav1e_frame)?;

        // Receive packets
        while let Some(packet) = self.ctx.receive_packet().ok() {
            write_ivf_frame(&mut self.output, &packet)?;
        }

        Ok(())
    }

    pub fn finish(&mut self) -> Result<()> {
        // Flush the encoder
        self.ctx.flush();

        // Receive and write remaining packets
        while let Some(packet) = self.ctx.receive_packet().ok() {
            write_ivf_frame(&mut self.output, &packet)?;
        }

        Ok(())
    }
}

// Helper function to convert image to rav1e frame
fn image_to_frame(img: &RgbaImage, frame: &mut Frame<u8>) -> anyhow::Result<()> {
    let (width, height) = img.dimensions();

    // Ensure the frame dimensions match the image dimensions
    // Y Plane
    let y_plane = &mut frame.planes[0];
    let y_width = y_plane.cfg.width;
    let y_stride = y_plane.cfg.stride;
    for y in 0..(height as usize) {
        let y_row = &mut y_plane.data_origin_mut()[y * y_stride..][..y_width];
        for x in 0..(width as usize) {
            let pixel = img.get_pixel(x as u32, y as u32);
            let r = pixel[0] as f32 / 255.0;
            let g = pixel[1] as f32 / 255.0;
            let b = pixel[2] as f32 / 255.0;
            let y_val = (0.299 * r + 0.587 * g + 0.114 * b) * 255.0;
            y_row[x] = y_val as u8;
        }
    }

    // U Plane (Subsampled)
    let u_plane = &mut frame.planes[1];
    let u_width = u_plane.cfg.width;
    let u_stride = u_plane.cfg.stride;
    for y in 0..(height as usize / 2) {
        let u_row = &mut u_plane.data_origin_mut()[y * u_stride..][..u_width];
        for x in 0..(width as usize / 2) {
            let mut r = 0.0;
            let mut g = 0.0;
            let mut b = 0.0;
            for dy in 0..2 {
                for dx in 0..2 {
                    let px = (2 * x + dx) as u32;
                    let py = (2 * y + dy) as u32;
                    if px < width && py < height {
                        let pixel = img.get_pixel(px, py);
                        r += pixel[0] as f32 / 255.0;
                        g += pixel[1] as f32 / 255.0;
                        b += pixel[2] as f32 / 255.0;
                    }
                }
            }
            r /= 4.0;
            g /= 4.0;
            b /= 4.0;
            let u_val = ((-0.169 * r - 0.331 * g + 0.5 * b) + 0.5) * 255.0;
            u_row[x] = u_val as u8;
        }
    }

    // V Plane (Subsampled)
    let v_plane = &mut frame.planes[2];
    let v_width = v_plane.cfg.width;
    let v_stride = v_plane.cfg.stride;
    for y in 0..(height as usize / 2) {
        let v_row = &mut v_plane.data_origin_mut()[y * v_stride..][..v_width];
        for x in 0..(width as usize / 2) {
            let mut r = 0.0;
            let mut g = 0.0;
            let mut b = 0.0;
            for dy in 0..2 {
                for dx in 0..2 {
                    let px = (2 * x + dx) as u32;
                    let py = (2 * y + dy) as u32;
                    if px < width && py < height {
                        let pixel = img.get_pixel(px, py);
                        r += pixel[0] as f32 / 255.0;
                        g += pixel[1] as f32 / 255.0;
                        b += pixel[2] as f32 / 255.0;
                    }
                }
            }
            r /= 4.0;
            g /= 4.0;
            b /= 4.0;
            let v_val = ((0.5 * r - 0.419 * g - 0.081 * b) + 0.5) * 255.0;
            v_row[x] = v_val as u8;
        }
    }

    Ok(())
}

// Function to write IVF header
fn write_ivf_header<W: Write>(
    w: &mut W,
    width: u32,
    height: u32,
    timebase_num: u32,
    timebase_den: u32,
) -> std::io::Result<()> {
    // IVF header is 32 bytes
    w.write_all(b"DKIF")?;                              // 0-3   file signature
    w.write_all(&0u16.to_le_bytes())?;                  // 4-5   version
    w.write_all(&32u16.to_le_bytes())?;                 // 6-7   header size
    w.write_all(b"AV01")?;                              // 8-11  fourcc
    w.write_all(&(width as u16).to_le_bytes())?;        // 12-13 width
    w.write_all(&(height as u16).to_le_bytes())?;       // 14-15 height
    w.write_all(&timebase_den.to_le_bytes())?;          // 16-19 timebase denominator (little endian)
    w.write_all(&timebase_num.to_le_bytes())?;          // 20-23 timebase numerator (little endian)
    w.write_all(&0u32.to_le_bytes())?;                  // 24-27 frame count (unknown)
    w.write_all(&0u32.to_le_bytes())?;                  // 28-31 unused
    Ok(())
}

// Function to write IVF frame
fn write_ivf_frame<W: Write>(w: &mut W, packet: &Packet<u8>) -> std::io::Result<()> {
    let size = packet.data.len() as u32;
    let timestamp = packet.input_frameno as u64;

    w.write_all(&size.to_le_bytes())?;
    w.write_all(&timestamp.to_le_bytes())?;
    w.write_all(&packet.data)?;

    Ok(())
}