Update rust_highlight/src/lib.rs

rust_highlight/src/lib.rs

@@ -1,83 +1,344 @@
+// rust_highlight/src/lib.rs
 use pyo3::prelude::*;
-use std::process::Command;
+use opencv::core::{Mat, Point, Scalar, CV_8UC3};
+use opencv::imgproc::{circle, get_text_size, line, put_text, HersheyFonts, LineTypes};
+use std::process::{Command, Stdio};
+use std::io::Write;
+use std::time::Instant;
 use std::path::Path;
+use std::f64::consts::PI;
 
 #[pyfunction]
-fn render_video(
-    id: usize,
-    image_path: String,
-    audio_path: 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
-        )));
-    }
+fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f64, clips_dir: String) -> PyResult<Option<String>> {
     if !Path::new(&audio_path).exists() {
-        return Err(pyo3::exceptions::PyFileNotFoundError::new_err(format!(
-            "Audio not found: {}",
-            audio_path
-        )));
+        return Err(pyo3::exceptions::PyFileNotFoundError::new_err(format!("Audio not found: {}", audio_path)));
+    }
+
+    let skip_spaces = false;
+    let fps: f64 = 30.0;
+    let animation_frames_per_char: usize = 2;
+    let width: i32 = 1280;
+    let height: i32 = 720;
+    let margin_x: i32 = 40;
+    let margin_y: i32 = 60;
+    let line_spacing: i32 = 8;
+    let font = HersheyFonts::FONT_HERSHEY_SIMPLEX as i32;
+    let default_font_scale: f64 = 1.5;
+    let header_font_scale: f64 = 2.0;
+    let default_thickness: i32 = 2;
+    let header_thickness: i32 = 3;
+    let default_text_color = Scalar::new(0.0, 0.0, 0.0, 0.0);  // BGR Black
+    let header_text_color = Scalar::new(255.0, 0.0, 0.0, 0.0);  // BGR Blue
+    let bg_color = Scalar::new(255.0, 255.0, 255.0, 0.0);  // BGR White
+    let ffmpeg_preset = "ultrafast";
+    let crf = "28";
+    // Pen settings
+    let pen_color = Scalar::new(0.0, 0.0, 255.0, 0.0);  // BGR Red
+    let pen_tip_radius: i32 = 5;
+    let pen_length: i32 = 20;
+    let pen_thickness: i32 = 2;
+    let pen_base_angle: i32 = 45;
+    let pen_movement_amplitude: i32 = 10;
+
+    let silent_video_name = format!("silent_video{}.mp4", id);
+    let silent_video_path = format!("{}/{}", clips_dir, silent_video_name);
+    let final_video_name = format!("clip{}.mp4", id);
+    let final_video_path = format!("{}/{}", clips_dir, final_video_name);
+
+    // Wrap text
+    let text_area_width = width - 2 * margin_x;
+    let (wrapped_lines, line_styles) = wrap_text_cv(&text, font, default_font_scale, default_thickness, text_area_width);
+
+    let full_text = wrapped_lines.join("\n");
+    if full_text.is_empty() {
+        println!("No text to animate.");
+        return Ok(None);
+    }
+
+    let visible_indices: Vec<usize> = if skip_spaces {
+        full_text.char_indices().filter(|&(_, ch)| ch != ' ' && ch != '\n' && ch != '\t').map(|(i, _)| i).collect()
+    } else {
+        (0..full_text.len()).collect()
+    };
+
+    let total_glyphs = visible_indices.len();
+    println!("Wrapped lines: {} lines, total glyphs (counted): {}", wrapped_lines.len(), total_glyphs);
+    if total_glyphs == 0 {
+        println!("No text to animate.");
+        return Ok(None);
     }
 
-    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})'"
-        ));
+    // Pre-calc line heights and y_positions
+    let mut line_heights: Vec<i32> = Vec::new();
+    let mut y_positions: Vec<i32> = Vec::new();
+    let mut y = margin_y;
+    for (i, line) in wrapped_lines.iter().enumerate() {
+        let is_header = line_styles[i];
+        let font_scale = if is_header { header_font_scale } else { default_font_scale };
+        let thickness = if is_header { header_thickness } else { default_thickness };
+        let line_for_size = if line.is_empty() { "Ay".to_string() } else { line.clone() };
+        let mut base_line = 0;
+        let size = get_text_size(&line_for_size, font, font_scale, thickness, &mut base_line).unwrap();
+        let h = size.height;
+        let lh = h + base_line + line_spacing;
+        line_heights.push(lh);
+        y_positions.push(y);
+        y += lh;
     }
 
