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. Here is the source code:

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

use serde::Deserialize;
use std::{fs, sync::LazyLock};

// TODO: Lock here might not be good, consider a RwLock kind of thing
pub static CONFIG: LazyLock<Config> = LazyLock::new(|| {
    let config_str = fs::read_to_string("config.toml")
        .expect(r"C:\Project\Workspace\Javascript\gpt4o-site\media-cycler-rs\config.toml");
    toml::from_str(&config_str)
        .expect("Failed to parse config.toml")
});

#[test]
fn test_deser()
{
    println!("{:#?}", *CONFIG)
}

#[derive(Debug, Deserialize)]
pub struct Config {
    pub output_file: String,
    pub input_dirs: Vec<String>,
    pub duration_s: u32,
    pub width: u32,
    pub height: u32,
    pub framerate: u32,
    pub quantizer: u32,
    pub speed: u32,
    pub bit_depth: u32,
    pub effects: Effects,
}

#[derive(Debug, Deserialize)]
pub struct Effects {
    pub invert_colors: InvertColorsEffect,
    pub color_flash: ColorFlashEffect,
    pub sparkle: SparkleEffect,
    pub spiral: SpiralEffect,
    pub transparency_strobe: TransparencyStrobeEffect,
    pub overlay_text: OverlayTextEffect,
    pub zoom_in: ZoomInEffect,
    pub kaleidoscope: KaleidoscopeEffect,
    pub bit_focus: BitFocusEffect,
    pub glow: GlowEffect,
    pub grow_shrink: GrowShrinkEffect,
    pub fade_out: FadeOutEffect,
    pub split_apart: SplitApartEffect,
    pub randomized_oval_fade: RandomizedOvalFadeEffect,
    pub glitch: GlitchEffect,
}

#[derive(Debug, Deserialize)]
pub struct InvertColorsEffect {
    pub weight: f64,
}

#[derive(Debug, Deserialize)]
pub struct ColorFlashEffect {
    pub weight: f64,
}

#[derive(Debug, Deserialize)]
pub struct SparkleEffect {
    pub weight: f64,
}

#[derive(Debug, Deserialize)]
pub struct SpiralEffect {
    pub weight: f64,
}


#[derive(Debug, Deserialize)]
pub struct TransparencyStrobeEffect {
    pub weight: f64,
}


#[derive(Debug, Deserialize)]
pub struct OverlayTextEffect {
    pub weight: f64,
    pub texts: Vec<String>
}


#[derive(Debug, Deserialize)]
pub struct ZoomInEffect {
    pub weight: f64,
}


#[derive(Debug, Deserialize)]
pub struct KaleidoscopeEffect {
    pub weight: f64,
}


#[derive(Debug, Deserialize)]
pub struct BitFocusEffect {
    pub weight: f64,
}


#[derive(Debug, Deserialize)]
pub struct GlowEffect {
    pub weight: f64,
}


#[derive(Debug, Deserialize)]
pub struct GrowShrinkEffect {
    pub weight: f64,
}


#[derive(Debug, Deserialize)]
pub struct FadeOutEffect {
    pub weight: f64,
}


#[derive(Debug, Deserialize)]
pub struct SplitApartEffect {
    pub weight: f64,
}

#[derive(Debug, Deserialize)]
pub struct RandomizedOvalFadeEffect {
    pub weight: f64,
}

#[derive(Debug, Deserialize)]
pub struct GlitchEffect {
    pub weight: f64,
}

impl Config
{
    pub fn total_frames(&self) -> usize { (self.framerate * self.duration_s) as usize }
    pub fn timebase_num(&self) -> u32 { 1 }
    pub fn timebase_den(&self) -> u32 { self.framerate }
}


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

use image::{GrayImage, ImageBuffer, Pixel, Rgba, RgbaImage};
use imageproc::{filter::Kernel, morphology::erode_mut};
use rand::prelude::*;
use rusttype::Font;
use noise::{Fbm, MultiFractal, NoiseFn, Perlin, RidgedMulti};
use image::Luma;

use crate::{config::CONFIG, 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,
    },
    RandomizedOvalFade {
        randomness: f32, // Controls the amount of shape distortion
    },
    Glitch,
}


