Update rust_highlight/src/lib.rs

rust_highlight/src/lib.rs

@@ -1,303 +1,83 @@
 use pyo3::prelude::*;
-use pyo3::exceptions::PyRuntimeError;
-use image::{ImageBuffer, Rgb};
-use imageproc::drawing::{draw_text_mut, draw_line_segment_mut, draw_filled_circle_mut};
-use ab_glyph::{FontVec, PxScale, Font as AbFont};
-use std::process::{Command, Stdio};
-use std::io::Write;
+use std::process::Command;
 use std::path::Path;
 
-const WIDTH: u32 = 1280;
-const HEIGHT: u32 = 720;
-const MARGIN_X: u32 = 40;
-const MARGIN_Y: u32 = 60;
-const LINE_SPACING: u32 = 8;
-const FPS: u32 = 30;
-const ANIMATION_FRAMES_PER_CHAR: usize = 2;
-const BG_COLOR: Rgb<u8> = Rgb([255, 255, 255]);
-const DEFAULT_TEXT_COLOR: Rgb<u8> = Rgb([0, 0, 0]);
-const HEADER_TEXT_COLOR: Rgb<u8> = Rgb([0, 0, 255]);
-const PEN_COLOR: Rgb<u8> = Rgb([255, 0, 0]);
-const FFMPEG_PRESET: &str = "ultrafast";
-const CRF: u32 = 28;
-
 #[pyfunction]
 fn render_video(
     id: usize,
-    text: String,
+    image_path: String,
     audio_path: String,
-    audio_duration: f64,
-    clips_dir: String,
+    duration: f64,
+    words: Vec<(String, (u32, u32, u32, u32))>,
+    output_dir: String,
 ) -> PyResult<String> {
+    // Validate paths
+    if !Path::new(&image_path).exists() {
+        return Err(pyo3::exceptions::PyFileNotFoundError::new_err(format!(
+            "Image not found: {}",
+            image_path
+        )));
+    }
     if !Path::new(&audio_path).exists() {
-        return Err(PyRuntimeError::new_err(format!("Audio not found: {}", audio_path)));
+        return Err(pyo3::exceptions::PyFileNotFoundError::new_err(format!(
+            "Audio not found: {}",
+            audio_path
+        )));
     }
 
-    let font_data: &[u8] = include_bytes!("../fonts/LiberationSans-Regular.ttf");
-    let font = FontVec::try_from_vec(font_data.to_vec())
-        .map_err(|_| PyRuntimeError::new_err("Failed to load font"))?;
-
-    let (wrapped_lines, line_styles) = wrap_text(&text, &font, WIDTH - 2 * MARGIN_X);
-    let (_line_heights, y_positions) = calculate_line_positions(&wrapped_lines, &line_styles, &font);
-
-    let full_text = wrapped_lines.join("\n");
-    let visible_indices: Vec<usize> = full_text
-        .char_indices()
-        .filter(|(_, ch)| *ch != ' ' && *ch != '\n' && *ch != '\t')
-        .map(|(i, _)| i)
-        .collect();
-
-    let total_frames = visible_indices.len() * ANIMATION_FRAMES_PER_CHAR;
-    println!("Rendering {} frames for {} glyphs", total_frames, visible_indices.len());
-
-    let silent_video_path = format!("{}/silent_video{}.mp4", clips_dir, id);
-    let mut ffmpeg = Command::new("ffmpeg")
-        .args(&[
-            "-y",
-            "-f", "rawvideo",
-            "-pix_fmt", "rgb24",
-            "-s", &format!("{}x{}", WIDTH, HEIGHT),
-            "-r", &FPS.to_string(),
-            "-i", "-",
-            "-an",
-            "-c:v", "libx264",
-            "-preset", FFMPEG_PRESET,
-            "-crf", &CRF.to_string(),
-            "-pix_fmt", "yuv420p",
-            &silent_video_path,
-        ])
-        .stdin(Stdio::piped())
-        .spawn()
-        .map_err(|e| PyRuntimeError::new_err(format!("Failed to spawn FFmpeg: {}", e)))?;
-
-    let mut stdin = ffmpeg.stdin.take().unwrap();
-
-    let mut prev_text_len = 0;
-    for &char_idx in visible_indices.iter() {
-        let visible_text = &full_text[..=char_idx];
-        
-        if visible_text.len() != prev_text_len {
-            let lines: Vec<&str> = visible_text.split('\n').collect();
-            let last_line = lines.last().unwrap_or(&"");
-            let line_idx = lines.len() - 1;
-            
-            let (pen_x, pen_y) = calculate_pen_position(
-                last_line,
-                &font,
-                line_idx,
-                &line_styles,
-                &y_positions,
-            );
-
-            for anim_step in 0..ANIMATION_FRAMES_PER_CHAR {
-                let frame = render_frame(
-                    visible_text,
-                    &line_styles,
-                    &y_positions,
-                    &font,
-                    pen_x,
-                    pen_y,
-                    anim_step as f32 / ANIMATION_FRAMES_PER_CHAR as f32,
-                );
-                
-                stdin.write_all(frame.as_raw())
-                    .map_err(|e| PyRuntimeError::new_err(format!("Failed to write frame: {}", e)))?;
-            }
-            
-            prev_text_len = visible_text.len();
-        }
+    let output_path = format!("{}/clip{}.mp4", output_dir, id);
+    let duration_str = duration.to_string();
+    let n_words = words.len() as f64;
+    let highlight_duration = duration / n_words;
+
+    // Build filter graph
+    let mut filters = vec![];
+    for (i, (_, (x, y, w, h))) in words.iter().enumerate() {
+        let start = i as f64 * highlight_duration;
+        let end = start + highlight_duration;
+        filters.push(format!(
+            "drawbox=x={x}:y={y}:w={w}:h={h}:color=yellow@0.5:t=fill:enable='between(t,{start:.2},{end:.2})'"
+        ));
     }
 
-    drop(stdin);
-    let output = ffmpeg.wait()
-        .map_err(|e| PyRuntimeError::new_err(format!("FFmpeg error: {}", e)))?;
+    let filter_chain = filters.join(",");
 
-    if !output.success() {
-        return Err(PyRuntimeError::new_err("FFmpeg encoding failed"));
-    }
-
-    let rendered_duration = total_frames as f64 / FPS as f64;
-    let speed_factor = rendered_duration / audio_duration;
-    
-    let final_video_path = format!("{}/clip{}.mp4", clips_dir, id);
-    
+    // FFmpeg command
     let status = Command::new("ffmpeg")
         .args(&[
             "-y",
-            "-i", &silent_video_path,
+            "-loop", "1",
+            "-i", &image_path,
             "-i", &audio_path,
-            "-filter:v", &format!("setpts={}*PTS", 1.0 / speed_factor),
+            "-vf", &filter_chain,
             "-c:v", "libx264",
-            "-preset", "ultrafast",
+            "-preset", "fast",
+            "-crf", "18",
+            "-pix_fmt", "yuv420p",
+            "-movflags", "+faststart",
             "-c:a", "aac",
-            "-shortest",
-            &final_video_path,
+            "-b:a", "192k",
+            "-r", "60",
+            "-fps_mode", "vfr",  // Replaces deprecated -vsync
+            "-t", &duration_str,
+            &output_path,
         ])
-        .status()
-        .map_err(|e| PyRuntimeError::new_err(format!("Failed to mux audio: {}", e)))?;
-
-    if !status.success() {
-        return Err(PyRuntimeError::new_err("Audio muxing failed"));
+        .status();
+
+    match status {
+        Ok(exit_status) if exit_status.success() => Ok(output_path),
+        Ok(_) => Err(pyo3::exceptions::PyRuntimeError::new_err(
+            "FFmpeg command failed. Check filter syntax.",
+        )),
+        Err(e) => Err(pyo3::exceptions::PyRuntimeError::new_err(format!(
+            "Failed to execute FFmpeg: {}",
+            e
+        ))),
     }
-
-    let _ = std::fs::remove_file(&silent_video_path);
-    println!("Final video saved: {}", final_video_path);
-    Ok(final_video_path)
-}
-
-fn wrap_text(text: &str, font: &FontVec, max_width: u32) -> (Vec<String>, Vec<bool>) {
-    let mut wrapped = Vec::new();
-    let mut styles = Vec::new();
-    
-    for line in text.lines() {
-        let is_header = line.trim_start().starts_with("###");
-        let clean_line = if is_header {
-            line.trim_start_matches("###").trim()
-        } else {
-            line
-        };
-
-        let scale = if is_header { PxScale::from(60.0) } else { PxScale::from(45.0) };
-        
-        if clean_line.is_empty() {
-            wrapped.push(String::new());
-            styles.push(false);
-            continue;
-        }
-
-        let words: Vec<&str> = clean_line.split_whitespace().collect();
-        let mut current = String::new();
-        
-        for word in words {
-            let candidate = if current.is_empty() {
-                word.to_string()
-            } else {
-                format!("{} {}", current, word)
-            };
-            
-            let width = measure_text(&candidate, font, scale);
-            
-            if width <= max_width as f32 {
-                current = candidate;
-            } else {
-                if !current.is_empty() {
-                    wrapped.push(current.clone());
-                    styles.push(is_header);
-                }
-                current = word.to_string();
-            }
-        }
-        
-        if !current.is_empty() {
-            wrapped.push(current);
-            styles.push(is_header);
-        }
-    }
-    
-    (wrapped, styles)
-}
-
-fn measure_text(text: &str, font: &FontVec, scale: PxScale) -> f32 {
-    use ab_glyph::ScaleFont;
-    let scaled = font.as_scaled(scale);
-    
-    text.chars()
-        .filter_map(|c| {
-            let glyph_id = font.glyph_id(c);
-            Some(scaled.h_advance(glyph_id))
-        })
-        .sum()
-}
-
-fn calculate_line_positions(lines: &[String], styles: &[bool], font: &FontVec) -> (Vec<u32>, Vec<u32>) {
-    use ab_glyph::ScaleFont;
-    let mut heights = Vec::new();
-    let mut positions = Vec::new();
-    let mut y = MARGIN_Y;
-    
-    for (i, _line) in lines.iter().enumerate() {
-        let scale = if styles[i] { PxScale::from(60.0) } else { PxScale::from(45.0) };
-        let scaled = font.as_scaled(scale);
-        let height = (scaled.ascent() - scaled.descent()).ceil() as u32 + LINE_SPACING;
-        
-        heights.push(height);
-        positions.push(y);
-        y += height;
-    }
-    
-    (heights, positions)
-}
-
-fn calculate_pen_position(
-    last_line: &str,
-    font: &FontVec,
-    line_idx: usize,
-    styles: &[bool],
-    y_positions: &[u32],
-) -> (f32, f32) {
-    let scale = if styles.get(line_idx).copied().unwrap_or(false) {
-        PxScale::from(60.0)
-    } else {
-        PxScale::from(45.0)
-    };
-    
-    let width = measure_text(last_line, font, scale);
-    let pen_x = MARGIN_X as f32 + width + 5.0;
-    let pen_y = y_positions.get(line_idx).copied().unwrap_or(MARGIN_Y) as f32 + 20.0;
-    
-    (pen_x, pen_y)
-}
-
-fn render_frame(
-    visible_text: &str,
-    line_styles: &[bool],
-    y_positions: &[u32],
-    font: &FontVec,
-    pen_x: f32,
-    pen_y: f32,
-    anim_offset: f32,
-) -> ImageBuffer<Rgb<u8>, Vec<u8>> {
-    let mut img = ImageBuffer::from_pixel(WIDTH, HEIGHT, BG_COLOR);
-    
-    let visible_lines: Vec<&str> = visible_text.split('\n').collect();
-    for (idx, line) in visible_lines.iter().enumerate() {
-        if idx >= line_styles.len() {
-            break;
-        }
-        
-        let is_header = line_styles[idx];
-        let scale = if is_header { PxScale::from(60.0) } else { PxScale::from(45.0) };
-        let color = if is_header { HEADER_TEXT_COLOR } else { DEFAULT_TEXT_COLOR };
-        let y = y_positions.get(idx).copied().unwrap_or(MARGIN_Y);
-        
-        if !line.is_empty() {
-            draw_text_mut(&mut img, color, MARGIN_X as i32, y as i32, scale, font, line);
-        }
-    }
-    
-    if pen_x > 0.0 {
-        let pen_offset = (anim_offset * std::f32::consts::PI).sin() * 10.0;
-        let pen_tip_y = pen_y + pen_offset;
-        
-        let angle = 45.0_f32.to_radians();
-        let pen_end_x = pen_x + 20.0 * angle.cos();
-        let pen_end_y = pen_tip_y - 20.0 * angle.sin();
-        
-        draw_line_segment_mut(
-            &mut img,
-            (pen_x, pen_tip_y),
-            (pen_end_x, pen_end_y),
-            PEN_COLOR,
-        );
-        
-        draw_filled_circle_mut(&mut img, (pen_x as i32, pen_tip_y as i32), 5, PEN_COLOR);
-    }
-    
-    img
 }
 
 #[pymodule]
-fn rust_video_gen(m: &Bound<PyModule>) -> PyResult<()> {
+fn rust_highlight(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(render_video, m)?)?;
     Ok(())
-}
+}