-    let filter_chain = filters.join(",");
-
-    // FFmpeg command
-    let status = Command::new("ffmpeg")
-        .args(&[
-            "-y",
-            "-loop", "1",
-            "-i", &image_path,
-            "-i", &audio_path,
-            "-vf", &filter_chain,
-            "-c:v", "libx264",
-            "-preset", "fast",
-            "-crf", "18",
-            "-pix_fmt", "yuv420p",
-            "-movflags", "+faststart",
-            "-c:a", "aac",
-            "-b:a", "192k",
-            "-r", "60",
-            "-fps_mode", "vfr",  // Replaces deprecated -vsync
-            "-t", &duration_str,
-            &output_path,
-        ])
-        .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
-        ))),
+    // FFmpeg for silent video
+    let mut child = Command::new("ffmpeg")
+        .arg("-y")
+        .arg("-f").arg("rawvideo")
+        .arg("-pix_fmt").arg("bgr24")
+        .arg("-s").arg(format!("{}x{}", width, height))
+        .arg("-r").arg(fps.to_string())
+        .arg("-i").arg("-")
+        .arg("-an")
+        .arg("-c:v").arg("libx264")
+        .arg("-preset").arg(ffmpeg_preset)
+        .arg("-crf").arg(crf)
+        .arg("-pix_fmt").arg("yuv420p")
+        .arg(&silent_video_path)
+        .stdin(Stdio::piped())
+        .spawn()
+        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to spawn FFmpeg: {}", e)))?;
+
+    let mut stdin = child.stdin.take().unwrap();
+
+    let t0 = Instant::now();
+    let mut frames_sent: usize = 0;
+    let mut prev_visible_sub = String::new();
+    let mut last_pen_x: i32 = 0;
+    let mut last_pen_y: i32 = 0;
+
+    for &idx_in_full in visible_indices.iter() {
+        let visible_sub = &full_text[0..=idx_in_full];
+        if visible_sub != prev_visible_sub {
+            let lines: Vec<&str> = visible_sub.split('\n').collect();
+            let last_line = lines.last().unwrap();
+            let line_idx = lines.len() - 1;
+            let is_header = line_styles[line_idx];
+            let font_scale = if is_header { header_font_scale } else { default_font_scale };
+            let thickness = if is_header { header_thickness } else { default_thickness };
+            let mut base_line = 0;
+            let size = get_text_size(last_line, font, font_scale, thickness, &mut base_line).unwrap();
+            let w = size.width;
+            let h = size.height;
+            let pen_x = margin_x + w + 5;
+            let pen_y = y_positions[line_idx] + h / 2;
+            last_pen_x = pen_x;
+            last_pen_y = pen_y;
+
+            for anim_step in 0..animation_frames_per_char {
+                let anim_offset = (anim_step as f64) / (animation_frames_per_char as f64);
+                let frame_img = render_frame(
+                    visible_sub,
+                    pen_x,
+                    pen_y,
+                    anim_offset,
+                    width as i32,
+                    height as i32,
+                    &wrapped_lines,
+                    &line_styles,
+                    &y_positions,
+                    margin_x,
+                    font,
+                    default_font_scale,
+                    header_font_scale,
+                    default_thickness,
+                    header_thickness,
+                    default_text_color,
+                    header_text_color,
+                    bg_color,
+                    pen_color,
+                    pen_tip_radius,
+                    pen_length,
+                    pen_thickness,
+                    pen_base_angle,
+                    pen_movement_amplitude,
+                )?;
+                stdin.write_all(frame_img.data_bytes().unwrap())
+                    .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("Failed to write frame: {}", e)))?;
+                frames_sent += 1;
+            }
+            prev_visible_sub = visible_sub.to_string();
+        }
+    }
+
+    drop(stdin);
+    child.wait().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("FFmpeg failed: {}", e)))?;
+
+    let elapsed = t0.elapsed().as_secs_f64();
+    println!("Frames sent: {}, elapsed time: {:.3} seconds", frames_sent, elapsed);
+
+    if !Path::new(&silent_video_path).exists() {
+        println!("Silent video generation failed.");
+        return Ok(None);
     }
