spectrust/src/main.rs

use clap::Parser;
use ncurses::*;
use rustfft::{algorithm::Radix4, num_complex::Complex, Fft, FftDirection};
use std::collections::VecDeque;
use std::io::{stdin, Read};
use std::fs::File;
use std::time::{Duration, Instant};
use byteorder::{LittleEndian, ReadBytesExt};

use std::process::{Command, Stdio};
use std::thread::sleep;

// Create screen buffer struct to track cell states
// Each cell stores a character and its state (whether it's filled or not)
#[derive(Clone, PartialEq)]
struct Cell {
    character: String,
    filled: bool,
}

// Function to handle terminal resizing
fn handle_terminal_resize(
    rows: &mut i32, 
    cols: &mut i32, 
    new_rows: i32, 
    new_cols: i32,
    max_width: i32,
    max_height: i32,
    safe_rows: &mut i32,
    safe_cols: &mut i32,
    empty_cell: &Cell,
    current_buffer: &mut Vec<Vec<Cell>>,
    next_buffer: &mut Vec<Vec<Cell>>,
    last_bars: &mut Vec<f32>,
    freq_mapping: &mut Vec<usize>,
    log_power: f32,
    fft_size: usize
) {
    // Update terminal dimensions
    *rows = new_rows;
    *cols = new_cols;
    
    // Update safe dimensions to account for new terminal size
    *safe_cols = new_cols.min(max_width);
    *safe_rows = new_rows.min(max_height);
    
    // Resize screen buffers
    *current_buffer = vec![vec![empty_cell.clone(); *safe_cols as usize]; *safe_rows as usize];
    *next_buffer = vec![vec![empty_cell.clone(); *safe_cols as usize]; *safe_rows as usize];
    
    // Resize last_bars to fit new screen width
    *last_bars = vec![0.0; *safe_cols as usize];
    
    // Recalculate frequency mapping for new terminal width
    freq_mapping.clear();
    freq_mapping.reserve(*safe_cols as usize);
    for i in 0..*safe_cols as usize {
        // Using a more balanced logarithmic scale to better distribute frequencies
        let log_pos = (i as f32 / *safe_cols 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));
    }
    
    // Completely clear the screen
    clear();
    
    // Apply changes
    refresh();
}

#[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.85)]
    drop_off: f32,
}

