spectrust/src/main.rs

use clap::Parser;
use crossterm::{
    cursor::{Hide, MoveTo},
    event::{self, Event, KeyCode, KeyEvent, KeyModifiers},
    execute, queue,
    style::Print,
    terminal::{
        self, Clear, ClearType, EnterAlternateScreen, LeaveAlternateScreen,
        disable_raw_mode, enable_raw_mode,
    },
};
use rustfft::{algorithm::Radix4, num_complex::Complex, Fft, FftDirection};
use std::collections::VecDeque;
use std::io::{stdout, Read, Write};
use std::fs::File;
use std::sync::mpsc;
use std::thread;
use std::time::{Duration, Instant};
use byteorder::{LittleEndian, ReadBytesExt};

mod pipewire_capture;

const EMPTY_CELL: Cell = Cell { character: ' ', filled: false };
const BLOCK_CELL: Cell = Cell { character: '█', filled: true };
const SUB_BLOCKS: [char; 7] = ['▁', '▂', '▃', '▄', '▅', '▆', '▇'];

#[derive(Clone, Copy, PartialEq, Debug)]
struct Cell {
    character: char,
    filled: bool,
}

struct ViewDiff {
    changes: Vec<(u16, u16, char)>,
}

impl ViewDiff {
    fn new() -> Self {
        ViewDiff { changes: Vec::with_capacity(512) }
    }

    fn clear(&mut self) {
        self.changes.clear();
    }

    fn add_change(&mut self, y: usize, x: usize, character: char) {
        self.changes.push((y as u16, x as u16, character));
    }

    fn apply(&self) {
        let mut out = stdout();
        for &(y, x, ch) in &self.changes {
            let _ = queue!(out, MoveTo(x, y), Print(ch));
        }
    }
}

fn render_bar(buf: &mut Vec<Vec<Cell>>, col: usize, sr: usize, bar_height_f: f32) {
    let full = (bar_height_f as usize).min(sr);
    for j in 0..full {
        buf[sr - 1 - j][col] = BLOCK_CELL;
    }
    let sub_idx = (bar_height_f.fract() * 8.0) as usize;
    if sub_idx > 0 && full < sr {
        buf[sr - 1 - full][col] = Cell { character: SUB_BLOCKS[sub_idx - 1], filled: true };
    }
}

fn build_freq_mapping(sc: usize, log_power: f32, fft_size: usize) -> Vec<usize> {
    (0..sc)
        .map(|i| {
            let log_pos = (i as f32 / sc as f32).powf(log_power * 0.8);
            ((log_pos * (fft_size / 2) as f32) as usize).min(fft_size / 2 - 1)
        })
        .collect()
}

fn handle_terminal_resize(
    rows: &mut u16,
    cols: &mut u16,
    new_rows: u16,
    new_cols: u16,
    max_width: u16,
    max_height: u16,
    safe_rows: &mut u16,
    safe_cols: &mut u16,
    current_buffer: &mut Vec<Vec<Cell>>,
    next_buffer: &mut Vec<Vec<Cell>>,
    last_bars: &mut Vec<f32>,
    freq_bands: &mut Vec<f32>,
    freq_mapping: &mut Vec<usize>,
    log_power: f32,
    fft_size: usize,
) {
    *rows = new_rows;
    *cols = new_cols;
    *safe_cols = new_cols.min(max_width);
    *safe_rows = new_rows.min(max_height);

    let sc = *safe_cols as usize;
    let sr = *safe_rows as usize;

    *current_buffer = vec![vec![EMPTY_CELL; sc]; sr];
    *next_buffer = vec![vec![EMPTY_CELL; sc]; sr];

    last_bars.clear();
    last_bars.resize(sc, 0.0);
    freq_bands.clear();
    freq_bands.resize(sc, 0.0);

    *freq_mapping = build_freq_mapping(sc, log_power, fft_size);

    let mut out = stdout();
    let _ = queue!(out, Clear(ClearType::All));
    let _ = out.flush();
}

