cyber-rain 0.1.0

use std::{
    io::{self, Write, stdout},
    thread,
    time::{Duration, Instant},
};

use clap::{Parser, ValueEnum};
use crossterm::{
    cursor::{Hide, MoveTo, Show},
    event::{self, Event, KeyCode, KeyEventKind, KeyModifiers},
    execute, queue,
    style::{Color, Print, ResetColor, SetForegroundColor},
    terminal::{self, Clear, ClearType, EnterAlternateScreen, LeaveAlternateScreen},
};
use rand::{Rng, SeedableRng, rngs::StdRng};

#[derive(Parser, Debug)]
#[command(
    name = "cyber-rain",
    version,
    about = "A Rust-powered digital rain terminal visualizer.",
    after_help = "Controls: q/Esc/Ctrl-C quit, Space or p pauses. Try: cyber-rain --palette rainbow --message 'follow the white rabbit'"
)]
struct Cli {
    /// Target frames per second.
    #[arg(long, default_value_t = 60, value_parser = clap::value_parser!(u64).range(1..=240))]
    fps: u64,

    /// Fraction of terminal columns that emit rain.
    #[arg(short, long, default_value_t = 0.82, value_parser = parse_density)]
    density: f32,

    /// Rain speed from 1 (slow) to 10 (fast).
    #[arg(short, long, default_value_t = 7, value_parser = clap::value_parser!(u16).range(1..=10))]
    speed: u16,

    /// Color theme for the rain.
    #[arg(short, long, value_enum, default_value_t = Palette::Green)]
    palette: Palette,

    /// Hide this message inside occasional falling streams.
    #[arg(short, long)]
    message: Option<String>,

    /// Use ASCII glyphs only instead of half-width katakana and symbols.
    #[arg(long)]
    ascii: bool,

    /// Seed the random generator for reproducible runs.
    #[arg(long)]
    seed: Option<u64>,

    /// Draw in the current terminal buffer instead of the alternate screen.
    #[arg(long)]
    no_alt_screen: bool,

    /// Disable occasional high-energy rain bursts.
    #[arg(long)]
    calm: bool,
}

#[derive(Clone, Copy, Debug, ValueEnum)]
enum Palette {
    Green,
    Cyan,
    Amber,
    Crimson,
    Ghost,
    Rainbow,
}

struct Config {
    fps: u64,
    density: f32,
    speed: u16,
    palette: Palette,
    message: Vec<char>,
    ascii: bool,
    seed: Option<u64>,
    alt_screen: bool,
    bursts: bool,
}

impl From<Cli> for Config {
    fn from(cli: Cli) -> Self {
        Self {
            fps: cli.fps,
            density: cli.density,
            speed: cli.speed,
            palette: cli.palette,
            message: cli.message.unwrap_or_default().chars().collect(),
            ascii: cli.ascii,
            seed: cli.seed,
            alt_screen: !cli.no_alt_screen,
            bursts: !cli.calm,
        }
    }
}

struct TerminalSession {
    alt_screen: bool,
}

impl TerminalSession {
    fn enter(stdout: &mut impl Write, alt_screen: bool) -> io::Result<Self> {
        terminal::enable_raw_mode()?;

        if alt_screen {
            execute!(stdout, EnterAlternateScreen, Clear(ClearType::All), Hide)?;
        } else {
            execute!(stdout, Clear(ClearType::All), Hide)?;
        }

        Ok(Self { alt_screen })
    }
}

impl Drop for TerminalSession {
    fn drop(&mut self) {
        let mut stdout = stdout();
        let _ = execute!(stdout, ResetColor, Show);

        if self.alt_screen {
            let _ = execute!(stdout, LeaveAlternateScreen);
        }

        let _ = terminal::disable_raw_mode();
    }
}

struct Matrix {
    config: Config,
    columns: Vec<Column>,
    rng: StdRng,
    width: u16,
    height: u16,
    frame: u64,
    burst_frames: u16,
}

impl Matrix {
    fn new(config: Config, width: u16, height: u16) -> Self {
        let mut rng = match config.seed {
            Some(seed) => StdRng::seed_from_u64(seed),
            None => StdRng::from_entropy(),
        };
        let columns = build_columns(width, height, &config, &mut rng);

        Self {
            config,
            columns,
            rng,
            width,
            height,
            frame: 0,
            burst_frames: 0,
        }
    }

