TalkEdit/src-tauri/src/caption_generator.rs

use std::process::Command;
use serde_json;
use serde::{Deserialize, Serialize};

#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct Word {
    pub word: String,
    pub start: f64,
    pub end: f64,
    pub confidence: f64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub speaker: Option<String>,
}

#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct CaptionStyle {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub font_name: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub font_size: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub font_color: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub bold: Option<bool>,
}

#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct CaptionContent {
    pub content: String,
}

#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct SaveCaptionsResult {
    pub output_path: String,
}

/// Generate SRT caption content
pub fn generate_srt(
    words: &[Word],
    deleted_indices: Option<&std::collections::HashSet<usize>>,
    words_per_line: usize,
) -> Result<String, String> {
    let python_exe = crate::paths::python_exe();
    let python_exe = python_exe.to_str().unwrap_or_default();
    let script_path = crate::paths::backend_script("caption_generator.py");
    let script_path = script_path.to_str().unwrap_or_default();

    let words_json = serde_json::to_string(words)
        .map_err(|e| format!("Failed to serialize words: {}", e))?;

    let deleted_json = match deleted_indices {
        Some(indices) => serde_json::to_string(indices)
            .map_err(|e| format!("Failed to serialize deleted indices: {}", e))?,
        None => "null".to_string(),
    };

    let output = Command::new(python_exe)
        .args(&[script_path, "generate_srt", &words_json, &deleted_json, &words_per_line.to_string()])
        .output()
        .map_err(|e| format!("Failed to run Python script: {}", e))?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Err(format!("Python script failed: {}", stderr));
    }

    let stdout = String::from_utf8_lossy(&output.stdout);
    let result: CaptionContent = serde_json::from_str(&stdout.trim())
        .map_err(|e| format!("Failed to parse JSON: {}", e))?;

    Ok(result.content)
}

/// Generate VTT caption content
pub fn generate_vtt(
    words: &[Word],
    deleted_indices: Option<&std::collections::HashSet<usize>>,
    words_per_line: usize,
) -> Result<String, String> {
    let python_exe = crate::paths::python_exe();
    let python_exe = python_exe.to_str().unwrap_or_default();
    let script_path = crate::paths::backend_script("caption_generator.py");
    let script_path = script_path.to_str().unwrap_or_default();

    let words_json = serde_json::to_string(words)
        .map_err(|e| format!("Failed to serialize words: {}", e))?;

    let deleted_json = match deleted_indices {
        Some(indices) => serde_json::to_string(indices)
            .map_err(|e| format!("Failed to serialize deleted indices: {}", e))?,
        None => "null".to_string(),
    };

    let output = Command::new(python_exe)
        .args(&[script_path, "generate_vtt", &words_json, &deleted_json, &words_per_line.to_string()])
        .output()
        .map_err(|e| format!("Failed to run Python script: {}", e))?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Err(format!("Python script failed: {}", stderr));
    }

    let stdout = String::from_utf8_lossy(&output.stdout);
    let result: CaptionContent = serde_json::from_str(&stdout.trim())
        .map_err(|e| format!("Failed to parse JSON: {}", e))?;

    Ok(result.content)
}

/// Generate ASS subtitle content
pub fn generate_ass(
    words: &[Word],
    deleted_indices: Option<&std::collections::HashSet<usize>>,
    words_per_line: usize,
    style: Option<&CaptionStyle>,
) -> Result<String, String> {
    let python_exe = crate::paths::python_exe();
    let python_exe = python_exe.to_str().unwrap_or_default();
    let script_path = crate::paths::backend_script("caption_generator.py");
    let script_path = script_path.to_str().unwrap_or_default();

    let words_json = serde_json::to_string(words)
        .map_err(|e| format!("Failed to serialize words: {}", e))?;

    let deleted_json = match deleted_indices {
        Some(indices) => serde_json::to_string(indices)
            .map_err(|e| format!("Failed to serialize deleted indices: {}", e))?,
        None => "null".to_string(),
    };

    let style_json = match style {
        Some(s) => serde_json::to_string(s)
            .map_err(|e| format!("Failed to serialize style: {}", e))?,
        None => "null".to_string(),
    };

    let output = Command::new(python_exe)
        .args(&[script_path, "generate_ass", &words_json, &deleted_json, &words_per_line.to_string(), &style_json])
        .output()
        .map_err(|e| format!("Failed to run Python script: {}", e))?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Err(format!("Python script failed: {}", stderr));
    }

    let stdout = String::from_utf8_lossy(&output.stdout);
    let result: CaptionContent = serde_json::from_str(&stdout.trim())
        .map_err(|e| format!("Failed to parse JSON: {}", e))?;

    Ok(result.content)
}

/// Save caption content to file
pub fn save_captions(content: &str, output_path: &str) -> Result<String, String> {
    let python_exe = crate::paths::python_exe();
    let python_exe = python_exe.to_str().unwrap_or_default();
    let script_path = crate::paths::backend_script("caption_generator.py");
    let script_path = script_path.to_str().unwrap_or_default();

    let output = Command::new(python_exe)
        .args(&[script_path, "save_captions", content, output_path])
        .output()
        .map_err(|e| format!("Failed to run Python script: {}", e))?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Err(format!("Python script failed: {}", stderr));
    }

    let stdout = String::from_utf8_lossy(&output.stdout);
    let result: SaveCaptionsResult = serde_json::from_str(&stdout.trim())
        .map_err(|e| format!("Failed to parse JSON: {}", e))?;

    Ok(result.output_path)
}