#[derive(Parser)]
#[command(name = "Spectrust")]
#[command(version)]
#[command(about = "A PCM spectrum analyzer for audio visualization",
          long_about = "Spectrust is a PCM spectrum analyzer that visualizes audio in your terminal.\n\
                       It can read from stdin, a file, or capture directly from PipeWire.\n\n\
                       Examples:\n\
                       - Direct PipeWire capture (native, low-latency): spectrust -p\n\
                       - From audio file: spectrust -i audio.pcm\n\
                       - From stdin: pw-record --raw - | spectrust\n\n\
                       Runtime controls:\n\
                       - q / Ctrl+C  quit\n\
                       - [ / ]       decrease / increase drop-off\n\
                       - , / .       decrease / increase log scaling\n\
                       - - / +       decrease / increase sensitivity")]
struct Cli {
    #[arg(short, long, help = "Target frames per second", default_value_t = 60)]
    fps: u16,

    #[arg(short, long, help = "Capture audio directly via native low-latency PipeWire (default sink monitor)")]
    pipewire: bool,

    #[arg(long, help = "PipeWire capture quantum in frames — lower is lower latency (more CPU/xrun risk)", default_value_t = 256)]
    latency: u32,

    #[arg(short, long, help = "Input file (reads from stdin if not provided)")]
    input: Option<String>,

    #[arg(short, long, help = "Logarithmic frequency scaling power (0.5–3.0; higher values give more space to bass)", default_value_t = 1.4)]
    log_power: f32,

    #[arg(short, long, help = "Bar drop-off factor per frame (0.0 = instant drop, 1.0 = no drop)", default_value_t = 0.75)]
    drop_off: f32,

    #[arg(long, help = "Sample rate in Hz — must match the audio source", default_value_t = 48000)]
    rate: u32,

    #[arg(long, help = "Number of audio channels — must match the audio source", default_value_t = 2)]
    channels: u8,
}