// EffectFactory to generate random effects
pub struct EffectFactory {
    pub rng: ThreadRng,
    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(),
            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(CONFIG.effects.invert_colors.weight) {
            effects.push(Effect::InvertColors);
        }
        if self.rng.gen_bool(CONFIG.effects.color_flash.weight) {
            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(CONFIG.effects.sparkle.weight) {
            effects.push(Effect::Sparkle);
        }
        if self.rng.gen_bool(CONFIG.effects.spiral.weight) {
            effects.push(Effect::Spiral);
        }
        if self.rng.gen_bool(CONFIG.effects.transparency_strobe.weight) {
            effects.push(Effect::TransparencyStrobe(self.rng.gen_range(0.2..=10.0)));
        }
        if self.rng.gen_bool(CONFIG.effects.overlay_text.weight) {
            let text = CONFIG.effects.overlay_text.texts.choose(&mut self.rng).unwrap().to_string();
            effects.push(Effect::OverlayText(text));
        }
        if self.rng.gen_bool(CONFIG.effects.zoom_in.weight) {
            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(CONFIG.effects.kaleidoscope.weight) {
            effects.push(Effect::Kaleidoscope);
        }
        if self.rng.gen_bool(CONFIG.effects.bit_focus.weight) {
            effects.push(Effect::BitFocus);
        }
        if self.rng.gen_bool(CONFIG.effects.glow.weight) {
            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(CONFIG.effects.grow_shrink.weight) {
            let speed = self.rng.gen_range(0.01..0.05);
            effects.push(Effect::GrowShrink { speed });
        }
        if self.rng.gen_bool(CONFIG.effects.fade_out.weight) {
            effects.push(Effect::FadeOut);
        }
        if self.rng.gen_bool(CONFIG.effects.split_apart.weight) {
            effects.push(Effect::SplitApart);
        }
        if self.rng.gen_bool(CONFIG.effects.sparkle.weight) {
            effects.push(Effect::Sparkle);
        }
        if self.rng.gen_bool(CONFIG.effects.spiral.weight) {
            effects.push(Effect::Spiral);
        }
        if self.rng.gen_bool(CONFIG.effects.glow.weight) {
            effects.push(Effect::Glow{max_intensity: self.rng.gen_range(1.5..3.0)});
        }
        if self.rng.gen_bool(CONFIG.effects.bit_focus.weight) {
            effects.push(Effect::BitFocus);
        }
        if self.rng.gen_bool(CONFIG.effects.color_flash.weight) {
            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(CONFIG.effects.invert_colors.weight / 2.0) {
            effects.push(Effect::InvertColors);
        }
        if self.rng.gen_bool(CONFIG.effects.randomized_oval_fade.weight) {
            let randomness = self.rng.gen_range(2.0..5.0);
            effects.push(Effect::RandomizedOvalFade { randomness });
        }
        if self.rng.gen_bool(CONFIG.effects.glitch.weight) {
            effects.push(Effect::Glitch);
        }

        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 = CONFIG.effects.overlay_text.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,
                    );
                },
                Effect::RandomizedOvalFade { randomness } => {
                    self.apply_randomized_oval_fade(image, *randomness);
                },
                Effect::Glitch => {
                    self.apply_glitch_effect(image);
                },
            }
        }
        // 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
        }
    }

    fn apply_randomized_oval_fade(&self, image: &mut RgbaImage, randomness: f32) {
        let fbm = Fbm::new()
            .set_octaves(6)
            .set_frequency(2.0)
            .set_lacunarity(2.0)
            .set_persistence(0.5);

        let ridged = RidgedMulti::new()
            .set_octaves(4)
            .set_frequency(1.0)
            .set_lacunarity(2.0)
            .set_persistence(0.5);

        let width = image.width() as f32;
        let height = image.height() as f32;
        let center_x = width / 2.0;
        let center_y = height / 2.0;
        let max_distance = ((center_x.powi(2) + center_y.powi(2)).sqrt()) as f32;
        let noise_scale = randomness * 5.0;
        let noise_threshold = 0.2;

        for y in 0..image.height() {
            for x in 0..image.width() {
                let dx = x as f32 - center_x;
                let dy = y as f32 - center_y;
                let distance = (dx.powi(2) + dy.powi(2)).sqrt();
                let distance_factor = (distance / max_distance).clamp(0.0, 1.0);
    
                // Get combined noise value
                let noise_value1 = fbm.get([
                    dx as f64 / width as f64 * noise_scale as f64,
                    dy as f64 / height as f64 * noise_scale as f64,
                ]) as f32;
    
                let noise_value2 = ridged.get([
                    dx as f64 / width as f64 * noise_scale as f64 * 0.5,
                    dy as f64 / height as f64 * noise_scale as f64 * 0.5,
                ]) as f32;
    
                // Combine noises
                let combined_noise = (noise_value1 + noise_value2) / 2.0;
    
                // Apply threshold
                let mask_value = if combined_noise > noise_threshold { 1.0 } else { 0.0 };
    
                let alpha = (1.0 - distance_factor.powf(0.5)) * mask_value;
                let edge_alpha = if distance_factor >= 1.0 { 0.0 } else { alpha };
    
                let pixel = image.get_pixel_mut(x, y);
                let color_shift = ((combined_noise - 0.5) * 2.0 * randomness).clamp(-0.5, 0.5);
                for i in 0..3 {
                    let color = pixel[i] as f32 / 255.0;
                    let shifted = (color + color_shift).clamp(0.0, 1.0);
                    pixel[i] = (shifted * 255.0) as u8;
                }
                pixel[3] = (pixel[3] as f32 * edge_alpha) as u8;
            }
        }
    }

    // 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
    }

    fn apply_glitch_effect(&mut self, image: &mut RgbaImage) {
        let width = image.width();
        let height = image.height();
        let shift = (self.rng.gen_range(-10..10), self.rng.gen_range(-10..10));
        
        // Clone the image to avoid borrowing issues
        let original = image.clone();
    
        for y in 0..height {
            for x in 0..width {
                let src_x = x.wrapping_add(shift.0 as u32) % width;
                let src_y = y.wrapping_add(shift.1 as u32) % height;
                let src_pixel = original.get_pixel(src_x, src_y);
    
                let pixel = image.get_pixel_mut(x, y);
                // Shift the red channel
                pixel[0] = src_pixel[0];
            }
        }
    }

    // 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) {
        if image.width() == 0 || image.height() == 0 
        {
            return;
        }
        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;

        if image.width() == 0 || image.height() == 0 { 
            return;
        }

        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::config::CONFIG;
