1use crate::driver::Screenshot;
14use crate::result::{ProbarError, ProbarResult};
15use gif::{Encoder, Frame, Repeat};
16use image::{DynamicImage, GenericImageView, ImageFormat};
17use serde::{Deserialize, Serialize};
18use std::path::Path;
19
20#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct GifConfig {
23 pub fps: u8,
25 pub width: u32,
27 pub height: u32,
29 pub quality: u8,
31 pub loop_count: u16,
33}
34
35impl Default for GifConfig {
36 fn default() -> Self {
37 Self {
38 fps: 10,
39 width: 800,
40 height: 600,
41 quality: 80,
42 loop_count: 0, }
44 }
45}
46
47impl GifConfig {
48 #[must_use]
50 pub fn new(width: u32, height: u32) -> Self {
51 Self {
52 width,
53 height,
54 ..Default::default()
55 }
56 }
57
58 #[must_use]
60 pub fn with_fps(mut self, fps: u8) -> Self {
61 self.fps = fps.clamp(1, 60);
62 self
63 }
64
65 #[must_use]
67 pub fn with_quality(mut self, quality: u8) -> Self {
68 self.quality = quality.clamp(1, 100);
69 self
70 }
71
72 #[must_use]
74 pub fn with_loop_count(mut self, count: u16) -> Self {
75 self.loop_count = count;
76 self
77 }
78
79 #[must_use]
81 pub fn frame_delay_cs(&self) -> u16 {
82 (100 / u16::from(self.fps.max(1))).max(1)
86 }
87}
88
89#[derive(Debug, Clone)]
91pub struct GifFrame {
92 pub data: Vec<u8>,
94 pub width: u32,
96 pub height: u32,
98 pub timestamp_ms: u64,
100}
101
102impl GifFrame {
103 #[must_use]
105 pub fn new(data: Vec<u8>, width: u32, height: u32, timestamp_ms: u64) -> Self {
106 Self {
107 data,
108 width,
109 height,
110 timestamp_ms,
111 }
112 }
113
114 pub fn from_screenshot(screenshot: &Screenshot, timestamp_ms: u64) -> ProbarResult<Self> {
116 let img = image::load_from_memory_with_format(&screenshot.data, ImageFormat::Png).map_err(
118 |e| ProbarError::ImageProcessing {
119 message: format!("Failed to decode screenshot: {e}"),
120 },
121 )?;
122
123 let (width, height) = img.dimensions();
124 let rgba = img.to_rgba8();
125
126 Ok(Self {
127 data: rgba.into_raw(),
128 width,
129 height,
130 timestamp_ms,
131 })
132 }
133}
134
135#[derive(Debug)]
152pub struct GifRecorder {
153 config: GifConfig,
154 frames: Vec<GifFrame>,
155 recording: bool,
156 start_time_ms: u64,
157 encoded_data: Option<Vec<u8>>,
158}
159
160impl GifRecorder {
161 #[must_use]
163 pub fn new(config: GifConfig) -> Self {
164 Self {
165 config,
166 frames: Vec::new(),
167 recording: false,
168 start_time_ms: 0,
169 encoded_data: None,
170 }
171 }
172
173 #[must_use]
175 pub fn config(&self) -> &GifConfig {
176 &self.config
177 }
178
179 #[must_use]
181 pub fn is_recording(&self) -> bool {
182 self.recording
183 }
184
185 #[must_use]
187 pub fn frame_count(&self) -> usize {
188 self.frames.len()
189 }
190
191 pub fn start(&mut self) -> ProbarResult<()> {
197 if self.recording {
198 return Err(ProbarError::InvalidState {
199 message: "GIF recording already in progress".to_string(),
200 });
201 }
202
203 self.frames.clear();
204 self.encoded_data = None;
205 self.recording = true;
206 self.start_time_ms = std::time::SystemTime::now()
207 .duration_since(std::time::UNIX_EPOCH)
208 .map(|d| d.as_millis() as u64)
209 .unwrap_or(0);
210
211 Ok(())
212 }
213
214 pub fn capture_frame(&mut self, screenshot: &Screenshot) -> ProbarResult<()> {
220 if !self.recording {
221 return Err(ProbarError::InvalidState {
222 message: "GIF recording not started".to_string(),
223 });
224 }
225
226 let current_time = std::time::SystemTime::now()
227 .duration_since(std::time::UNIX_EPOCH)
228 .map(|d| d.as_millis() as u64)
229 .unwrap_or(0);
230
231 let timestamp_ms = current_time.saturating_sub(self.start_time_ms);
232
233 let frame = GifFrame::from_screenshot(screenshot, timestamp_ms)?;
234 self.frames.push(frame);
235
236 Ok(())
237 }
238
239 pub fn add_frame(&mut self, frame: GifFrame) -> ProbarResult<()> {
245 if !self.recording {
246 return Err(ProbarError::InvalidState {
247 message: "GIF recording not started".to_string(),
248 });
249 }
250
251 self.frames.push(frame);
252 Ok(())
253 }
254
255 pub fn stop(&mut self) -> ProbarResult<Vec<u8>> {
261 if !self.recording {
262 return Err(ProbarError::InvalidState {
263 message: "GIF recording not started".to_string(),
264 });
265 }
266
267 self.recording = false;
268
269 if self.frames.is_empty() {
270 return Err(ProbarError::InvalidState {
271 message: "No frames captured".to_string(),
272 });
273 }
274
275 let encoded = self.encode_gif()?;
276 self.encoded_data = Some(encoded.clone());
277
278 Ok(encoded)
279 }
280
281 #[must_use]
283 pub fn encoded_data(&self) -> Option<&[u8]> {
284 self.encoded_data.as_deref()
285 }
286
287 pub fn save(&self, path: &Path) -> ProbarResult<()> {
293 let data = self
294 .encoded_data
295 .as_ref()
296 .ok_or_else(|| ProbarError::InvalidState {
297 message: "No encoded GIF data. Call stop() first.".to_string(),
298 })?;
299
300 std::fs::write(path, data)?;
301
302 Ok(())
303 }
304
305 fn encode_gif(&self) -> ProbarResult<Vec<u8>> {
307 let mut output = Vec::new();
308
309 let width = self.config.width as u16;
311 let height = self.config.height as u16;
312
313 {
314 let mut encoder = Encoder::new(&mut output, width, height, &[]).map_err(|e| {
315 ProbarError::ImageProcessing {
316 message: format!("Failed to create GIF encoder: {e}"),
317 }
318 })?;
319
320 let repeat = if self.config.loop_count == 0 {
322 Repeat::Infinite
323 } else {
324 Repeat::Finite(self.config.loop_count)
325 };
326 encoder
327 .set_repeat(repeat)
328 .map_err(|e| ProbarError::ImageProcessing {
329 message: format!("Failed to set GIF repeat: {e}"),
330 })?;
331
332 let frame_delay = self.config.frame_delay_cs();
333
334 for gif_frame in &self.frames {
335 let rgba_data = self.resize_frame(gif_frame)?;
337
338 let mut frame = Frame::from_rgba_speed(
340 width,
341 height,
342 &mut rgba_data.clone(),
343 self.quality_to_speed(),
344 );
345 frame.delay = frame_delay;
346
347 encoder
348 .write_frame(&frame)
349 .map_err(|e| ProbarError::ImageProcessing {
350 message: format!("Failed to write GIF frame: {e}"),
351 })?;
352 }
353 }
354
355 Ok(output)
356 }
357
358 fn resize_frame(&self, frame: &GifFrame) -> ProbarResult<Vec<u8>> {
360 if frame.width == self.config.width && frame.height == self.config.height {
361 return Ok(frame.data.clone());
362 }
363
364 let img = DynamicImage::ImageRgba8(
366 image::RgbaImage::from_raw(frame.width, frame.height, frame.data.clone()).ok_or_else(
367 || ProbarError::ImageProcessing {
368 message: "Invalid frame data dimensions".to_string(),
369 },
370 )?,
371 );
372
373 let resized = img.resize_exact(
375 self.config.width,
376 self.config.height,
377 image::imageops::FilterType::Triangle,
378 );
379
380 Ok(resized.to_rgba8().into_raw())
381 }
382
383 fn quality_to_speed(&self) -> i32 {
385 let normalized = i32::from(100 - self.config.quality);
389 (normalized * 29 / 100 + 1).clamp(1, 30)
390 }
391}
392
393#[cfg(test)]
398#[allow(clippy::unwrap_used, clippy::expect_used)]
399mod tests {
400 use super::*;
401 use image::{ImageFormat, Rgba};
402 use std::io::Cursor;
403
404 mod gif_config_tests {
405 use super::*;
406
407 #[test]
408 fn test_default_config() {
409 let config = GifConfig::default();
410 assert_eq!(config.fps, 10);
411 assert_eq!(config.width, 800);
412 assert_eq!(config.height, 600);
413 assert_eq!(config.quality, 80);
414 assert_eq!(config.loop_count, 0);
415 }
416
417 #[test]
418 fn test_new_config() {
419 let config = GifConfig::new(1920, 1080);
420 assert_eq!(config.width, 1920);
421 assert_eq!(config.height, 1080);
422 }
423
424 #[test]
425 fn test_with_fps() {
426 let config = GifConfig::default().with_fps(30);
427 assert_eq!(config.fps, 30);
428 }
429
430 #[test]
431 fn test_fps_clamping() {
432 let config = GifConfig::default().with_fps(100);
433 assert_eq!(config.fps, 60); let config = GifConfig::default().with_fps(0);
436 assert_eq!(config.fps, 1); }
438
439 #[test]
440 fn test_with_quality() {
441 let config = GifConfig::default().with_quality(50);
442 assert_eq!(config.quality, 50);
443 }
444
445 #[test]
446 fn test_quality_clamping() {
447 let config = GifConfig::default().with_quality(150);
448 assert_eq!(config.quality, 100);
449
450 let config = GifConfig::default().with_quality(0);
451 assert_eq!(config.quality, 1);
452 }
453
454 #[test]
455 fn test_with_loop_count() {
456 let config = GifConfig::default().with_loop_count(3);
457 assert_eq!(config.loop_count, 3);
458 }
459
460 #[test]
461 fn test_frame_delay_calculation() {
462 let config = GifConfig::default().with_fps(10);
463 assert_eq!(config.frame_delay_cs(), 10); let config = GifConfig::default().with_fps(20);
466 assert_eq!(config.frame_delay_cs(), 5); let config = GifConfig::default().with_fps(1);
469 assert_eq!(config.frame_delay_cs(), 100); }
471 }
472
473 mod gif_frame_tests {
474 use super::*;
475
476 #[test]
477 fn test_new_frame() {
478 let data = vec![255, 0, 0, 255]; let frame = GifFrame::new(data.clone(), 1, 1, 100);
480
481 assert_eq!(frame.data, data);
482 assert_eq!(frame.width, 1);
483 assert_eq!(frame.height, 1);
484 assert_eq!(frame.timestamp_ms, 100);
485 }
486
487 #[test]
488 fn test_frame_from_screenshot() {
489 let mut img = image::RgbaImage::new(2, 2);
491 img.put_pixel(0, 0, Rgba([255, 0, 0, 255])); img.put_pixel(1, 0, Rgba([0, 255, 0, 255])); img.put_pixel(0, 1, Rgba([0, 0, 255, 255])); img.put_pixel(1, 1, Rgba([255, 255, 255, 255])); let mut png_data = Vec::new();
497 img.write_to(&mut Cursor::new(&mut png_data), ImageFormat::Png)
498 .unwrap();
499
500 let screenshot = Screenshot::new(png_data, 2, 2);
501 let frame = GifFrame::from_screenshot(&screenshot, 500).unwrap();
502
503 assert_eq!(frame.width, 2);
504 assert_eq!(frame.height, 2);
505 assert_eq!(frame.timestamp_ms, 500);
506 assert_eq!(frame.data.len(), 16); }
508 }
509
510 mod gif_recorder_tests {
511 use super::*;
512
513 fn create_test_screenshot(width: u32, height: u32, color: [u8; 4]) -> Screenshot {
514 let mut img = image::RgbaImage::new(width, height);
515 for pixel in img.pixels_mut() {
516 *pixel = Rgba(color);
517 }
518
519 let mut png_data = Vec::new();
520 img.write_to(&mut Cursor::new(&mut png_data), ImageFormat::Png)
521 .unwrap();
522
523 Screenshot::new(png_data, width, height)
524 }
525
526 #[test]
527 fn test_new_recorder() {
528 let config = GifConfig::new(800, 600);
529 let recorder = GifRecorder::new(config);
530
531 assert_eq!(recorder.config().width, 800);
532 assert_eq!(recorder.config().height, 600);
533 assert!(!recorder.is_recording());
534 assert_eq!(recorder.frame_count(), 0);
535 }
536
537 #[test]
538 fn test_start_recording() {
539 let mut recorder = GifRecorder::new(GifConfig::default());
540
541 assert!(recorder.start().is_ok());
542 assert!(recorder.is_recording());
543 }
544
545 #[test]
546 fn test_start_recording_twice_fails() {
547 let mut recorder = GifRecorder::new(GifConfig::default());
548
549 recorder.start().unwrap();
550 let result = recorder.start();
551
552 assert!(result.is_err());
553 }
554
555 #[test]
556 fn test_capture_frame() {
557 let mut recorder = GifRecorder::new(GifConfig::new(100, 100));
558 recorder.start().unwrap();
559
560 let screenshot = create_test_screenshot(100, 100, [255, 0, 0, 255]);
561 let result = recorder.capture_frame(&screenshot);
562
563 assert!(result.is_ok());
564 assert_eq!(recorder.frame_count(), 1);
565 }
566
567 #[test]
568 fn test_capture_frame_without_start_fails() {
569 let mut recorder = GifRecorder::new(GifConfig::default());
570 let screenshot = create_test_screenshot(100, 100, [255, 0, 0, 255]);
571
572 let result = recorder.capture_frame(&screenshot);
573 assert!(result.is_err());
574 }
575
576 #[test]
577 fn test_add_frame() {
578 let mut recorder = GifRecorder::new(GifConfig::new(100, 100));
579 recorder.start().unwrap();
580
581 let frame = GifFrame::new(vec![255, 0, 0, 255], 1, 1, 0);
582 let result = recorder.add_frame(frame);
583
584 assert!(result.is_ok());
585 assert_eq!(recorder.frame_count(), 1);
586 }
587
588 #[test]
589 fn test_stop_recording() {
590 let mut recorder = GifRecorder::new(GifConfig::new(10, 10));
591 recorder.start().unwrap();
592
593 let screenshot = create_test_screenshot(10, 10, [255, 0, 0, 255]);
594 recorder.capture_frame(&screenshot).unwrap();
595
596 let result = recorder.stop();
597 assert!(result.is_ok());
598 assert!(!recorder.is_recording());
599 assert!(recorder.encoded_data().is_some());
600 }
601
602 #[test]
603 fn test_stop_without_frames_fails() {
604 let mut recorder = GifRecorder::new(GifConfig::default());
605 recorder.start().unwrap();
606
607 let result = recorder.stop();
608 assert!(result.is_err());
609 }
610
611 #[test]
612 fn test_stop_without_start_fails() {
613 let mut recorder = GifRecorder::new(GifConfig::default());
614
615 let result = recorder.stop();
616 assert!(result.is_err());
617 }
618
619 #[test]
620 fn test_save_gif() {
621 let mut recorder = GifRecorder::new(GifConfig::new(10, 10));
622 recorder.start().unwrap();
623
624 let screenshot = create_test_screenshot(10, 10, [255, 0, 0, 255]);
625 recorder.capture_frame(&screenshot).unwrap();
626 recorder.stop().unwrap();
627
628 let temp_dir = tempfile::tempdir().unwrap();
629 let path = temp_dir.path().join("test.gif");
630
631 let result = recorder.save(&path);
632 assert!(result.is_ok());
633 assert!(path.exists());
634
635 let data = std::fs::read(&path).unwrap();
637 assert_eq!(&data[0..6], b"GIF89a");
638 }
639
640 #[test]
641 fn test_save_without_encoding_fails() {
642 let recorder = GifRecorder::new(GifConfig::default());
643 let temp_dir = tempfile::tempdir().unwrap();
644 let path = temp_dir.path().join("test.gif");
645
646 let result = recorder.save(&path);
647 assert!(result.is_err());
648 }
649
650 #[test]
651 fn test_multiple_frames() {
652 let mut recorder = GifRecorder::new(GifConfig::new(10, 10).with_fps(10));
653 recorder.start().unwrap();
654
655 for color in [[255, 0, 0, 255], [0, 255, 0, 255], [0, 0, 255, 255]] {
657 let screenshot = create_test_screenshot(10, 10, color);
658 recorder.capture_frame(&screenshot).unwrap();
659 }
660
661 assert_eq!(recorder.frame_count(), 3);
662
663 let gif_data = recorder.stop().unwrap();
664 assert!(!gif_data.is_empty());
665 assert_eq!(&gif_data[0..6], b"GIF89a");
666 }
667
668 #[test]
669 fn test_frame_resizing() {
670 let mut recorder = GifRecorder::new(GifConfig::new(50, 50));
671 recorder.start().unwrap();
672
673 let screenshot = create_test_screenshot(100, 100, [255, 0, 0, 255]);
675 recorder.capture_frame(&screenshot).unwrap();
676
677 let gif_data = recorder.stop().unwrap();
678 assert!(!gif_data.is_empty());
679
680 let width = u16::from_le_bytes([gif_data[6], gif_data[7]]);
682 let height = u16::from_le_bytes([gif_data[8], gif_data[9]]);
683 assert_eq!(width, 50);
684 assert_eq!(height, 50);
685 }
686 }
687
688 mod property_tests {
689 use super::*;
690 use proptest::prelude::*;
691
692 proptest! {
693 #[test]
694 fn prop_config_dimensions_preserved(width in 1u32..4096, height in 1u32..4096) {
695 let config = GifConfig::new(width, height);
696 prop_assert_eq!(config.width, width);
697 prop_assert_eq!(config.height, height);
698 }
699
700 #[test]
701 fn prop_fps_always_valid(fps in 0u8..=255) {
702 let config = GifConfig::default().with_fps(fps);
703 prop_assert!(config.fps >= 1);
704 prop_assert!(config.fps <= 60);
705 }
706
707 #[test]
708 fn prop_quality_always_valid(quality in 0u8..=255) {
709 let config = GifConfig::default().with_quality(quality);
710 prop_assert!(config.quality >= 1);
711 prop_assert!(config.quality <= 100);
712 }
713
714 #[test]
715 fn prop_frame_delay_always_positive(fps in 1u8..=60) {
716 let config = GifConfig::default().with_fps(fps);
717 prop_assert!(config.frame_delay_cs() >= 1);
718 }
719 }
720 }
721}