+
+    let rendered_duration = frames_sent as f64 / fps;
+    println!("Rendered video duration: {:.3}s, Audio duration: {:.3}s", rendered_duration, duration);
+
+    let multiplier = if rendered_duration > 0.0 && duration > 0.0 {
+        duration / rendered_duration
+    } else {
+        1.0
+    };
+
+    // Combine with audio
+    let mut combine_child = Command::new("ffmpeg")
+        .arg("-y")
+        .arg("-i").arg(&silent_video_path)
+        .arg("-i").arg(&audio_path)
+        .arg("-filter:v").arg(format!("setpts=PTS*{}", multiplier))
+        .arg("-c:v").arg("libx264")
+        .arg("-preset").arg("ultrafast")
+        .arg("-crf").arg("28")
+        .arg("-pix_fmt").arg("yuv420p")
+        .arg("-c:a").arg("aac")
+        .arg("-shortest")
+        .arg("-map").arg("0:v:0")
+        .arg("-map").arg("1:a:0")
+        .arg(&final_video_path)
+        .spawn()
+        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to spawn FFmpeg for combine: {}", e)))?;
+
+    combine_child.wait().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("FFmpeg combine failed: {}", e)))?;
+
+    // Clean up
+    std::fs::remove_file(&silent_video_path)
+        .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("Failed to remove silent video: {}", e)))?;
+
+    Ok(Some(final_video_path))
+}
+
+fn wrap_text_cv(text: &str, font: i32, default_font_scale: f64, default_thickness: i32, max_width: i32) -> (Vec<String>, Vec<bool>) {
+    let mut wrapped_lines: Vec<String> = Vec::new();
+    let mut styles: Vec<bool> = Vec::new();
+    for para in text.lines() {
+        let trimmed = para.trim_start();
+        let mut is_header = trimmed.starts_with("###");
+        let mut para_str = para.to_string();
+        if is_header {
+            para_str = trimmed[3..].trim().to_string();
+        }
+        let font_scale = if is_header { header_font_scale } else { default_font_scale };  // Use globals or pass
+        let thickness = if is_header { header_thickness } else { default_thickness };
+        if para_str.is_empty() {
+            wrapped_lines.push("".to_string());
+            styles.push(false);
+            continue;
+        }
+        let words: Vec<&str> = para_str.split_whitespace().collect();
+        let mut cur = String::new();
+        for &w in &words {
+            let candidate = if cur.is_empty() { w.to_string() } else { format!("{} {}", cur, w) };
+            let mut base_line = 0;
+            let size = get_text_size(&candidate, font, font_scale, thickness, &mut base_line).unwrap();
+            if size.width <= max_width {
+                cur = candidate;
+            } else {
+                if !cur.is_empty() {
+                    wrapped_lines.push(cur);
+                    styles.push(is_header);
+                    cur = String::new();
+                }
+                let mut base_line_single = 0;
+                let size_single = get_text_size(w, font, font_scale, thickness, &mut base_line_single).unwrap();
+                if size_single.width > max_width {
+                    let mut chunk = String::new();
+                    for ch in w.chars() {
+                        let cand2 = format!("{}{}", chunk, ch);
+                        let mut base_line_ch = 0;
+                        let size_ch = get_text_size(&cand2, font, font_scale, thickness, &mut base_line_ch).unwrap();
+                        if size_ch.width <= max_width {
+                            chunk = cand2;
+                        } else {
+                            wrapped_lines.push(chunk);
+                            styles.push(is_header);
+                            chunk = ch.to_string();
+                        }
+                    }
+                    if !chunk.is_empty() {
+                        cur = chunk;
+                    }
+                } else {
+                    cur = w.to_string();
+                }
+            }
+        }
+        if !cur.is_empty() {
+            wrapped_lines.push(cur);
+            styles.push(is_header);
+        }
+    }
+    (wrapped_lines, styles)
+}
+
+fn render_frame(
+    visible_text: &str,
+    pen_x: i32,
+    pen_y: i32,
+    anim_offset: f64,
+    width: i32,
+    height: i32,
+    line_styles: &Vec<bool>,
+    y_positions: &Vec<i32>,
+    margin_x: i32,
+    font: i32,
+    default_font_scale: f64,
+    header_font_scale: f64,
+    default_thickness: i32,
+    header_thickness: i32,
+    default_text_color: Scalar,
+    header_text_color: Scalar,
+    bg_color: Scalar,
+    pen_color: Scalar,
+    pen_tip_radius: i32,
+    pen_length: i32,
+    pen_thickness: i32,
+    pen_base_angle: i32,
+    pen_movement_amplitude: i32,
+) -> PyResult<Mat> {
+    let mut img = Mat::new_rows_cols_with_default(height, width, CV_8UC3, bg_color)
+        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to create Mat: {}", e)))?;
+
+    let lines: Vec<&str> = visible_text.split('\n').collect();
+    for (idx, &line) in lines.iter().enumerate() {
+        let is_header = line_styles[idx];
+        let font_scale = if is_header { header_font_scale } else { default_font_scale };
+        let thickness = if is_header { header_thickness } else { default_thickness };
+        let color = if is_header { header_text_color } else { default_text_color };
+        let x = margin_x;
+        let y = y_positions[idx];
+        let mut base_line = 0;
+        let size = get_text_size(line, font, font_scale, thickness, &mut base_line).unwrap();
+        let h = size.height;
+        let y_draw = y + h;
+        if !line.is_empty() {
+            put_text(&mut img, line, Point::new(x, y_draw), font, font_scale, color, thickness, LineTypes::LINE_AA as i32, false)
+                .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to put text: {}", e)))?;
+        }
+    }
+
+    if pen_x > 0 {
+        let offset_y = (pen_movement_amplitude as f64 * (anim_offset * PI).sin()) as i32;
+        let pen_tip_y = pen_y + offset_y;
+        let angle_rad = (pen_base_angle as f64).to_radians();
+        let pen_end_x = pen_x + (pen_length as f64 * angle_rad.cos()) as i32;
+        let pen_end_y = pen_tip_y - (pen_length as f64 * angle_rad.sin()) as i32;
+        line(&mut img, Point::new(pen_x, pen_tip_y), Point::new(pen_end_x, pen_end_y), pen_color, pen_thickness, LineTypes::LINE_8 as i32, 0)
+            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to draw line: {}", e)))?;
+        circle(&mut img, Point::new(pen_x, pen_tip_y), pen_tip_radius, pen_color, -1, LineTypes::LINE_8 as i32, 0)
+            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to draw circle: {}", e)))?;
+    }
+
+    Ok(img)
 }
 
 #[pymodule]
 fn rust_highlight(_py: Python, m: &PyModule) -> PyResult<()> {
-    m.add_function(wrap_pyfunction!(render_video, m)?)?;
+    m.add_function(wrap_pyfunction!(generate_video_clip, m)?)?;
     Ok(())
 }