splink_client/src/main.rs

#![warn(clippy::pedantic)]
#![feature(never_type)]

use std::{
    f32::consts::FRAC_PI_2,
    io::{stdin, stdout, Read},
    net::{Ipv4Addr, SocketAddr, UdpSocket},
    path::PathBuf,
    thread::sleep,
    time::{Duration, Instant},
};

use bracket_color::prelude::{HSV, RGB};
use clap::{builder::TypedValueParser, Parser, Subcommand, ValueEnum};
use image::{imageops::FilterType, io::Reader as ImageReader, Pixel, Rgb, RgbImage};
use rand::Rng;

use splink_client::{send_frame, Layout, SenderError};

/// Blinkenwall v3 prototype client
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
    /// Timeout in ms for frame completion responses (0 to disable)
    #[clap(long, default_value_t = 10)]
    response_timeout: u64,

    /// Local address and port to bind to
    #[clap(long = "bind")]
    bind_addr: Option<SocketAddr>,

    /// Controller's address
    remote_addr: SocketAddr,

    /// The action to perform
    #[clap(subcommand)]
    action: Action,
}

#[derive(Clone, Copy)]
struct ColorParser;
impl TypedValueParser for ColorParser {
    type Value = Rgb<u8>;

    fn parse_ref(
        &self,
        _cmd: &clap::Command,
        _arg: Option<&clap::Arg>,
        value: &std::ffi::OsStr,
    ) -> Result<Self::Value, clap::Error> {
        let s = value
            .to_str()
            .ok_or(clap::Error::raw(clap::ErrorKind::InvalidUtf8, ""))?;
        let s = s.strip_prefix('#').unwrap_or(s);

        if s.len() != 6 {
            return Err(clap::Error::raw(
                clap::ErrorKind::InvalidValue,
                "Must be a 6-digit hexadecimal number",
            ));
        }

        let mut channels = [0; 3];
        for i in 0..3 {
            channels[i] = u8::from_str_radix(&s[i * 2..i * 2 + 2], 16).map_err(|_| {
                clap::Error::raw(clap::ErrorKind::InvalidValue, "Invalid hex literal")
            })?;
        }

        Ok(Rgb(channels))
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Subcommand)]
enum Action {
    Animation {
        #[clap(value_enum)]
        animation: Animation,
    },
    Solid {
        #[clap(value_parser = ColorParser)]
        color: Rgb<u8>,
    },
    Image {
        path: PathBuf,
    },
    Video,
    Clear,
}

#[derive(Clone, Debug, PartialEq, Eq, ValueEnum)]
enum Animation {
    Rainbow,
    Bling,
    Strobe,
}

fn bling(layout: Layout, frame: u32) -> RgbImage {
    #![allow(
        clippy::cast_precision_loss,
        clippy::cast_lossless,
        clippy::cast_possible_truncation,
        clippy::cast_sign_loss
    )]

    let w = layout.width_px();
    let h = layout.height_px();

    let frame = frame % 60;
    let frame = frame as f32 / 3.0;

    RgbImage::from_fn(w, h, |x, y| {
        let dist_from_center = {
            let x = x as f32;
            let y = y as f32;

            let center_x = w as f32 / 2.0;
            let center_y = h as f32 / 2.0;

            ((x - center_x).powf(2.0) + (y - center_y).powf(2.0)).sqrt()
        };

        let radius = frame.powf(1.8);

        let dist_from_radius = dist_from_center - radius;
        let dist = 0.3 * dist_from_radius;

        Rgb([100, 0, 100]).map(|x| {
            (if dist < 0.0 {
                (1.0 - (dist.abs() / 10.0)).max(0.0)
            } else {
                dist.min(FRAC_PI_2).cos()
            } * x as f32) as u8
        })
    })
}