use crate::transitions::{apply_reassembly_transition, TransitionEffect};
use crate::{config::Config, transitions::Transition};
use crate::effects::EffectFactory;
use crate::types::{ImageFrames, ImageState};
use image::{imageops, AnimationDecoder, DynamicImage, RgbaImage};
use std::collections::HashMap;
use std::fs::File;
use std::path::PathBuf;
use std::cell::RefCell;

pub fn render_frame(
    frame_number: usize,
    active_primary_images: &[ImageState],
    active_overlay_images: &[ImageState],
    active_global_effects: &[ImageState],
    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  {
        // Obtain the current frame
        let current_frame_image = state.frames.get_current_frame(frame_number, state.start_frame, state.end_frame).unwrap_or_else(|| RgbaImage::new(1, 1));

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

        let mut img = img.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 {
        // Similar handling as above
        let current_frame_image = state.frames.get_current_frame(frame_number, state.start_frame, state.end_frame).unwrap_or_else(|| RgbaImage::new(1, 1));

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

        let mut img = img.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,
                    );
                    // Overlay the transition_frame onto the main frame
                    image::imageops::overlay(&mut frame, &transition_frame, 0, 0);
                },
                // Handle other transition effects...
            }
        }
    }

    frame
}

// // Function to get the appropriate frame from an animated GIF
// fn get_gif_frame(
//     path: &PathBuf,
//     frame_number: usize,
//     start_frame: usize,
//     end_frame: usize,
//     _total_duration_ms: u64,
//     frame_count: usize,
//     frame_cache: &mut HashMap<(PathBuf, usize), RgbaImage>,
// ) -> Option<RgbaImage> {
//     let relative_frame = frame_number - start_frame;
//     let total_duration_frames = end_frame - start_frame;

//     // Calculate the corresponding frame index in the GIF
//     let progress = relative_frame as f64 / total_duration_frames as f64;
//     let gif_frame_index = (progress * frame_count as f64).floor() as usize % frame_count;

//     // Check if frame is in cache
//     let cache_key = (path.clone(), gif_frame_index);
//     if let Some(cached_frame) = frame_cache.get(&cache_key) {
//         return Some(cached_frame.clone());
//     }

//     // Load the frame
//     if let Ok(file) = File::open(path) {
//         if let Ok(decoder) = image::codecs::gif::GifDecoder::new(file) {
//             let frames = decoder.into_frames();
//             let mut frames_iter = frames.into_iter();

//             // Advance to the desired frame
//             for current_frame_index in 0..=gif_frame_index {
//                 if let Some(frame_result) = frames_iter.next() {
//                     if current_frame_index == gif_frame_index {
//                         if let Ok(frame) = frame_result {
//                             let image = frame.into_buffer();

//                             // Cache the frame
//                             frame_cache.insert(cache_key, image.clone());

//                             return Some(image);
//                         }
//                     }
//                 } else {
//                     break;
//                 }
//             }
//         }
//     }

//     None
// }
////// src/image_loader.rs //////

use anyhow::{Error, Result};
use image::{self, RgbaImage};
use image::AnimationDecoder;
use std::fs::{self, File};
use std::path::PathBuf;
use ffmpeg_next::{self as ffmpeg, codec, format, frame, media, Rational};

pub enum Media {
    Static(RgbaImage),
    Animated {
        path: PathBuf,
        total_duration_ms: u64, // total duration in milliseconds
        frame_count: usize,
    },
    Video {
        path: PathBuf,
        total_duration_ms: u64,
        frame_rate: f64,
        width: u32,
        height: u32,
    },
}

struct VideoMetadata {
    duration_ms: u64,
    frame_rate: f64,
    width: u32,
    height: u32,
}

pub struct VideoClip {
    pub path: PathBuf,
    ictx: format::context::Input,
    decoder: codec::decoder::Video,
    stream_index: usize,
    scaler: ffmpeg_next::software::scaling::Context,
    next_packet: Option<ffmpeg_next::Packet>,
    pending_frames: Vec<frame::Video>,
    frame_finished: bool,
    pub total_duration_ms: u64,
    pub frame_rate: f64,
    pub width: u32,
    pub height: u32,
    pub current_pts: i64,
    pub time_base: Rational,
}

impl VideoClip {
    pub fn open(path: &PathBuf) -> Result<Self> {
        // Initialize FFmpeg and open the input file
        ffmpeg_next::init()?;
        let ictx = format::input(path)?;
        let input_stream = ictx
            .streams()
            .best(media::Type::Video)
            .ok_or_else(|| anyhow::anyhow!("Could not find video stream"))?;
        let stream_index = input_stream.index();

        // Get codec parameters and open the decoder
        let context_decoder = codec::Context::from_parameters(input_stream.parameters())?;
        let decoder = context_decoder.decoder().video()?;

        // Create a scaler to convert frames to RGBA
        let scaler = ffmpeg_next::software::scaling::Context::get(
            decoder.format(),
            decoder.width(),
            decoder.height(),
            ffmpeg_next::format::Pixel::RGBA,
            decoder.width(),
            decoder.height(),
            ffmpeg_next::software::scaling::flag::Flags::BILINEAR,
        )?;

        // Get video metadata
        let duration_ms = input_stream.duration() as u64;
        let frame_rate = input_stream.avg_frame_rate().0 as f64 / input_stream.avg_frame_rate().1 as f64;
        let width = decoder.width();
        let height = decoder.height();
        let time_base = input_stream.time_base();

        Ok(Self {
            path: path.clone(),
            ictx,
            decoder,
            stream_index,
            scaler,
            next_packet: None,
            pending_frames: Vec::new(),
            frame_finished: false,
            total_duration_ms: duration_ms,
            frame_rate,
            width: width as u32,
            height: height as u32,
            current_pts: 0,
            time_base,
        })
    }

