sing_detect/detect.rs

//! Singing-segment detection.
//!
//! Rust owns the fast/IO parts (audio extraction via ffmpeg, progress events,
//! and all file exports). The Python sidecar owns *only* the ML step
//! (inaSpeechSegmenter + TensorFlow), which has no native equivalent.
//!
//! Protocol: the sidecar prints one JSON object per line to stdout:
//!   {"type":"progress","message":"…"}
//!   {"type":"result","segments":[{index,start,end,duration}, …],
//!                     "stats":{"music":123.4,"speech":456.7,…}}
//!   {"type":"error","message":"…"}

use std::io::{BufRead, BufReader};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

use serde::{Deserialize, Serialize};
use serde_json::Value;
use tauri::{AppHandle, Emitter};

use crate::tools::hide_window;
use crate::AppState;

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Segment {
    pub index: u32,
    pub start: f64,
    pub end: f64,
    pub duration: f64,
}

#[derive(Clone, Serialize)]
pub struct DetectResult {
    pub segments: Vec<Segment>,
    pub stats: Value,
}

#[derive(Clone, Serialize)]
struct DetectProgress {
    message: String,
}

fn emit(app: &AppHandle, message: impl Into<String>) {
    let _ = app.emit("detect-progress", DetectProgress { message: message.into() });
}

/// Locate the sidecar script: bundled resources first, dev tree as fallback.
fn sidecar_script() -> Option<PathBuf> {
    if let Ok(p) = std::env::var("STREAM_STUDIO_SIDECAR") {
        let pb = PathBuf::from(p);
        if pb.exists() {
            return Some(pb);
        }
    }
    if let Ok(exe) = std::env::current_exe() {
        if let Some(base) = exe.parent() {
            for sub in ["sidecar", "resources/sidecar", "../sidecar"] {
                let cand = base.join(sub).join("detect.py");
                if cand.exists() {
                    return Some(cand);
                }
            }
        }
    }
    // Dev fallback: <crate>/../sidecar/detect.py
    let dev = Path::new(env!("CARGO_MANIFEST_DIR")).join("../sidecar/detect.py");
    if dev.exists() {
        return Some(dev);
    }
    None
}

/// Build the python invocation. Prefers the Windows `py -3.10` launcher
/// (TensorFlow does not yet support 3.12+), then a venv, then plain `python`.
fn python_command(script: &Path) -> Command {
    if let Ok(custom) = std::env::var("STREAM_STUDIO_PYTHON") {
        let mut c = Command::new(custom);
        c.arg(script);
        return c;
    }
    #[cfg(windows)]
    {
        // The py launcher lets us pin 3.10 regardless of the default install.
        let mut c = Command::new("py");
        c.args(["-3.10"]).arg(script);
        return c;
    }
    #[allow(unreachable_code)]
    {
        let mut c = Command::new("python");
        c.arg(script);
        c
    }
}

/// Extract 16 kHz mono WAV (the format inaSpeechSegmenter expects).
fn extract_audio(app: &AppHandle, ffmpeg: &str, video: &str, wav: &Path) -> Result<(), String> {
    emit(app, "Extracting audio (16 kHz mono)…");
    let mut cmd = Command::new(ffmpeg);
    cmd.args([
        "-y", "-i", video,
        "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1",
        "-hide_banner", "-loglevel", "error",
        &wav.to_string_lossy(),
    ]);
    cmd.stdout(Stdio::null()).stderr(Stdio::piped());
    hide_window(&mut cmd);
    let out = cmd.output().map_err(|e| format!("ffmpeg audio extract failed: {e}"))?;
    if !out.status.success() {
        return Err(format!("Audio extraction failed:\n{}", String::from_utf8_lossy(&out.stderr)));
    }
    Ok(())
}

pub fn run(
    app: &AppHandle,
    state: &AppState,
    ffmpeg: &str,
    video: String,
    gap: f64,
    min_duration: f64,
    padding: f64,
) -> Result<DetectResult, String> {
    state.cancel.store(false, std::sync::atomic::Ordering::SeqCst);

    let script = sidecar_script().ok_or(
        "Detection sidecar not found. Expected sidecar/detect.py next to the app.",
    )?;

    // Temp WAV beside the source video.
    let wav = Path::new(&video)
        .with_extension("")
        .to_string_lossy()
        .to_string();
    let wav = PathBuf::from(format!("{wav}.stream_studio.wav"));

    extract_audio(app, ffmpeg, &video, &wav)?;

    let mut cmd = python_command(&script);
    cmd.args([
        "--wav", &wav.to_string_lossy(),
        "--gap", &gap.to_string(),
        "--min-duration", &min_duration.to_string(),
        "--padding", &padding.to_string(),
    ]);
    cmd.stdout(Stdio::piped()).stderr(Stdio::piped());
    cmd.env("PYTHONUNBUFFERED", "1").env("PYTHONIOENCODING", "utf-8");
    hide_window(&mut cmd);

    emit(app, "Loading detection model…");
    let mut child = cmd.spawn().map_err(|e| {
        format!("Could not start Python sidecar ({e}). Is Python 3.10 + inaSpeechSegmenter installed?")
    })?;
    let stdout = child.stdout.take().ok_or("No sidecar stdout")?;
    let stderr = child.stderr.take();
    *state.proc.lock().unwrap() = Some(child);

    let mut result: Option<DetectResult> = None;
    let mut err_msg: Option<String> = None;

    for line in BufReader::new(stdout).lines() {
        let line = match line {
            Ok(l) => l,
            Err(_) => break,
        };
        let line = line.trim();
        if line.is_empty() {
            continue;
        }
        let parsed: Value = match serde_json::from_str(line) {
            Ok(v) => v,
            Err(_) => continue, // ignore stray non-JSON output
        };
        match parsed.get("type").and_then(|t| t.as_str()) {
            Some("progress") => {
                if let Some(m) = parsed.get("message").and_then(|m| m.as_str()) {
                    emit(app, m);
                }
            }
            Some("result") => {
                let segments: Vec<Segment> = serde_json::from_value(
                    parsed.get("segments").cloned().unwrap_or(Value::Array(vec![])),
                )
                .unwrap_or_default();
                let stats = parsed.get("stats").cloned().unwrap_or(Value::Null);
                result = Some(DetectResult { segments, stats });
            }
            Some("error") => {
                err_msg = parsed
                    .get("message")
                    .and_then(|m| m.as_str())
                    .map(|s| s.to_string());
            }
            _ => {}
        }
    }

    let mut child = state.proc.lock().unwrap().take().ok_or("Sidecar vanished")?;
    let status = child.wait().map_err(|e| format!("sidecar wait failed: {e}"))?;
    let mut stderr_text = String::new();
    if let Some(mut se) = stderr {
        use std::io::Read;
        let _ = se.read_to_string(&mut stderr_text);
    }

    let _ = std::fs::remove_file(&wav);

    if state.cancel.load(std::sync::atomic::Ordering::SeqCst) {
        return Err("__cancelled__".into());
    }
    if let Some(m) = err_msg {
        return Err(m);
    }
    if let Some(r) = result {
        return Ok(r);
    }
    if !status.success() {
        let tail: String = stderr_text.chars().rev().take(800).collect::<String>().chars().rev().collect();
        return Err(format!("Detection failed.\n{tail}"));
    }
    Err("Detection produced no result.".into())
}