fn main() {
    let cli = Cli::parse();
    
    // Initialize ncurses with proper locale for UTF-8 support
    setlocale(LcCategory::all, "");  // Set locale before ncurses init for UTF-8 support
    initscr();
    noecho();
    curs_set(CURSOR_VISIBILITY::CURSOR_INVISIBLE);
    keypad(stdscr(), true);
    // Enable handling of window resize signals
    timeout(0); // Non-blocking input for getch()
    
    // Get terminal dimensions
    let mut rows: i32 = 0;
    let mut cols: i32 = 0;
    getmaxyx(stdscr(), &mut rows, &mut cols);
    
    let sample_rate: usize = 48000;
    let bit_depth: usize = 16;
    let samples_per_frame: usize = sample_rate / usize::from(cli.fps);
    let fft_size: usize = 1024;
    
    // Set up FFT processor
    let fft = Radix4::new(fft_size, FftDirection::Forward);
    
    // Set up audio input source
    let mut input_buffer = vec![0; samples_per_frame * (bit_depth / 8)];
    let mut complex_buffer = vec![Complex { re: 0.0, im: 0.0 }; fft_size];
    
    // Set up input reader based on command line arguments
    enum AudioSource {
        StdIn(std::io::StdinLock<'static>),
        File(File),
        PipeWire(std::process::ChildStdout),
    }
    
    // Initialize display without a title bar
    refresh();
    
    // Set up audio source
    let mut audio_source = if cli.pipewire {
        // Launch PipeWire capture process with raw output format
        let process = Command::new("pw-record")
            .args(["--format=s16", "--rate=48000", "--channels=1", "--raw", "-"])
            .stdout(Stdio::piped())
            .spawn()
            .expect("Failed to start pw-record");
            
        let stdout = process.stdout.expect("Failed to capture pw-record stdout");
        refresh();
        AudioSource::PipeWire(stdout)
    } else if let Some(filename) = &cli.input {
        // Read from file
        let file = File::open(filename).expect("Failed to open input file");
        AudioSource::File(file)
    } else {
        // Read from stdin
        AudioSource::StdIn(stdin().lock())
    };
    
    // Define maximum sizes to avoid overflow
    let max_width = 1000;
    let max_height = 1000;
    
    // Calculate safe dimensions that will be updated when terminal resizes
    let mut safe_cols = cols.min(max_width);
    let mut safe_rows = rows.min(max_height);
    
    // Initialize screen buffers with safe dimensions
    let empty_cell = Cell { character: " ".to_string(), filled: false };
    let mut current_buffer: Vec<Vec<Cell>> = vec![vec![empty_cell.clone(); safe_cols as usize]; safe_rows as usize];
    let mut next_buffer: Vec<Vec<Cell>> = vec![vec![empty_cell.clone(); safe_cols as usize]; safe_rows as usize];
    
    // For smooth drop-off
    let mut last_bars: Vec<f32> = vec![0.0; safe_cols as usize];
    let drop_factor: f32 = cli.drop_off.max(0.0).min(1.0); // Clamp between 0 and 1

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

    let mut last_frame_time = Instant::now();

    // Store max amplitudes for brief stability (but not smoothing)
    let mut max_magnitudes: VecDeque<f32> = VecDeque::with_capacity(usize::from(cli.fps / 2));
    
    // Frequency scaling with safe dimensions - balanced logarithmic mapping for better audio visualization
    let mut freq_mapping: Vec<usize> = Vec::with_capacity(safe_cols as usize);
    for i in 0..safe_cols as usize {
        // Using a more balanced logarithmic scale to better distribute frequencies
        // Lower log_power to give more space to high frequencies
        let log_pos = (i as f32 / safe_cols as f32).powf(cli.log_power * 0.8); // Reduced power for better balance
        let index = (log_pos * (fft_size / 2) as f32) as usize;
        freq_mapping.push(index.min(fft_size / 2 - 1));
    }

    // Track if terminal has been resized
    let mut last_terminal_resize = Instant::now();
    let resize_check_interval = Duration::from_millis(500); // Check resize every 500ms
    
    loop {
        // Check for terminal resize - polling approach for terminals that don't send KEY_RESIZE events
        if last_terminal_resize.elapsed() >= resize_check_interval {
            let mut new_rows = 0;
            let mut new_cols = 0;
            getmaxyx(stdscr(), &mut new_rows, &mut new_cols);
            
            // Handle resize operation
            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, 
                                      &empty_cell, &mut current_buffer, &mut next_buffer, 
                                      &mut last_bars, &mut freq_mapping, cli.log_power, fft_size);
            }
            
            last_terminal_resize = Instant::now();
        }
        
        // Check for character input - direct resize event
        let ch = getch();
        if ch == KEY_RESIZE {
            // Terminal was resized, update dimensions
            let mut new_rows = 0;
            let mut new_cols = 0;
            getmaxyx(stdscr(), &mut new_rows, &mut new_cols);
            
            // Handle resize operation
            handle_terminal_resize(&mut rows, &mut cols, new_rows, new_cols, 
                                  max_width, max_height, &mut safe_rows, &mut safe_cols, 
                                  &empty_cell, &mut current_buffer, &mut next_buffer, 
                                  &mut last_bars, &mut freq_mapping, cli.log_power, fft_size);
        }
        
        // Process audio frames at regular intervals
        if last_frame_time.elapsed() >= interval {
            // Clear the entire screen before rendering the new frame (prevents lingering characters)
            erase();

            // No welcome message - keeping the display clean as requested
            
            // Reset next buffer for this frame - using safe dimensions
            for y in 0..safe_rows as usize {
                for x in 0..safe_cols as usize {
                    next_buffer[y][x] = empty_cell.clone();
                    // Also reset the current buffer to ensure clean state
                    current_buffer[y][x] = empty_cell.clone();
                }
            }
            
            // Read and process all available data for this frame to prevent latency
            let mut all_samples = Vec::new();
            let mut buffer_is_empty = false;
            
            // Process all available data until buffer is empty or we have an error
            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(stdout) => stdout.read(&mut input_buffer),
                };
                
                match read_result {
                    Ok(0) => {
                        // No data read but not EOF - buffer is empty
                        buffer_is_empty = true;
                        break;
                    },
                    Ok(n) => {
                        // Resize input buffer to match actual bytes read
                        let actual_input = &input_buffer[0..n];
                        
                        // Convert raw bytes to audio samples
                        let mut cursor = std::io::Cursor::new(actual_input);
                        
                        // Read as many complete samples as possible
                        while cursor.position() < (n as u64 - 1) {
                            match cursor.read_i16::<LittleEndian>() {
                                Ok(sample) => all_samples.push(sample as f32 / 32768.0),
                                Err(_) => break,
                            }
                        }
                        
                        // If we've collected enough samples, we can stop to avoid too much processing
                        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 currently empty, but might have data later
                        buffer_is_empty = true;
                        break;
                    },
                    Err(e) => {
                        // Error occurred
                        // Show error message in the buffer - using safe dimensions
                        let error_msg = format!("Error: {}", e);
                        for (i, c) in error_msg.chars().enumerate() {
                            if i < safe_cols as usize {
                                next_buffer[0][i] = Cell { 
                                    character: c.to_string(), 
                                    filled: true 
                                };
                            }
                        }
                        
                        // Render only the changed cells - using safe dimensions
                        for y in 0..safe_rows as usize {
                            for x in 0..safe_cols as usize {
                                if current_buffer[y][x] != next_buffer[y][x] {
                                    let _ = mvprintw(y as i32, x as i32, &next_buffer[y][x].character);
                                }
                            }
                        }
                        
                        // Swap buffers
                        std::mem::swap(&mut current_buffer, &mut next_buffer);
                        
                        refresh();
                        break;
                    }
                }
            }
            
            if !all_samples.is_empty() {
                // Process all samples by averaging if we have more than fft_size
                if all_samples.len() > fft_size {
                    // Calculate how many samples to average per FFT bin
                    let samples_per_bin = all_samples.len() / fft_size;
                    let remainder = all_samples.len() % fft_size;
                    
                    // Average samples to fill complex buffer (ensures we don't lose data)
                    for i in 0..fft_size {
                        let start = i * samples_per_bin;
                        let end = if i < remainder {
                            start + samples_per_bin + 1 // Distribute remainder
                        } else {
                            start + samples_per_bin
                        };
                        
                        // Average the samples in this bin
                        let bin_samples = &all_samples[start..end];
                        let avg = bin_samples.iter().sum::<f32>() / bin_samples.len() as f32;
                        
                        complex_buffer[i].re = avg;
                        complex_buffer[i].im = 0.0;
                    }
                } else {
                    // If we have fewer samples than fft_size, use them directly
                    for i in 0..fft_size {
                        if i < all_samples.len() {
                            complex_buffer[i].re = all_samples[i];
                            complex_buffer[i].im = 0.0;
                        } else {
                            complex_buffer[i].re = 0.0;
                            complex_buffer[i].im = 0.0;
                        }
                    }
                }

                // Apply window function (Hann window) to reduce spectral leakage
                for i in 0..fft_size {
                    let window = 0.5 * (1.0 - (2.0 * std::f32::consts::PI * i as f32 / fft_size as f32).cos());
                    complex_buffer[i].re *= window;
                }

                fft.process(&mut complex_buffer);
                
                // Calculate frequency bands using our mapping - with safe dimensions
                let mut freq_bands: Vec<f32> = vec![0.0; safe_cols as usize];
                
                // We only use the first half of FFT results (Nyquist theorem)
                for (i, &mapping_idx) in freq_mapping.iter().enumerate() {
                    if i < freq_bands.len() && mapping_idx < fft_size / 2 {
                        // Get magnitude and apply frequency-dependent scaling for better visualization
                        let magnitude = complex_buffer[mapping_idx].norm();
                        
                        // Boost higher frequencies to make them more visible
                        // The higher the frequency, the more we boost it
                        let freq_boost = 1.0 + (mapping_idx as f32 / (fft_size / 2) as f32) * 4.0;
                        
                        // Amplify the signal with the frequency boost
                        let scaled_magnitude = (magnitude * 80.0 * freq_boost).powf(0.5);
                        
                        // Direct response for punchier visualization
                        freq_bands[i] = if scaled_magnitude > last_bars[i] {
                            // Immediate rise for maximum punch
                            scaled_magnitude
                        } else {
                            // Quick fall with drop factor
                            last_bars[i] * drop_factor
                        };
                    }
                }
                
                // Update last bars for next frame
                last_bars = freq_bands.clone();

                // Calculate max magnitude over all frequencies for normalization
                let current_max = freq_bands.iter().cloned().fold(0.0, f32::max);
                
                // Make sure we have a reasonable minimum value to avoid empty display
                let current_max = if current_max < 0.001 { 1.0 } else { current_max };
                
                // Store this frame's maximum for very minimal smoothing
                max_magnitudes.push_front(current_max);
                
                // Keep just half a second of history - enough to prevent extreme flickering
                // but not enough to create significant smoothing
                if max_magnitudes.len() > usize::from(cli.fps) / 2 {
                    max_magnitudes.pop_back();
                }
                
                // Direct maximum for most punchy visualization
                // Just take the current frame's maximum with minimal safety buffer
                let max_magnitude = if !max_magnitudes.is_empty() {
                    max_magnitudes.iter().cloned().fold(0.0, f32::max).max(0.1)
                } else {
                    current_max.max(0.1)
                };
                
                // Update next buffer with the new bar states - using safe dimensions
                for (i, &magnitude) in freq_bands.iter().enumerate() {
                    if i < safe_cols as usize {
                        let normalized_magnitude = magnitude / max_magnitude;
                        let bar_height = (normalized_magnitude * safe_rows as f32) as i32;
                        
                        for j in 0..safe_rows as i32 {
                            let y = safe_rows - 1 - j;
                            if j < bar_height {
                                // Simple solid bar - one character wide
                                let char_to_use = "▒"; // Using medium shade block (U+2592) for better compatibility
                                
                                if y >= 0 && y < safe_rows && i < safe_cols as usize {
                                    next_buffer[y as usize][i] = Cell { 
                                        character: char_to_use.to_string(), 
                                        filled: true 
                                    };
                                }
                            }
                        }
                    }
                }
                
                // Render only the cells that changed (diff and patch) - using safe dimensions
                for y in 0..safe_rows as usize {
                    for x in 0..safe_cols as usize {
                        if current_buffer[y][x] != next_buffer[y][x] {
                            let _ = mvprintw(y as i32, x as i32, &next_buffer[y][x].character);
                        }
                    }
                }
                
                // Swap buffers for next frame
                std::mem::swap(&mut current_buffer, &mut next_buffer);
                
                last_frame_time = Instant::now();
                refresh();
            } else if buffer_is_empty {
                // Reset next buffer - using safe dimensions
                for y in 0..safe_rows as usize {
                    for x in 0..safe_cols as usize {
                        next_buffer[y][x] = empty_cell.clone();
                    }
                }
                
                // Update buffer with reduced heights - using safe dimensions
                for i in 0..last_bars.len() {
                    if i < safe_cols as usize {
                        last_bars[i] *= drop_factor;
                        
                        let bar_height = (last_bars[i] * safe_rows as f32) as i32;
                        for j in 0..bar_height {
                            let y = safe_rows - 1 - j;
                            if y >= 0 && y < safe_rows && i < safe_cols as usize {
                                // Simple solid bar - one character wide
                                let char_to_use = "▒"; // Using medium shade block (U+2592) for better compatibility
                                
                                next_buffer[y as usize][i] = Cell { 
                                    character: char_to_use.to_string(), 
                                    filled: true 
                                };
                            }
                        }
                    }
                }
                
                // Render only the cells that changed - using safe dimensions
                for y in 0..safe_rows as usize {
                    for x in 0..safe_cols as usize {
                        if current_buffer[y][x] != next_buffer[y][x] {
                            let _ = mvprintw(y as i32, x as i32, &next_buffer[y][x].character);
                        }
                    }
                }
                
                // Swap buffers
                std::mem::swap(&mut current_buffer, &mut next_buffer);
                
                refresh();
            }
        }
    }
}