fn rainbow(layout: Layout, frame: u32) -> RgbImage {
    #![allow(
        clippy::cast_precision_loss,
        clippy::cast_lossless,
        clippy::cast_possible_truncation,
        clippy::cast_sign_loss
    )]

    let w = layout.width_px();
    let h = layout.height_px();

    let tightness = 100.0;
    let slowness = 200;
    let brightness = 0.1;

    RgbImage::from_fn(w, h, |x, y| {
        let RGB { r, g, b } = HSV::from_f32(
            (x + y) as f32 / tightness % 1.0 + (frame % slowness) as f32 / slowness as f32,
            1.0,
            brightness,
        )
        .to_rgb();
        Rgb([(r * 255.0) as u8, (g * 255.0) as u8, (b * 255.0) as u8])
    })
}

fn strobe(layout: Layout, frame: u32) -> RgbImage {
    let w = layout.width_px();
    let h = layout.height_px();

    let brightness = 120;
    let period = 15;

    RgbImage::from_fn(w, h, |_, _| {
        if frame % period < period / 2 {
            Rgb([brightness; 3])
        } else {
            Rgb([0, 0, 0])
        }
    })
}

fn print_image(image: &RgbImage) {
    let _hide = termion::cursor::HideCursor::from(stdout());

    print!("{}", termion::cursor::Goto(1, 1));
    for row in image.rows() {
        for &Rgb([r, g, b]) in row {
            print!("{}  ", termion::color::Rgb(r, g, b).bg_string());
        }
        println!("{}", termion::color::Reset.bg_str());
    }
}

fn animate<F: Fn(Layout, u32) -> RgbImage>(
    socket: &UdpSocket,
    layout: Layout,
    generator: F,
) -> Result<!, anyhow::Error> {
    print!("{}", termion::clear::All);

    let mut frame: u32 = rand::thread_rng().gen();
    loop {
        let start = Instant::now();
        let image = generator(layout, frame);

        print_image(&image);

        match send_frame(socket, layout, frame, &image) {
            Ok(()) | Err(SenderError::ConfirmationTimeout) => {}
            Err(e) => return Err(e.into()),
        };

        println!("{:?}", Instant::now().duration_since(start));
        frame += 1;
    }
}

fn main() -> anyhow::Result<()> {
    let args = Args::parse();

    let bind_addr = args
        .bind_addr
        .unwrap_or_else(|| (Ipv4Addr::UNSPECIFIED, 0).into());
    let socket = UdpSocket::bind(bind_addr)?;

    socket.connect(args.remote_addr)?;
    socket.set_read_timeout(match args.response_timeout {
        0 => None,
        t => Some(Duration::from_millis(t)),
    })?;

    let layout = Layout {
        gang_len: 8,
        num_gangs: 32,

        num_panels_h: 6,
        num_panels_v: 1,

        total_strands: 24,
        first_strand_index: 8,
    };

    let image = match args.action {
        Action::Solid { color } => {
            RgbImage::from_pixel(layout.width_px(), layout.height_px(), color)
        }
        Action::Clear => RgbImage::new(layout.width_px(), layout.height_px()),
        Action::Image { path } => ImageReader::open(path)?
            .decode()?
            .resize(layout.width_px(), layout.height_px(), FilterType::Gaussian)
            .into_rgb8(),
        Action::Animation { animation } => {
            animate(
                &socket,
                layout,
                match animation {
                    Animation::Rainbow => rainbow,
                    Animation::Bling => bling,
                    Animation::Strobe => strobe,
                },
            )?;
        }
        Action::Video => {
            print!("{}", termion::clear::All);

            let mut frame_num: u32 = rand::thread_rng().gen();
            loop {
                let w = layout.width_px();
                let h = layout.height_px();

                let mut buf = vec![0; w as usize * h as usize * 3];
                stdin().lock().read_exact(&mut buf)?;
                for c in &mut buf {
                    *c /= 6;
                }
                let image = RgbImage::from_vec(w, h, buf).unwrap();

                print_image(&image);

                match send_frame(&socket, layout, frame_num, &image) {
                    Ok(()) | Err(SenderError::ConfirmationTimeout) => {}
                    Err(e) => return Err(e.into()),
                };
                frame_num += 1;

                sleep(Duration::from_millis(30));
            }
        }
    };

    let frame_num: u32 = rand::thread_rng().gen();
    send_frame(&socket, layout, frame_num, &image)?;

    Ok(())
}