    pub fn get_next_frame(&mut self) -> Result<Option<RgbaImage>> {
        // If we have frames pending from previous decode
        if let Some(frame) = self.pending_frames.pop() {
            self.current_pts += 1;
            return self.convert_frame_to_image(frame);
        }

        // Read packets and decode frames
        while let Some((stream, packet)) = self.ictx.packets().next() {
            if stream.index() == self.stream_index {
                self.decoder.send_packet(&packet)?;
                let mut frame = ffmpeg::util::frame::Video::empty();
                while let Ok(()) = self.decoder.receive_frame(&mut frame) {
                    self.current_pts += 1;
                    return self.convert_frame_to_image(frame);
                }
            }
        }

        // No more packets; flush the decoder
        self.decoder.send_eof()?;
        let mut frame = ffmpeg::util::frame::Video::empty();
        while let Ok(()) = self.decoder.receive_frame(&mut frame) {
            self.current_pts += 1;
            return self.convert_frame_to_image(frame);
        }

        Ok(None)
    }

    fn convert_frame_to_image(&mut self, frame: frame::Video) -> Result<Option<RgbaImage>> {
        let mut rgb_frame = frame::Video::empty();
        self.scaler.run(&frame, &mut rgb_frame)?;

        let data = rgb_frame.data(0);
        let width = rgb_frame.width() as u32;
        let height = rgb_frame.height() as u32;
        let stride = rgb_frame.stride(0);

        // Create RgbaImage
        let mut image = RgbaImage::new(width, height);
        let buf = image.as_mut();

        // Copy data from ffmpeg frame to image buffer
        for y in 0..height as usize {
            let src_start = y * stride as usize;
            let src_end = src_start + (width as usize * 4); // 4 bytes per pixel (RGBA)
            let src_row = &data[src_start..src_end];

            let dst_start = y * (width as usize * 4);
            let dst_end = dst_start + (width as usize * 4);

            buf[dst_start..dst_end].copy_from_slice(src_row);
        }

        Ok(Some(image))
    }
}

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

    // Initialize ffmpeg
    ffmpeg::init().unwrap();

    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" => {
                        // Load static image
                        if let Ok(img) = image::open(&path) {
                            media.push(Media::Static(img.to_rgba8()));
                        }
                    },
                    "gif" => {
                        // Load animated GIF metadata only
                        if let Ok(file) = File::open(&path) {
                            if let Ok(decoder) = image::codecs::gif::GifDecoder::new(file) {
                                let frames = decoder.into_frames();
                                let mut total_duration_ms = 0u64;
                                let mut frame_count = 0usize;

                                // We need to get frame durations without loading frames into memory
                                for frame_result in frames.into_iter() {
                                    if let Ok(frame) = frame_result {
                                        let delay = frame.delay().numer_denom_ms().0;
                                        total_duration_ms += delay as u64;
                                        frame_count += 1;
                                    } else {
                                        break;
                                    }
                                }

                                // Store the path and minimal metadata
                                media.push(Media::Animated {
                                    path: path.clone(),
                                    total_duration_ms,
                                    frame_count,
                                });
                            }
                        }
                    },
                    "mp4" | "mov" | "avi" | "mkv" => {
                        // Load video files
                        let path_clone = path.clone();
                        if let Ok(metadata) = get_video_metadata(&path) {
                            media.push(Media::Video {
                                path: path_clone,
                                total_duration_ms: metadata.duration_ms,
                                frame_rate: metadata.frame_rate,
                                width: metadata.width,
                                height: metadata.height,
                            });
                        }
                    },
                    _ => {
                        println!("Unsupported file type: {:?}", path.extension());
                    }
                }
            }
        }
    }

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

fn get_video_metadata(path: &PathBuf) -> Result<VideoMetadata> {
    let ictx = ffmpeg::format::input(&path)?;
    let input: ffmpeg_next::Stream<'_> = ictx
        .streams()
        .best(ffmpeg::media::Type::Video)
        .ok_or_else(|| ffmpeg::Error::StreamNotFound)?;
    let codec_params = input.parameters();
    let (width, height) = if codec_params.medium() == ffmpeg::media::Type::Video {
        let width = unsafe {(*codec_params.as_ptr()).width} as u32;
        let height = unsafe {(*codec_params.as_ptr()).height} as u32;
        (width, height)
    } else { 
        return Err(Error::msg("Not a video"))
    };
    let time_base = input.time_base();
    let duration = input.duration();
    let frame_rate = input.avg_frame_rate();

    let duration_ms = ((duration as f64 * time_base.numerator() as f64 / time_base.denominator() as f64) * 1000.0) as u64;
    let frame_rate = frame_rate.numerator() as f64 / frame_rate.denominator() as f64;

    Ok(VideoMetadata {
        duration_ms,
        frame_rate,
        width,
        height,
    })
}