    fn resize(&mut self, width: u16, height: u16, stdout: &mut impl Write) -> io::Result<()> {
        self.width = width;
        self.height = height;
        self.columns = build_columns(width, height, &self.config, &mut self.rng);
        queue!(stdout, Clear(ClearType::All))?;
        Ok(())
    }

    fn draw_frame(&mut self, stdout: &mut impl Write, paused: bool) -> io::Result<()> {
        if paused {
            self.draw_status(stdout, "paused")?;
            return Ok(());
        }

        self.frame = self.frame.wrapping_add(1);

        if self.config.bursts && self.burst_frames == 0 && self.rng.gen_ratio(1, 520) {
            self.burst_frames = self.rng.gen_range(24..=72);
        }

        let burst_active = self.burst_frames > 0;
        let frame = self.frame;
        let config = &self.config;
        let height = self.height;
        let rng = &mut self.rng;

        for column in &mut self.columns {
            let force_step = burst_active && rng.gen_bool(0.38);
            column.tick(stdout, config, height, frame, rng, force_step)?;
        }

        if self.burst_frames > 0 {
            self.burst_frames -= 1;
        }

        self.draw_status(stdout, if burst_active { "burst" } else { "live" })?;
        Ok(())
    }

    fn draw_status(&self, stdout: &mut impl Write, state: &str) -> io::Result<()> {
        if self.height < 2 || self.width < 28 {
            return Ok(());
        }

        let label = format!(
            " cyber-rain | {state} | {:?} | q quit | space pause ",
            self.config.palette
        );
        let clipped: String = label.chars().take(self.width as usize).collect();

        queue!(
            stdout,
            MoveTo(0, self.height - 1),
            SetForegroundColor(Color::DarkGrey),
            Print(format!("{clipped:<width$}", width = self.width as usize)),
            ResetColor
        )?;

        Ok(())
    }
}

struct Column {
    x: u16,
    y: i16,
    trail: u16,
    step_every: u16,
    clock: u16,
    hot: bool,
    glyphs: Vec<char>,
    message: Option<Vec<char>>,
    message_index: usize,
}

impl Column {
    fn random(x: u16, height: u16, config: &Config, rng: &mut StdRng) -> Self {
        let mut column = Self {
            x,
            y: 0,
            trail: 8,
            step_every: 1,
            clock: rng.gen_range(0..=u16::MAX),
            hot: false,
            glyphs: Vec::new(),
            message: None,
            message_index: 0,
        };
        column.respawn(height, config, rng);
        column
    }

    fn respawn(&mut self, height: u16, config: &Config, rng: &mut StdRng) {
        let height = height.max(2);
        let max_trail = (height / 2).clamp(6, 34);
        let base_delay = match config.speed {
            1 => 9,
            2 => 7,
            3 => 6,
            4 => 5,
            5 => 4,
            6 | 7 => 3,
            8 | 9 => 2,
            _ => 1,
        };
        let jitter = rng.gen_range(0..=1);

        self.y = -(rng.gen_range(0..height) as i16);
        self.trail = rng.gen_range(4..=max_trail);
        self.step_every = (base_delay + jitter).max(1);
        self.hot = rng.gen_bool(0.09);
        self.message_index = 0;
        self.glyphs = (0..=self.trail)
            .map(|_| random_glyph(config.ascii, rng))
            .collect();
        self.message = if !config.message.is_empty() && rng.gen_ratio(1, 11) {
            Some(config.message.clone())
        } else {
            None
        };
    }

    fn tick(
        &mut self,
        stdout: &mut impl Write,
        config: &Config,
        height: u16,
        frame: u64,
        rng: &mut StdRng,
        force_step: bool,
    ) -> io::Result<()> {
        self.clock = self.clock.wrapping_add(1);

        let should_step = self.clock.checked_rem(self.step_every) == Some(0);
        if !force_step && !should_step {
            return Ok(());
        }

        self.y += 1;
        self.advance_glyphs(config, rng);
        self.clear_tail(stdout, height)?;
        self.draw_trail(stdout, config, height, frame, rng, force_step)?;

        if self.y - self.trail as i16 > height as i16 {
            self.respawn(height, config, rng);
        }

        Ok(())
    }

    fn clear_tail(&self, stdout: &mut impl Write, height: u16) -> io::Result<()> {
        let row = self.y - self.trail as i16 - 1;

        if row >= 0 && row < height as i16 {
            queue!(stdout, MoveTo(self.x, row as u16), Print(" "))?;
        }

        Ok(())
    }

