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::time::{Duration, Instant};
use byteorder::{LittleEndian, ReadBytesExt};
use std::process::{Command, Stdio};

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 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.clear();
    for i in 0..sc {
        let log_pos = (i as f32 / sc as f32).powf(log_power * 0.8);
        let index = (log_pos * (fft_size / 2) as f32) as usize;
        freq_mapping.push(index.min(fft_size / 2 - 1));
    }

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

#[derive(Parser)]
#[command(name = "Spectrust")]
#[command(version = "0.1.0")]
#[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: spectrust -p\n\
                       - From audio file: spectrust -i audio.pcm\n\
                       - From stdin: pw-record --raw - | spectrust")]
struct Cli {
    #[arg(value_name = "FPS", default_value_t = 60)]
    fps: u16,

    #[arg(short, long, help = "Use PipeWire to capture audio directly")]
    pipewire: bool,

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

    #[arg(short, long, help = "Logarithmic scaling power (higher values emphasize lower frequencies)", default_value_t = 1.4)]
    log_power: f32,

    #[arg(short, long, help = "Drop-off factor for bar animation (0.0-1.0)", default_value_t = 0.75)]
    drop_off: f32,
}

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 = 48000;
    let bit_depth: usize = 16;
    let num_channels: usize = 2;
    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 input_buffer = vec![0u8; samples_per_frame * num_channels * (bit_depth / 8)];
    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();

    enum AudioSource {
        StdIn(std::io::StdinLock<'static>),
        File(File),
        PipeWire(std::process::ChildStdout),
    }

    let _ = out.flush();

    let mut audio_source = if cli.pipewire {
        let process = Command::new("pw-record")
            .args(["--format=s16", "--rate=48000", "--channels=2", "--raw", "-"])
            .stdout(Stdio::piped())
            .spawn()
            .expect("failed to start pw-record");
        let pipe = process.stdout.expect("failed to capture pw-record stdout");
        let _ = out.flush();
        AudioSource::PipeWire(pipe)
    } else if let Some(filename) = &cli.input {
        let file = File::open(filename).expect("failed to open input file");
        AudioSource::File(file)
    } else {
        AudioSource::StdIn(std::io::stdin().lock())
    };

    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 drop_factor: f32 = cli.drop_off.clamp(0.0, 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> = (0..sc)
        .map(|i| {
            let log_pos = (i as f32 / sc as f32).powf(cli.log_power * 0.8);
            ((log_pos * (fft_size / 2) as f32) as usize).min(fft_size / 2 - 1)
        })
        .collect();

    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 {
        if event::poll(Duration::ZERO).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::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,
                        cli.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,
                        cli.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();
            let mut buffer_is_empty = false;

            loop {
                let read_result = match &mut audio_source {
                    AudioSource::StdIn(input) => input.read(&mut input_buffer),
                    AudioSource::File(file) => file.read(&mut input_buffer),
                    AudioSource::PipeWire(pipe) => pipe.read(&mut input_buffer),
                };

                match read_result {
                    Ok(0) => {
                        buffer_is_empty = true;
                        break;
                    }
                    Ok(n) => {
                        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 {
                                all_samples.push((frame_sum / valid_channels as f32) / 32768.0);
                            }
                        }
                        if all_samples.len() >= fft_size * 2 {
                            break;
                        }
                    }
                    Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => break,
                    Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                        buffer_is_empty = true;
                        break;
                    }
                    Err(e) => {
                        let error_msg = format!("Error: {}", e);
                        for (i, c) in error_msg.chars().enumerate() {
                            if i < sc {
                                next_buffer[0][i] = Cell { character: c, filled: true };
                            }
                        }
                        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);
                        let _ = out.flush();
                        break;
                    }
                }
            }

            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).powf(0.5);
                        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() {
                    render_bar(&mut next_buffer, i, sr, (magnitude / max_magnitude) * sr as f32);
                }

                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 if buffer_is_empty {
                for i in 0..sc {
                    last_bars[i] *= drop_factor;
                    render_bar(&mut next_buffer, i, sr, (last_bars[i] / max_magnitude) * sr as f32);
                }

                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();
}