pub fn get_video_frame_at_timestamp(mp4_video_path: &PathBuf, timestamp_secs: u64) -> Result<RgbaImage> {
    // Initialize ffmpeg
    ffmpeg::init()?;

    // Open the input file
    let mut ictx = ffmpeg::format::input(mp4_video_path)?;

    // Find the best video stream
    let input_stream = ictx
        .streams()
        .best(ffmpeg::media::Type::Video)
        .ok_or_else(|| Error::msg("Could not find video stream"))?;

    let video_stream_index = input_stream.index();

    // Get codec parameters and open the decoder
    let codec_params = input_stream.parameters();
    let context_decoder = ffmpeg::codec::context::Context::from_parameters(codec_params)?;
    let mut decoder = context_decoder.decoder().video()?;

    // Get the stream's time base
    let tb = input_stream.time_base();

    // Convert the timestamp from seconds to the stream's time base units
    use ffmpeg_next::util::mathematics::rescale::Rescale;
    use ffmpeg::util::rational::Rational;

    // Check for timestamp overflow when casting to i64
    if timestamp_secs > i64::MAX as u64 {
        return Err(Error::msg("Timestamp is too large"));
    }
    let timestamp_secs_i64 = timestamp_secs as i64;

    let timestamp_in_stream_tb = timestamp_secs_i64.rescale(Rational(1, 1), tb);

    // Use av_seek_frame via FFI to seek to the desired timestamp
    unsafe {
        let ictx_ptr = ictx.as_mut_ptr();

        // Use AVSEEK_FLAG_BACKWARD to seek to the closest keyframe before the desired timestamp
        let ret = ffmpeg::ffi::av_seek_frame(
            ictx_ptr,
            video_stream_index as i32,
            timestamp_in_stream_tb,
            ffmpeg::ffi::AVSEEK_FLAG_BACKWARD,
        );
        if ret < 0 {
            return Err(Error::msg("Error seeking"));
        }
    }

    // Flush the decoder to clear any buffered frames
    decoder.flush();

    // Create a scaling context to convert frames to RGBA
    use ffmpeg::software::scaling::context::Context as ScalingContext;
    use ffmpeg::software::scaling::flag::Flags;

    let decoder_format = get_modern_pixel_format(decoder.format());
    let mut scaler = ScalingContext::get(
        decoder_format,
        decoder.width(),
        decoder.height(),
        ffmpeg::format::Pixel::RGBA,
        decoder.width(),
        decoder.height(),
        Flags::BILINEAR,
    )?;

    // Prepare frames
    let mut frame = ffmpeg::util::frame::Video::empty();
    let mut rgb_frame = ffmpeg::util::frame::Video::empty();

    // Read packets and decode frames
    for (stream, packet) in ictx.packets() {
        if stream.index() == video_stream_index {
            // Send the packet to the decoder
            decoder.send_packet(&packet)?;

            // Receive and process decoded frames
            while decoder.receive_frame(&mut frame).is_ok() {
                // Set color range for deprecated formats
                match decoder.format() {
                    ffmpeg::format::Pixel::YUVJ420P |
                    ffmpeg::format::Pixel::YUVJ422P |
                    ffmpeg::format::Pixel::YUVJ444P |
                    ffmpeg::format::Pixel::YUVJ440P => {
                        unsafe {
                            (*frame.as_mut_ptr()).color_range = ffmpeg::ffi::AVColorRange::AVCOL_RANGE_JPEG;
                        }
                    },
                    _ => {},
                }
                // Check if the frame PTS is greater than or equal to the desired timestamp
                if let Some(frame_pts) = frame.pts() {
                    if frame_pts >= timestamp_in_stream_tb {
                        // Convert the frame to RGBA format
                        scaler.run(&frame, &mut rgb_frame)?;

                        // Extract frame data
                        let data = rgb_frame.data(0);
                        let width = rgb_frame.width() as u32;
                        let height = rgb_frame.height() as u32;
                        let stride = rgb_frame.stride(0);

                        // Create RgbaImage
                        let mut image = RgbaImage::new(width, height);
                        let buf = image.as_mut();

                        // Copy data from ffmpeg frame to image buffer
                        for y in 0..height as usize {
                            let src_start = y * stride as usize;
                            let src_end = src_start + (width as usize * 4); // 4 bytes per pixel (RGBA)
                            let src_row = &data[src_start..src_end];

                            let dst_start = y * (width as usize * 4);
                            let dst_end = dst_start + (width as usize * 4);

                            buf[dst_start..dst_end].copy_from_slice(src_row);
                        }

                        return Ok(image);
                    }
                }
            }
        }
    }

    Err(Error::msg("Frame not found"))
}

fn get_modern_pixel_format(format: ffmpeg::format::Pixel) -> ffmpeg::format::Pixel {
    match format {
        ffmpeg::format::Pixel::YUVJ420P => ffmpeg::format::Pixel::YUV420P,
        ffmpeg::format::Pixel::YUVJ422P => ffmpeg::format::Pixel::YUV422P,
        ffmpeg::format::Pixel::YUVJ444P => ffmpeg::format::Pixel::YUV444P,
        ffmpeg::format::Pixel::YUVJ440P => ffmpeg::format::Pixel::YUV440P,
        _ => format,
    }
}
////// src/main.rs //////