fn main() {
    let cli = Cli::parse();

    enable_raw_mode().expect("failed to enable raw mode");
    let mut out = stdout();
    execute!(out, EnterAlternateScreen, Hide).expect("failed to initialize terminal");

    let default_hook = std::panic::take_hook();
    std::panic::set_hook(Box::new(move |info| {
        let _ = disable_raw_mode();
        let _ = execute!(stdout(), LeaveAlternateScreen);
        default_hook(info);
    }));

    let (mut cols, mut rows) = terminal::size().expect("failed to get terminal size");

    let sample_rate: usize = cli.rate as usize;
    let bit_depth: usize = 16;
    let num_channels: usize = cli.channels as usize;
    let samples_per_frame: usize = sample_rate / usize::from(cli.fps);
    let fft_size: usize = 1024;

    let fft = Radix4::new(fft_size, FftDirection::Forward);

    let mut complex_buffer = vec![Complex { re: 0.0f32, im: 0.0f32 }; fft_size];

    let hann_window: Vec<f32> = (0..fft_size)
        .map(|i| 0.5 * (1.0 - (2.0 * std::f32::consts::PI * i as f32 / fft_size as f32).cos()))
        .collect();

    let _ = out.flush();

    let (tx, rx) = mpsc::channel::<Vec<f32>>();

    if cli.pipewire {
        // Native low-latency capture: pushes mono f32 straight into `tx`.
        if let Err(e) = pipewire_capture::spawn(tx, cli.rate, cli.channels, cli.latency) {
            let _ = execute!(out, LeaveAlternateScreen);
            let _ = disable_raw_mode();
            eprintln!("error: {e}");
            std::process::exit(1);
        }
    } else {
        let reader: Box<dyn Read + Send + 'static> = if let Some(filename) = &cli.input {
            Box::new(File::open(filename).expect("failed to open input file"))
        } else {
            Box::new(std::io::stdin())
        };

        let _audio_thread = thread::spawn(move || {
            let mut reader = reader;
            let mut input_buffer = vec![0u8; samples_per_frame * num_channels * (bit_depth / 8)];
            loop {
                match reader.read(&mut input_buffer) {
                    Ok(0) => break,
                    Ok(n) => {
                        let mut samples = Vec::with_capacity(n / (num_channels * (bit_depth / 8)));
                        let mut cursor = std::io::Cursor::new(&input_buffer[..n]);
                        while cursor.position() < (n as u64 - (num_channels * (bit_depth / 8)) as u64 + 1) {
                            let mut frame_sum = 0.0f32;
                            let mut valid_channels = 0u32;
                            for _ in 0..num_channels {
                                match cursor.read_i16::<LittleEndian>() {
                                    Ok(sample) => {
                                        frame_sum += sample as f32;
                                        valid_channels += 1;
                                    }
                                    Err(_) => break,
                                }
                            }
                            if valid_channels > 0 {
                                samples.push((frame_sum / valid_channels as f32) / 32768.0);
                            }
                        }
                        if !samples.is_empty() && tx.send(samples).is_err() {
                            break;
                        }
                    }
                    Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                        thread::sleep(Duration::from_millis(1));
                    }
                    Err(_) => break,
                }
            }
        });
    }

    let max_width: u16 = 1000;
    let max_height: u16 = 1000;

    let mut safe_cols = cols.min(max_width);
    let mut safe_rows = rows.min(max_height);

    let mut sc = safe_cols as usize;
    let mut sr = safe_rows as usize;

    let mut current_buffer: Vec<Vec<Cell>> = vec![vec![EMPTY_CELL; sc]; sr];
    let mut next_buffer: Vec<Vec<Cell>> = vec![vec![EMPTY_CELL; sc]; sr];

    let mut last_bars: Vec<f32> = vec![0.0; sc];
    let mut freq_bands: Vec<f32> = vec![0.0; sc];
    let mut drop_factor: f32 = cli.drop_off.clamp(0.0, 1.0);
    let mut log_power: f32 = cli.log_power;
    let mut sensitivity: f32 = 1.0;

    let interval = Duration::from_secs_f32(1.0 / f32::from(cli.fps));
    let mut last_frame_time = Instant::now();

    let mut max_magnitudes: VecDeque<f32> = VecDeque::with_capacity(usize::from(cli.fps / 2));

    let mut freq_mapping: Vec<usize> = build_freq_mapping(sc, log_power, fft_size);

    let mut last_terminal_resize = Instant::now();
    let resize_check_interval = Duration::from_millis(500);

    let mut all_samples: Vec<f32> = Vec::with_capacity(fft_size * 4);
    let mut view_diff = ViewDiff::new();
    let mut max_magnitude: f32 = 1.0;

    loop {
        let time_to_next_frame = interval.saturating_sub(last_frame_time.elapsed());
        if event::poll(time_to_next_frame).unwrap_or(false) {
            match event::read() {
                Ok(Event::Key(KeyEvent { code: KeyCode::Char('q'), .. })) => break,
                Ok(Event::Key(KeyEvent {
                    code: KeyCode::Char('c'),
                    modifiers: KeyModifiers::CONTROL,
                    ..
                })) => break,
                Ok(Event::Key(KeyEvent { code: KeyCode::Char('['), .. })) => {
                    drop_factor = (drop_factor - 0.05).max(0.0);
                }
                Ok(Event::Key(KeyEvent { code: KeyCode::Char(']'), .. })) => {
                    drop_factor = (drop_factor + 0.05).min(1.0);
                }
                Ok(Event::Key(KeyEvent { code: KeyCode::Char(','), .. })) => {
                    log_power = (log_power - 0.1).max(0.5);
                    freq_mapping = build_freq_mapping(sc, log_power, fft_size);
                }
                Ok(Event::Key(KeyEvent { code: KeyCode::Char('.'), .. })) => {
                    log_power = (log_power + 0.1).min(3.0);
                    freq_mapping = build_freq_mapping(sc, log_power, fft_size);
                }
                Ok(Event::Key(KeyEvent { code: KeyCode::Char('-'), .. })) => {
                    sensitivity = (sensitivity / 1.2).max(0.1);
                }
                Ok(Event::Key(KeyEvent { code: KeyCode::Char('+' | '='), .. })) => {
                    sensitivity = (sensitivity * 1.2).min(10.0);
                }
                Ok(Event::Resize(new_cols, new_rows)) => {
                    handle_terminal_resize(
                        &mut rows, &mut cols, new_rows, new_cols,
                        max_width, max_height, &mut safe_rows, &mut safe_cols,
                        &mut current_buffer, &mut next_buffer,
                        &mut last_bars, &mut freq_bands, &mut freq_mapping,
                        log_power, fft_size,
                    );
                    sc = safe_cols as usize;
                    sr = safe_rows as usize;
                }
                _ => {}
            }
        }

        if last_terminal_resize.elapsed() >= resize_check_interval {
            if let Ok((new_cols, new_rows)) = terminal::size() {
                if new_rows != rows || new_cols != cols {
                    handle_terminal_resize(
                        &mut rows, &mut cols, new_rows, new_cols,
                        max_width, max_height, &mut safe_rows, &mut safe_cols,
                        &mut current_buffer, &mut next_buffer,
                        &mut last_bars, &mut freq_bands, &mut freq_mapping,
                        log_power, fft_size,
                    );
                    sc = safe_cols as usize;
                    sr = safe_rows as usize;
                }
            }
            last_terminal_resize = Instant::now();
        }

        if last_frame_time.elapsed() >= interval {
            for row in next_buffer.iter_mut() {
                row.fill(EMPTY_CELL);
            }

            all_samples.clear();
            loop {
                match rx.try_recv() {
                    Ok(batch) => all_samples.extend_from_slice(&batch),
                    Err(_) => break,
                }
            }
            // Latency guard: if the audio thread got ahead, drop old samples
            if all_samples.len() > fft_size * 2 {
                let excess = all_samples.len() - fft_size * 2;
                all_samples.drain(..excess);
            }

            if !all_samples.is_empty() {
                if all_samples.len() > fft_size {
                    let spb = all_samples.len() / fft_size;
                    let rem = all_samples.len() % fft_size;
                    for i in 0..fft_size {
                        let start = i * spb;
                        let end = start + spb + if i < rem { 1 } else { 0 };
                        let bin = &all_samples[start..end];
                        complex_buffer[i].re = bin.iter().sum::<f32>() / bin.len() as f32;
                        complex_buffer[i].im = 0.0;
                    }
                } else {
                    for i in 0..fft_size {
                        complex_buffer[i].re = if i < all_samples.len() { all_samples[i] } else { 0.0 };
                        complex_buffer[i].im = 0.0;
                    }
                }

                for i in 0..fft_size {
                    complex_buffer[i].re *= hann_window[i];
                }

                fft.process(&mut complex_buffer);

                for (i, &mapping_idx) in freq_mapping.iter().enumerate() {
                    if mapping_idx < fft_size / 2 {
                        let magnitude = complex_buffer[mapping_idx].norm();
                        let freq_boost = 1.0 + (mapping_idx as f32 / (fft_size / 2) as f32) * 4.0;
                        let scaled = (magnitude * 80.0 * freq_boost).sqrt();
                        freq_bands[i] = if scaled > last_bars[i] { scaled } else { last_bars[i] * drop_factor };
                    }
                }

                std::mem::swap(&mut last_bars, &mut freq_bands);

                let current_max = last_bars.iter().copied().fold(0.0f32, f32::max);
                let current_max = if current_max < 0.001 { 1.0 } else { current_max };

                max_magnitudes.push_front(current_max);
                if max_magnitudes.len() > usize::from(cli.fps) / 2 {
                    max_magnitudes.pop_back();
                }

                max_magnitude = max_magnitudes.iter().copied().fold(0.0f32, f32::max).max(0.1);

                for (i, &magnitude) in last_bars[..sc].iter().enumerate() {
                    let height = ((magnitude / max_magnitude) * sr as f32 * sensitivity).min(sr as f32);
                    render_bar(&mut next_buffer, i, sr, height);
                }

                view_diff.clear();
                for y in 0..sr {
                    for x in 0..sc {
                        if current_buffer[y][x] != next_buffer[y][x] {
                            view_diff.add_change(y, x, next_buffer[y][x].character);
                        }
                    }
                }
                view_diff.apply();
                std::mem::swap(&mut current_buffer, &mut next_buffer);
                last_frame_time = Instant::now();
                let _ = out.flush();
            } else {
                for i in 0..sc {
                    last_bars[i] *= drop_factor;
                    let height = ((last_bars[i] / max_magnitude) * sr as f32 * sensitivity).min(sr as f32);
                    render_bar(&mut next_buffer, i, sr, height);
                }

                view_diff.clear();
                for y in 0..sr {
                    for x in 0..sc {
                        if current_buffer[y][x] != next_buffer[y][x] {
                            view_diff.add_change(y, x, next_buffer[y][x].character);
                        }
                    }
                }
                view_diff.apply();
                std::mem::swap(&mut current_buffer, &mut next_buffer);
                last_frame_time = Instant::now();
                let _ = out.flush();
            }
        }
    }

    let _ = execute!(out, LeaveAlternateScreen);
    let _ = disable_raw_mode();
}