    fn draw_trail(
        &mut self,
        stdout: &mut impl Write,
        config: &Config,
        height: u16,
        frame: u64,
        _rng: &mut StdRng,
        burst_hot: bool,
    ) -> io::Result<()> {
        for offset in (1..=self.trail).rev() {
            let row = self.y - offset as i16;

            if row < 0 || row >= height as i16 {
                continue;
            }

            let intensity = tail_intensity(offset, self.trail);
            let color =
                config
                    .palette
                    .color(intensity, self.x, row as u16, frame, self.hot || burst_hot);
            queue!(
                stdout,
                MoveTo(self.x, row as u16),
                SetForegroundColor(color),
                Print(self.glyphs.get(offset as usize).copied().unwrap_or('?'))
            )?;
        }

        if self.y >= 0 && self.y < height as i16 {
            let color = config
                .palette
                .color(255, self.x, self.y as u16, frame, true);
            queue!(
                stdout,
                MoveTo(self.x, self.y as u16),
                SetForegroundColor(color),
                Print(self.glyphs.first().copied().unwrap_or('?')),
                ResetColor
            )?;
        }

        Ok(())
    }

    fn advance_glyphs(&mut self, config: &Config, rng: &mut StdRng) {
        let expected_len = self.trail as usize + 1;
        if self.glyphs.len() != expected_len {
            self.glyphs = (0..expected_len)
                .map(|_| random_glyph(config.ascii, rng))
                .collect();
        }

        let head = self.next_head_glyph(config, rng);
        self.glyphs.rotate_right(1);
        self.glyphs[0] = head;

        if self.glyphs.len() > 3 && rng.gen_bool(0.16) {
            let index = rng.gen_range(1..self.glyphs.len());
            self.glyphs[index] = random_glyph(config.ascii, rng);
        }
    }

    fn next_head_glyph(&mut self, config: &Config, rng: &mut StdRng) -> char {
        if let Some(chars) = self.message.as_ref() {
            let ch = chars[self.message_index % chars.len()];
            self.message_index += 1;

            if self.message_index >= chars.len() {
                self.message = None;
                self.message_index = 0;
            }

            return ch;
        }

        random_glyph(config.ascii, rng)
    }
}

impl Palette {
    fn color(&self, intensity: u8, x: u16, y: u16, frame: u64, hot: bool) -> Color {
        if hot && intensity > 220 {
            return match self {
                Self::Amber => Color::Rgb {
                    r: 255,
                    g: 252,
                    b: 194,
                },
                Self::Crimson => Color::Rgb {
                    r: 255,
                    g: 220,
                    b: 220,
                },
                Self::Ghost => Color::White,
                _ => Color::Rgb {
                    r: 225,
                    g: 255,
                    b: 225,
                },
            };
        }

        match self {
            Self::Green => Color::Rgb {
                r: intensity / 5,
                g: intensity,
                b: intensity / 4,
            },
            Self::Cyan => Color::Rgb {
                r: intensity / 5,
                g: intensity,
                b: intensity,
            },
            Self::Amber => Color::Rgb {
                r: intensity,
                g: intensity.saturating_mul(3) / 5,
                b: intensity / 9,
            },
            Self::Crimson => Color::Rgb {
                r: intensity,
                g: intensity / 6,
                b: intensity / 4,
            },
            Self::Ghost => Color::Rgb {
                r: intensity,
                g: intensity,
                b: intensity,
            },
            Self::Rainbow => rainbow_color(intensity, x, y, frame),
        }
    }
}

fn build_columns(width: u16, height: u16, config: &Config, rng: &mut StdRng) -> Vec<Column> {
    let mut columns = Vec::with_capacity((width as f32 * config.density) as usize);

    for x in 0..width {
        if rng.gen_bool(config.density as f64) {
            columns.push(Column::random(x, height, config, rng));
        }
    }

    columns
}

fn random_glyph(ascii: bool, rng: &mut StdRng) -> char {
    const ASCII: &[u8] = b"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz@#$%&*+=?";

    if ascii {
        return ASCII[rng.gen_range(0..ASCII.len())] as char;
    }

    match rng.gen_range(0..100) {
        0..=56 => char::from_u32(rng.gen_range(0xFF66..=0xFF9D)).unwrap_or('?'),
        57..=76 => char::from_u32(rng.gen_range(0x30..=0x39)).unwrap_or('?'),
        77..=92 => ASCII[rng.gen_range(0..ASCII.len())] as char,
        _ => char::from_u32(rng.gen_range(0x2500..=0x257F)).unwrap_or('?'),
    }
}

