mudmouth/src/workers.rs

use crate::{
    config::Config,
    events::{LogEvent, parse_line},
    kokoro::KokoroClient,
};
use anyhow::Result;
use rodio::{Decoder, OutputStreamBuilder, Sink};
use std::{
    collections::{HashMap, VecDeque},
    fs::File,
    io::{BufRead, BufReader, Cursor, Seek, SeekFrom},
    path::{Path, PathBuf},
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
        mpsc::{self, Receiver, Sender},
    },
    thread,
    time::{Duration, Instant},
};
use tokio::sync::mpsc as tokio_mpsc;

#[derive(Debug, Clone)]
pub enum WorkerEvent {
    Log(LogEvent),
    Activity(String),
    KokoroStatus(KokoroStatus),
    Voices(Vec<String>),
    QueueDepth(usize),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KokoroStatus {
    Checking,
    Connected,
    Unreachable,
}

#[derive(Debug)]
pub struct SynthesisJob {
    pub text: String,
    pub voice: String,
    pub speed: f32,
    pub repeat_key: String,
    pub observed_at: Instant,
}

#[derive(Debug)]
struct RepeatState {
    last_seen: Instant,
    count: usize,
}

#[derive(Debug)]
pub enum AudioCommand {
    Play(Vec<u8>),
    SetVolume(f32),
}

pub fn spawn_log_tailer(
    path: PathBuf,
    sender: Sender<WorkerEvent>,
    generation: Arc<AtomicUsize>,
    expected_generation: usize,
) {
    thread::spawn(move || {
        let mut position = None;
        while generation.load(Ordering::Relaxed) == expected_generation {
            match read_new_lines(&path, &mut position) {
                Ok(lines) => {
                    for line in lines {
                        if let Some(event) = parse_line(&line) {
                            let _ = sender.send(WorkerEvent::Log(event));
                        }
                    }
                }
                Err(error) => {
                    let _ = sender.send(WorkerEvent::Activity(format!("log watch: {error}")));
                }
            }
            thread::sleep(Duration::from_millis(350));
        }
    });
}

fn read_new_lines(path: &Path, position: &mut Option<u64>) -> Result<Vec<String>> {
    let mut file = File::open(path)?;
    let length = file.metadata()?.len();
    let offset = position.unwrap_or(length);
    let offset = if length < offset { 0 } else { offset };
    file.seek(SeekFrom::Start(offset))?;
    let reader = BufReader::new(&file);
    let lines = reader.lines().collect::<std::io::Result<Vec<_>>>()?;
    *position = Some(length);
    Ok(lines)
}

pub fn spawn_audio_worker(volume: f32, event_sender: Sender<WorkerEvent>) -> Sender<AudioCommand> {
    let (sender, receiver): (Sender<AudioCommand>, Receiver<AudioCommand>) = mpsc::channel();
    thread::spawn(move || {
        let stream = match OutputStreamBuilder::open_default_stream() {
            Ok(stream) => stream,
            Err(error) => {
                let _ = event_sender.send(WorkerEvent::Activity(format!(
                    "audio output unavailable: {error}"
                )));
                return;
            }
        };
        let mut volume = volume;
        while let Ok(command) = receiver.recv() {
            let wav = match command {
                AudioCommand::Play(wav) => wav,
                AudioCommand::SetVolume(updated_volume) => {
                    volume = updated_volume;
                    continue;
                }
            };
            let _ = event_sender.send(WorkerEvent::QueueDepth(1));
            match Decoder::try_from(Cursor::new(wav)) {
                Ok(source) => {
                    let sink = Sink::connect_new(stream.mixer());
                    sink.set_volume(volume);
                    sink.append(source);
                    sink.sleep_until_end();
                }
                Err(error) => {
                    let _ = event_sender.send(WorkerEvent::Activity(format!(
                        "audio decode failed: {error}"
                    )));
                }
            }
            let _ = event_sender.send(WorkerEvent::QueueDepth(0));
        }
    });
    sender
}

pub fn spawn_kokoro_worker(
    config: Config,
    mut jobs: tokio_mpsc::Receiver<SynthesisJob>,
    audio_sender: Sender<AudioCommand>,
    event_sender: Sender<WorkerEvent>,
) -> tokio_mpsc::Sender<Config> {
    let (control_sender, mut control_receiver) = tokio_mpsc::channel::<Config>(4);
    tokio::spawn(async move {
        let mut config = config;
        let mut client = match KokoroClient::new(&config.kokoro_url) {
            Ok(client) => client,
            Err(error) => {
                let _ = event_sender.send(WorkerEvent::Activity(format!("Kokoro config: {error}")));
                return;
            }
        };
        let mut voices = Vec::new();
        let mut last_health_check = Instant::now() - Duration::from_secs(30);
        let mut repeats = HashMap::<String, RepeatState>::new();
        let mut queued = VecDeque::new();

        loop {
            while let Ok(updated_config) = control_receiver.try_recv() {
                match KokoroClient::new(&updated_config.kokoro_url) {
                    Ok(updated_client) => {
                        config = updated_config;
                        client = updated_client;
                        voices.clear();
                        last_health_check = Instant::now() - Duration::from_secs(30);
                        let _ = event_sender.send(WorkerEvent::Activity(
                            "Kokoro settings updated; validating endpoint…".into(),
                        ));
                    }
                    Err(error) => {
                        let _ = event_sender.send(WorkerEvent::Activity(format!(
                            "Kokoro endpoint rejected: {error}"
                        )));
                    }
                }
            }

            if last_health_check.elapsed() >= Duration::from_secs(5) {
                last_health_check = Instant::now();
                let _ = event_sender.send(WorkerEvent::KokoroStatus(KokoroStatus::Checking));
                match client.health().await {
                    Ok(()) => match client.voices().await {
                        Ok(found) => {
                            if found != voices {
                                voices = found.clone();
                                let _ = event_sender.send(WorkerEvent::Voices(found));
                            }
                            let _ = event_sender
                                .send(WorkerEvent::KokoroStatus(KokoroStatus::Connected));
                        }
                        Err(error) => {
                            let _ = event_sender
                                .send(WorkerEvent::Activity(format!("voice list failed: {error}")));
                            let _ = event_sender
                                .send(WorkerEvent::KokoroStatus(KokoroStatus::Unreachable));
                        }
                    },
                    Err(_) => {
                        let _ =
                            event_sender.send(WorkerEvent::KokoroStatus(KokoroStatus::Unreachable));
                    }
                }
            }

            while let Ok(job) = jobs.try_recv() {
                if queued.len() >= config.max_queue_length {
                    let _ = event_sender.send(WorkerEvent::Activity(
                        "speech queue full; line dropped".into(),
                    ));
                } else {
                    let mut job = job;
                    job.speed =
                        repeat_speed(&mut repeats, &job.repeat_key, &config, job.observed_at);
                    queued.push_back(job);
                }
            }

            if let Some(job) = queued.pop_front() {
                let _ = event_sender.send(WorkerEvent::QueueDepth(queued.len() + 1));
                match client
                    .synthesize_wav(&job.text, &job.voice, job.speed)
                    .await
                {
                    Ok(wav) => {
                        if audio_sender.send(AudioCommand::Play(wav)).is_err() {
                            let _ = event_sender.send(WorkerEvent::Activity(
                                "audio worker stopped; unable to play speech".into(),
                            ));
                        }
                    }
                    Err(error) => {
                        let _ = event_sender.send(WorkerEvent::Activity(format!(
                            "Kokoro speech failed: {error}"
                        )));
                    }
                }
                let _ = event_sender.send(WorkerEvent::QueueDepth(queued.len()));
            }

            tokio::time::sleep(Duration::from_millis(100)).await;
        }
    });
    control_sender
}

fn repeat_speed(
    repeats: &mut HashMap<String, RepeatState>,
    key: &str,
    config: &Config,
    now: Instant,
) -> f32 {
    let window = Duration::from_secs(config.repeat_window_seconds);
    let state = repeats.entry(key.to_owned()).or_insert(RepeatState {
        last_seen: now,
        count: 0,
    });
    if now.duration_since(state.last_seen) >= window {
        state.count = 0;
    }
    state.last_seen = now;
    state.count += 1;

    match state.count {
        1 => 1.0,
        2 => config.repeat_second_speed.clamp(0.25, 4.0),
        _ => config.repeat_third_plus_speed.clamp(0.25, 4.0),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use tempfile::tempdir;

    #[test]
    fn tailer_starts_at_end_and_recovers_after_truncation() {
        let directory = tempdir().unwrap();
        let path = directory.path().join("eqlog.txt");
        fs::write(&path, "old\n").unwrap();
        let mut position = None;
        assert!(read_new_lines(&path, &mut position).unwrap().is_empty());
        fs::write(&path, "old\nnew\n").unwrap();
        assert_eq!(read_new_lines(&path, &mut position).unwrap(), vec!["new"]);
        fs::write(&path, "fresh\n").unwrap();
        assert_eq!(read_new_lines(&path, &mut position).unwrap(), vec!["fresh"]);
    }

    #[test]
    fn repeated_dialogue_speeds_up_then_resets() {
        let config = Config::default();
        let now = Instant::now();
        let mut repeats = HashMap::new();

        assert_eq!(
            repeat_speed(&mut repeats, "npc:turga\u{1f}hello", &config, now),
            1.0
        );
        assert_eq!(
            repeat_speed(
                &mut repeats,
                "npc:turga\u{1f}hello",
                &config,
                now + Duration::from_secs(1),
            ),
            1.5
        );
        assert_eq!(
            repeat_speed(
                &mut repeats,
                "npc:turga\u{1f}hello",
                &config,
                now + Duration::from_secs(2),
            ),
            2.0
        );
        assert_eq!(
            repeat_speed(
                &mut repeats,
                "npc:turga\u{1f}hello",
                &config,
                now + Duration::from_secs(18),
            ),
            1.0
        );
    }
}