use config::CONFIG;
use indicatif::{ProgressBar, ProgressStyle};

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

fn main() -> anyhow::Result<()> {
    // 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(&CONFIG.input_dirs).unwrap();

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

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

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

    // 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,
            &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 std::cell::RefCell;
use std::rc::Rc;

use crate::config::CONFIG;
use crate::effects::{self, Effect};
use crate::image_loader::{Media, VideoClip};
use crate::transitions::{Transition, TransitionEffect};
use crate::types::{ImageFrames, ImageState};
use image::RgbaImage;
use rand::prelude::*;
use rand::seq::SliceRandom;

pub struct ImageScheduler {
    primary_images: Vec<ImageState>,
    overlay_images: Vec<ImageState>,
    rng: ThreadRng,
    global_effects: Vec<ImageState>,
    transitions: Vec<Transition>,
    overlay_schedule: Vec<usize>,
}

impl ImageScheduler {
    pub fn new(
        images: &[Media],
        effect_factory: &mut effects::EffectFactory,
    ) -> Self {
        let mut rng = rand::thread_rng();
        let mut primary_images = Vec::new();
        let mut overlay_images = Vec::new();
        let mut global_effects = Vec::new();

        // Schedule primary images
        Self::schedule_images(
            images,
            effect_factory,
            &mut rng,
            CONFIG.total_frames() / 5,
            Self::generate_weighted_duration,
            (0.5, 1.5),
            |ef| ef.generate_random_effects(),
            &mut primary_images,
        );

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

        // Schedule overlay images
        Self::schedule_overlay_images(
            images,
            effect_factory,
            &mut rng,
            overlay_schedule.clone(),
            &mut overlay_images,
        );

        // Schedule larger, slower overlay images
        Self::schedule_images(
            images,
            effect_factory,
            &mut rng.clone(),
            CONFIG.total_frames() / 100,
            |rng, max| {
                let min_duration = CONFIG.framerate * 2;
                let max_duration = (CONFIG.framerate * 5).min(max as u32);
                if max_duration >= min_duration {
                    rng.gen_range(min_duration..=max_duration) as usize
                } else if max_duration > 0 {
                    rng.gen_range(1..=max_duration) as usize
                } else {
                    0
                }
            },
            (0.5, 1.5),
            |_| vec![Effect::LowOpacityFlash {
                flash_duration: rng.clone().gen_range(1..=3),
                flash_interval: rng.clone().gen_range(30..=100),
            }],
            &mut overlay_images,
        );

        // Schedule global effects
        Self::schedule_global_effects(
            effect_factory,
            &mut rng,
            &mut global_effects,
        );

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

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

        scheduler
    }

    fn schedule_images(
        images: &[Media],
        effect_factory: &mut effects::EffectFactory,
        rng: &mut ThreadRng,
        num_images: usize,
        duration_generator: impl Fn(&mut ThreadRng, usize) -> usize,
        scale_range: (f32, f32),
        effect_generator: impl Fn(&mut effects::EffectFactory) -> Vec<Effect>,
        scheduled_images: &mut Vec<ImageState>,
    ) {
        for _ in 0..num_images {
            // Randomly select an image
            let media = images.choose(rng).unwrap();

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

            // Maximum duration in frames
            let max_possible_duration = CONFIG.total_frames() - start_frame;
            let duration = duration_generator(rng, max_possible_duration);
            if duration == 0 {
                continue;
            }

            // Create image state
            if let Some(image_state) = Self::create_image_state(
                media,
                effect_factory,
                rng,
                start_frame,
                duration,
                scale_range,
                &effect_generator,
            ) {
                scheduled_images.push(image_state);
            }
        }
    }

    fn schedule_overlay_images(
        images: &[Media],
        effect_factory: &mut effects::EffectFactory,
        rng: &mut ThreadRng,
        overlay_schedule: Vec<usize>,
        overlay_images: &mut Vec<ImageState>,
    ) {
        for &frame_number in &overlay_schedule {
            // Randomly select an image
            let media = images.choose(rng).unwrap();

            // Short durations for overlay images
            let max_possible_duration = CONFIG.total_frames() - frame_number;
            let max_duration = 10.min(max_possible_duration);
            let duration = Self::generate_short_duration(rng, max_duration);
            if duration == 0 {
                continue;
            }

            // Create image state
            if let Some(image_state) = Self::create_image_state(
                media,
                effect_factory,
                rng,
                frame_number,
                duration,
                (0.05, 0.2),
                &|ef| ef.generate_overlay_effects(),
            ) {
                overlay_images.push(image_state);
            }
        }
    }

    fn schedule_global_effects(
        effect_factory: &mut effects::EffectFactory,
        rng: &mut ThreadRng,
        global_effects: &mut Vec<ImageState>,
    ) {
        for _ in 0..(CONFIG.total_frames() / 50) {
            // 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(rng, max_duration);
            if duration == 0 {
                continue;
            }
            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 {
                frames: ImageFrames::Static(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);
        }
    }

    fn create_image_state(
        media: &Media,
        effect_factory: &mut effects::EffectFactory,
        rng: &mut ThreadRng,
        start_frame: usize,
        duration: usize,
        scale_range: (f32, f32),
        effect_generator: &impl Fn(&mut effects::EffectFactory) -> Vec<Effect>,
    ) -> Option<ImageState> {
        // Process media into frames and durations
        let (frames, any_frame, video_start_time) =
            Self::process_media_to_frames(media, duration, rng)?;

        // Calculate scale
        let scale = Self::calculate_scale(&any_frame, rng, scale_range);

        // Calculate position
        let position = Self::calculate_position(&any_frame, rng);

        // Generate effects
        let effects = effect_generator(effect_factory);

        // Create image state
        Some(ImageState {
            frames,
            position,
            scale,
            start_frame,
            end_frame: start_frame + duration,
            effects,
        })
    }

    fn process_media_to_frames(
        media: &Media,
        duration: usize,
        rng: &mut ThreadRng,
    ) -> Option<(ImageFrames, RgbaImage, Option<u64>)> {
        let mut video_start_time = None;

        let frames: ImageFrames = match media {
            Media::Static(img) => {
                ImageFrames::Static(img.clone())
            },
            Media::Animated { path, total_duration_ms, frame_count } => {
                ImageFrames::Animated {
                    path: path.clone(),
                    total_duration_ms: *total_duration_ms,
                    frame_count: *frame_count,
                }
            },
            Media::Video { path, total_duration_ms, frame_rate, width, height } => {
                let max_start_time_ms = *total_duration_ms as i64
                    - (duration as i64 * 1000 / CONFIG.framerate as i64);
                let video_start_time_ms = if max_start_time_ms > 0 {
                    rng.gen_range(0..=max_start_time_ms as u64)
                } else {
                    0
                };
                video_start_time = Some(video_start_time_ms);
                ImageFrames::Video {
                    clip: Rc::new(RefCell::new(VideoClip::open(&path).ok()?)),
                }
            }
        };

        let any_frame = frames.get_first_frame().unwrap_or_else(|| {
            // In case of error, use a placeholder image
            RgbaImage::new(1, 1)
        });

        Some((frames, any_frame, video_start_time))
    }

    fn calculate_scale(
        any_frame: &RgbaImage,
        rng: &mut ThreadRng,
        scale_range: (f32, f32),
    ) -> f32 {
        // You can customize this function further if needed
        let max_scale_x = CONFIG.width as f32 / any_frame.width() as f32;
        let max_scale_y = CONFIG.height as f32 / any_frame.height() as f32;
        let max_scale = max_scale_x.min(max_scale_y).min(scale_range.1); // Allow up to 1.5x scaling
        let min_scale = scale_range.0.min(max_scale);
        if max_scale > min_scale {
            rng.gen_range(min_scale..=max_scale)
        } else {
            min_scale
        }
    }
    

    fn calculate_position(
        any_frame: &RgbaImage,
        rng: &mut ThreadRng,
    ) -> (i64, i64) {
        let max_x = CONFIG.width.saturating_sub(any_frame.width() as u32) as i64;
        let max_y = CONFIG.height.saturating_sub(any_frame.height() as u32) 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
        };
        (x, y)
    }

    // 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 = CONFIG.framerate; // 1 second
                let max_transition_duration = 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 = CONFIG.width;
                let frame_height = CONFIG.height;

                // Resize current image
                // let from_image_resized: Vec<_> = current_state.frames.iter().map(|img| {
                //     image::imageops::resize(
                //         img,
                //         frame_width as u32,
                //         frame_height as u32,
                //         image::imageops::FilterType::Lanczos3,
                //     )
                // }).collect();

                // // Resize next image
                // let to_image_resized: Vec<_> = next_state.frames.iter().map(|img| {
                //     image::imageops::resize(
                //         img,
                //         frame_width as u32,
                //         frame_height as u32,
                //         image::imageops::FilterType::Lanczos3,
                //     )
                // }).collect();

                // Create the transition
                let transition = Transition {
                    from_frames: current_state.frames.clone(),//: from_image_resized,
                    to_frames: next_state.frames.clone(),//: 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, CONFIG}, types::ImageFrames};

#[derive(Clone)]
pub struct Transition {
    pub from_frames: ImageFrames,
    pub to_frames: ImageFrames,
    pub start_frame: usize,
    pub end_frame: usize,
    pub duration: usize,
    pub effect: TransitionEffect,
    pub particles: Vec<Particle>,
    pub initialized: bool, 
}

#[derive(Clone)]
pub enum TransitionEffect {
    Reassembly,
    // You can add other transition effects in the future
}

// Define the Particle struct
#[derive(Clone)]
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,
    // 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 first_frame = if let Some(first_frame) = self.from_frames.get_first_frame()
        {
            first_frame
        }
        else
        {
            return
        };
        let width = first_frame.width();
        let height = first_frame.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;
                }

                // 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,
                    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,
) -> RgbaImage {
    let progress = (frame_number - transition.start_frame) as f32 / transition.duration as f32;

    // Get the current frames
    let from_image = transition.from_frames.get_current_frame(frame_number, transition.start_frame, transition.end_frame)
        .unwrap_or_else(|| RgbaImage::new(CONFIG.width, CONFIG.height));
    let to_image = transition.to_frames.get_current_frame(frame_number, transition.start_frame, transition.end_frame)
        .unwrap_or_else(|| RgbaImage::new(CONFIG.width, CONFIG.height));

    let width = CONFIG.width;
    let height = CONFIG.height;

    // Initialize particles
    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));

        // Extract blocks from current frames
        let src_block = extract_block(&from_image, particle.src_x, particle.src_y, particle.block_size);
        let dst_block = extract_block(&to_image, particle.dst_x, particle.dst_y, particle.block_size);

        // Blend and transform blocks
        let blended_block = blend_blocks(
            &src_block,
            &dst_block,
            progress,
            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 extract_block(image: &RgbaImage, x: u32, y: u32, block_size: u32) -> RgbaImage {
    let width = image.width();
    let height = image.height();
    let block_width = block_size.min(width - x);
    let block_height = block_size.min(height - y);

    image::imageops::crop_imm(
        image,
        x,
        y,
        block_width,
        block_height,
    ).to_image()
}

fn blend_blocks(
    src_block: &RgbaImage,
    dst_block: &RgbaImage,
    progress: f32,
    rotation: f32,
    scale: f32,
    alpha: f32,
) -> RgbaImage {
    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 = get_pixel_safe(src_block, x, y);
        let dst_pixel = get_pixel_safe(dst_block, 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::{AnimationDecoder, RgbaImage};
use std::{cell::RefCell, fs::File, path::PathBuf, rc::Rc};

use crate::image_loader::VideoClip;

#[derive(Clone)]
pub struct ImageState {
    pub frames: ImageFrames,
    pub position: (i64, i64),
    pub scale: f32,
    pub start_frame: usize,
    pub end_frame: usize,
    pub effects: Vec<crate::effects::Effect>,
}

#[derive(Clone)]
pub enum ImageFrames {
    Static(RgbaImage),
    Animated {
        path: PathBuf,
        total_duration_ms: u64, // total duration in milliseconds
        frame_count: usize,
    },
    Video {
        clip: Rc<RefCell<VideoClip>>,
    },
}

impl ImageFrames
{
    pub fn get_first_frame(&self) -> Option<RgbaImage> {
        match self
        {
            ImageFrames::Static(img) => Some(img.clone()),
            ImageFrames::Animated { path, total_duration_ms: _, frame_count:_ } => 
            {
                // Load the first frame of the animated GIF
                if let Ok(file) = File::open(path) {
                    if let Ok(decoder) = image::codecs::gif::GifDecoder::new(file) {
                        let frames = decoder.into_frames();
                        let mut frames_iter = frames.into_iter();
                        if let Some(frame_result) = frames_iter.next() {
                            if let Ok(frame) = frame_result {
                                return Some(frame.into_buffer());
                            }
                        }
                    }
                }
                None
            },
            ImageFrames::Video { clip } => {
                if let Ok(frame) = crate::image_loader::get_video_frame_at_timestamp(&clip.borrow().path, 0 as u64) {
                    return Some(frame);
                }
                None
            },
        }
    }

    pub fn get_current_frame(&self, frame_number: usize, start_frame: usize, end_frame: usize) -> Option<RgbaImage> 
    {
        match self {
            ImageFrames::Static(img) => Some(img.clone()),
            ImageFrames::Animated { path, frame_count, .. } => {
                let relative_frame = frame_number - start_frame;
                let total_duration_frames = end_frame - start_frame;
                let progress = relative_frame as f64 / total_duration_frames as f64;
                let gif_frame_index = (progress * *frame_count as f64).floor() as usize % *frame_count;

                if let Ok(file) = File::open(path) {
                    if let Ok(decoder) = image::codecs::gif::GifDecoder::new(file) {
                        let frames = decoder.into_frames();
                        let mut frames_iter = frames.into_iter();
                        for current_frame_index in 0..=gif_frame_index {
                            if let Some(frame_result) = frames_iter.next() {
                                if current_frame_index == gif_frame_index {
                                    if let Ok(frame) = frame_result {
                                        return Some(frame.into_buffer());
                                    }
                                }
                            } else {
                                break;
                            }
                        }
                    }
                }
                None
            },
            ImageFrames::Video { clip } => {
                let mut clip = clip.borrow_mut();
                // Calculate the desired timestamp based on frame_number
                let relative_frame = frame_number - start_frame;
                let time_base: f64 = clip.time_base.into();
                let desired_pts: i64 = (relative_frame as f64 / clip.frame_rate as f64 / time_base).round() as i64;
                while clip.current_pts < desired_pts {
                    if clip.get_next_frame().ok().is_none() {
                        break;
                    }
                    clip.current_pts += 1;
                }
                clip.get_next_frame().ok().flatten()
            },
        }
    }
}

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

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

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

impl VideoEncoder {
    pub fn new() -> 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 as usize; // Adjust quality (0-255)
        encoder_cfg.speed_settings = SpeedSettings::from_preset(CONFIG.speed as u8);
        encoder_cfg.bit_depth = CONFIG.bit_depth as usize;
        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(())
}
////// end of source code //////

Video generation goes at a decent pace for the first 50-100 frames, but then slows down dramatically, and takes far too long. Diagnose why this happens and lay out what the solution is.