fn tail_intensity(offset: u16, trail: u16) -> u8 {
    let trail = trail.max(1);
    let distance = offset.saturating_sub(1) as f32 / trail as f32;
    let fade = 1.0 - distance;
    (28.0 + fade.powf(1.65) * 178.0) as u8
}

fn rainbow_color(intensity: u8, x: u16, y: u16, frame: u64) -> Color {
    let wheel = ((frame as u16 / 2)
        .wrapping_add(x.wrapping_mul(7))
        .wrapping_add(y.wrapping_mul(3)))
        % 1536;
    let sector = wheel / 256;
    let offset = (wheel % 256) as u8;
    let down = 255_u8.saturating_sub(offset);
    let (r, g, b) = match sector {
        0 => (255, offset, 0),
        1 => (down, 255, 0),
        2 => (0, 255, offset),
        3 => (0, down, 255),
        4 => (offset, 0, 255),
        _ => (255, 0, down),
    };

    Color::Rgb {
        r: scale_channel(r, intensity),
        g: scale_channel(g, intensity),
        b: scale_channel(b, intensity),
    }
}

fn scale_channel(channel: u8, intensity: u8) -> u8 {
    ((channel as u16 * intensity as u16) / 255) as u8
}

fn parse_density(value: &str) -> Result<f32, String> {
    let density = value
        .parse::<f32>()
        .map_err(|_| "density must be a number between 0.05 and 1.0".to_string())?;

    if (0.05..=1.0).contains(&density) {
        Ok(density)
    } else {
        Err("density must be between 0.05 and 1.0".to_string())
    }
}

fn handle_event(
    matrix: &mut Matrix,
    stdout: &mut impl Write,
    paused: &mut bool,
) -> io::Result<bool> {
    match event::read()? {
        Event::Key(key) if key.kind == KeyEventKind::Press => match key.code {
            KeyCode::Esc | KeyCode::Char('q') => Ok(true),
            KeyCode::Char('c') if key.modifiers.contains(KeyModifiers::CONTROL) => Ok(true),
            KeyCode::Char(' ') | KeyCode::Char('p') => {
                *paused = !*paused;
                Ok(false)
            }
            _ => Ok(false),
        },
        Event::Resize(width, height) => {
            matrix.resize(width, height, stdout)?;
            Ok(false)
        }
        _ => Ok(false),
    }
}

fn run() -> io::Result<()> {
    let config = Config::from(Cli::parse());
    let frame_delay = Duration::from_nanos(1_000_000_000 / config.fps);
    let stdout = stdout();
    let mut stdout = io::BufWriter::with_capacity(256 * 1024, stdout);
    let (width, height) = terminal::size()?;
    let alt_screen = config.alt_screen;
    let mut matrix = Matrix::new(config, width, height);
    let _terminal = TerminalSession::enter(&mut stdout, alt_screen)?;
    let mut paused = false;
    let mut next_frame = Instant::now();

    loop {
        while event::poll(Duration::from_millis(0))? {
            if handle_event(&mut matrix, &mut stdout, &mut paused)? {
                return Ok(());
            }
        }

        matrix.draw_frame(&mut stdout, paused)?;
        stdout.flush()?;

        next_frame += frame_delay;
        let now = Instant::now();
        if let Some(remaining) = next_frame.checked_duration_since(now) {
            thread::sleep(remaining);
        } else {
            next_frame = now;
        }
    }
}

fn main() -> io::Result<()> {
    run()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn density_parser_rejects_out_of_range_values() {
        assert!(parse_density("0.5").is_ok());
        assert!(parse_density("0.01").is_err());
        assert!(parse_density("1.2").is_err());
    }

    #[test]
    fn tail_gets_brighter_near_the_head() {
        assert!(tail_intensity(1, 20) > tail_intensity(20, 20));
    }

    #[test]
    fn seeded_ascii_glyphs_are_ascii() {
        let mut rng = StdRng::seed_from_u64(7);

        for _ in 0..128 {
            assert!(random_glyph(true, &mut rng).is_ascii());
        }
    }
}