Fearless Embedded Rust: Driving a Lego Car with a Pico W

#![no_std]
#![no_main]

mod slice_writer;

use {
    core::fmt::Write,
    embassy_executor::Spawner,
    embassy_futures::join::{
        join,
        join3,
        join_array,
    },
    embassy_rp::{
        bind_interrupts,
        gpio::{
            Level,
            Output,
        },
        peripherals::USB,
        usb::{
            Driver,
            InterruptHandler,
        },
    },
    embassy_time::Timer,
    embassy_usb::{
        class::cdc_acm::{
            CdcAcmClass,
            State,
        },
        Builder,
    },
    slice_writer::SliceWriter,
    panic_halt as _,
};

// Embedded devices react to hardware events through interrupts,
// lightweight callbacks triggered by peripherals such as USB controllers
// and timers. This macro statically binds the USB hardware interrupts
// to the Embassy handler
bind_interrupts!(struct Irqs {
    USBCTRL_IRQ => InterruptHandler<USB>;
});

#[embassy_executor::main]
async fn main(_spawner: Spawner) {
    // Initialize the RP2040 peripherals. The embassy_rp::init() function
    // configures the clocks, USB bus, and GPIOs.
    let peripherals = embassy_rp::init(Default::default());

    // Initializes the USB driver by passing the Irqs interrupt structure
    let driver = Driver::new(peripherals.USB, Irqs);

    // Configuration of pin GP15 (physical pin 20) as output
    let mut led = Output::new(peripherals.PIN_15, Level::Low);

    // Configuration of the USB device, with the CDC ACM class
    // for serial communication
    let mut config = embassy_usb::Config::new(0xc0de, 0xcafe);
    config.manufacturer = Some("dystroy");
    config.product = Some("Pico Mobile - log");
    config.serial_number = Some("12345678");
    config.max_power = 100;
    config.max_packet_size_0 = 64;
    let mut config_descriptor = [0; 256];
    let mut bos_descriptor = [0; 256];
    let mut control_buf = [0; 64];
    let mut state = State::new();
    let mut builder = Builder::new(
        driver,
        config,
        &mut config_descriptor,
        &mut bos_descriptor,
        &mut [],
        &mut control_buf,
    );
    let mut class = CdcAcmClass::new(&mut builder, &mut state, 64);
    let mut usb = builder.build();

    // Execution of all async tasks: USB, log, and LED blinking
    let usb_fut = usb.run();

    let log_fut = async {
        let mut buffer = [0u8; 100];
        let mut writer = SliceWriter::new(&mut buffer);
        let mut n = 0;
        loop {
            // Wait for the serial port to be opened on the computer side
            // (e.g., via tio)
            class.wait_connection().await;
            n += 1;
            writer.reset();
            if write!(writer, "Message n° {n}\r\n").is_ok() {
                let _ = class
                    .write_packet(writer.as_bytes())
                    .await;
            }
            Timer::after_secs(1).await;
        }
    };

    let blink_fut = async {
        loop {
            led.set_high();
            Timer::after_millis(5000).await;
            led.set_low();
            Timer::after_millis(5000).await;
        }
    };

    join3(usb_fut, log_fut, blink_fut).await;
}

use core::fmt::Write;

/// Wraps a mutable byte slice and implements the `Write` trait for it,
/// allowing formatted strings to be written into the slice.
pub struct SliceWriter<'a> {
    buf: &'a mut [u8],
    cursor: usize,
}

impl SliceWriter<'_> {
    pub fn new(buf: &mut [u8]) -> SliceWriter<'_> {
        SliceWriter { buf, cursor: 0 }
    }
    pub fn reset(&mut self) {
        self.cursor = 0;
    }
    pub fn as_bytes(&self) -> &[u8] {
        &self.buf[..self.cursor]
    }
}

impl<'a> Write for SliceWriter<'a> {
    fn write_str(&mut self, s: &str) -> core::fmt::Result {
        let bytes = s.as_bytes();
        if self.cursor + bytes.len() > self.buf.len() {
            return Err(core::fmt::Error);
        }
        self.buf[self.cursor..self.cursor + bytes.len()]
            .copy_from_slice(bytes);
        self.cursor += bytes.len();
        Ok(())
    }
}

use core::fmt::Write;

/// Wraps a mutable byte slice and implements the `Write` trait for it,
/// allowing formatted strings to be written into the slice.
pub struct SliceWriter<'a> {
    buf: &'a mut [u8],
    cursor: usize,
}

impl SliceWriter<'_> {
    pub fn new(buf: &mut [u8]) -> SliceWriter<'_> {
        SliceWriter { buf, cursor: 0 }
    }
    pub fn reset(&mut self) {
        self.cursor = 0;
    }
    pub fn as_bytes(&self) -> &[u8] {
        &self.buf[..self.cursor]
    }
}

impl<'a> Write for SliceWriter<'a> {
    fn write_str(&mut self, s: &str) -> core::fmt::Result {
        let bytes = s.as_bytes();
        if self.cursor + bytes.len() > self.buf.len() {
            return Err(core::fmt::Error);
        }
        self.buf[self.cursor..self.cursor + bytes.len()]
            .copy_from_slice(bytes);
        self.cursor += bytes.len();
        Ok(())
    }
}

use core::str::FromStr;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Command {
    Blink(u8), // number of times to blink
    Bye,
}

impl FromStr for Command {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut parts = s.split_whitespace();
        let command = parts.next().ok_or("empty command")?;
        match command {
            "blink" => {
                let count_str = parts.next()
                    .ok_or("Blink command missing a count")?;
                let count = count_str
                    .parse::<u8>()
                    .map_err(|_| "count not parsable as u8")?;
                Ok(Command::Blink(count))
            }
            "bye" => Ok(Command::Bye),
            _ => Err("unknown command"),
        }
    }
}

#![no_std]
#![no_main]

mod command;
mod slice_writer;
use {
    command::*,
    slice_writer::*,
};

use {
    core::{
        fmt::Write as _,
        str::from_utf8,
    },
    cyw43::{
        aligned_bytes,
        JoinOptions,
    },
    cyw43_pio::{
        PioSpi,
        DEFAULT_CLOCK_DIVIDER,
    },
    embassy_executor::Spawner,
    embassy_net::{
        tcp::TcpSocket,
        Config,
        StackResources,
    },
    embassy_rp::{
        bind_interrupts,
        clocks::RoscRng,
        dma,
        gpio::{
            Level,
            Output,
        },
        peripherals::{
            DMA_CH0,
            DMA_CH1,
            PIO0,
            USB,
        },
        pio::{
            InterruptHandler,
            Pio,
        },
        usb::{
            Driver as UsbDriver,
            InterruptHandler as UsbIrqHandler,
        },
    },
    embassy_time::{
        Duration,
        Timer,
    },
    embedded_io_async::Write,
    log::{
        info,
        warn,
    },
    panic_halt as _,
    static_cell::StaticCell,
};

bind_interrupts!(struct Irqs {
    PIO0_IRQ_0 => InterruptHandler<PIO0>;
    DMA_IRQ_0 => dma::InterruptHandler<DMA_CH0>, dma::InterruptHandler<DMA_CH1>;
    USBCTRL_IRQ => UsbIrqHandler<USB>;
});

const WIFI_NETWORK: &str = env!("WIFI_SSID");
const WIFI_PASSWORD: &str = env!("WIFI_PASSWORD");

#[embassy_executor::task]
async fn cyw43_task(
    runner: cyw43::Runner<
        'static,
        cyw43::SpiBus<Output<'static>, PioSpi<'static, PIO0, 0>>,
        cyw43::Cyw43439,
    >,
) -> ! {
    runner.run().await
}

#[embassy_executor::task]
async fn net_task(mut runner: embassy_net::Runner<'static, cyw43::NetDriver<'static>>) -> ! {
    runner.run().await
}

#[embassy_executor::task]
async fn logger_task(driver: UsbDriver<'static, USB>) {
    embassy_usb_logger::run!(1024, log::LevelFilter::Info, driver);
}

#[embassy_executor::task]
async fn tick() {
    let mut i = 0;
    loop {
        Timer::after_millis(10_000).await;
        i += 1;
        log::info!("tick {i}");
    }
}

/// Helper function to handle fatal errors in async tasks.
/// Logs the error, waits for a short duration so that the message can
/// arrive to the computer, and then panics.
async fn expect<T, E: core::fmt::Display>(res: Result<T, E>) -> T {
    match res {
        Ok(v) => v,
        Err(e) => {
            log::error!("Fatal error: {}", e);
            Timer::after_millis(100).await;
            panic!("Fatal error: {}", e);
        }
    }
}

#[embassy_executor::main]
async fn main(spawner: Spawner) {
    let p = embassy_rp::init(Default::default());

    // Logger USB — must be before everything else
    let usb_driver = UsbDriver::new(p.USB, Irqs);
    spawner.spawn(expect(logger_task(usb_driver)).await);

    info!("Starting...");

    // LED on pin 15, active=high
    let mut led = Output::new(p.PIN_15, Level::Low);

    spawner.spawn(expect(tick()).await);

    let mut rng = RoscRng;

    let fw = aligned_bytes!("../../cyw43-firmware/43439A0.bin");
    let clm = aligned_bytes!("../../cyw43-firmware/43439A0_clm.bin");
    let nvram = aligned_bytes!("../../cyw43-firmware/nvram_rp2040.bin");

    let pwr = Output::new(p.PIN_23, Level::Low);
    let cs = Output::new(p.PIN_25, Level::High);
    let mut pio = Pio::new(p.PIO0, Irqs);
    let spi = PioSpi::new(
        &mut pio.common,
        pio.sm0,
        DEFAULT_CLOCK_DIVIDER,
        pio.irq0,
        cs,
        p.PIN_24,
        p.PIN_29,
        dma::Channel::new(p.DMA_CH0, Irqs),
        dma::Channel::new(p.DMA_CH1, Irqs),
    );

    static STATE: StaticCell<cyw43::State> = StaticCell::new();
    let state = STATE.init(cyw43::State::new());
    let (net_device, mut control, runner) = cyw43::new(state, pwr, spi, fw, nvram).await;
    spawner.spawn(expect(cyw43_task(runner)).await);

    control.init(clm).await;
    control
        .set_power_management(cyw43::PowerManagementMode::PowerSave)
        .await;

    let config = Config::dhcpv4(Default::default());

    // Generate random seed
    let seed = rng.next_u64();

    // Init network stack
    static RESOURCES: StaticCell<StackResources<3>> = StaticCell::new();
    let (stack, runner) = embassy_net::new(
        net_device,
        config,
        RESOURCES.init(StackResources::new()),
        seed,
    );

    spawner.spawn(expect(net_task(runner)).await);

    while let Err(err) = control
        .join(WIFI_NETWORK, JoinOptions::new(WIFI_PASSWORD.as_bytes()))
        .await
    {
        info!("join failed: {:?}", err);
    }

    info!("waiting for link...");
    stack.wait_link_up().await;

    info!("waiting for DHCP...");
    stack.wait_config_up().await;

    // And now we can use it!
    info!("Stack is up!");

    let mut rx_buffer = [0; 4096];
    let mut tx_buffer = [0; 4096];
    let mut buf = [0; 4096];

    loop {
        let mut socket = TcpSocket::new(stack, &mut rx_buffer, &mut tx_buffer);
        socket.set_timeout(Some(Duration::from_secs(10)));

        control.gpio_set(0, false).await;
        info!("Listening on TCP:1234...");
        led.set_high();
        Timer::after_millis(100).await;
        led.set_low();
        if let Err(e) = socket.accept(1234).await {
            warn!("accept error: {:?}", e);
            continue;
        }

        info!("Received connection from {:?}", socket.remote_endpoint());
        control.gpio_set(0, true).await;

        loop {
            let _ = socket.write_all(b"\n> ").await;
            let n = match socket.read(&mut buf).await {
                Ok(0) => {
                    warn!("read EOF");
                    break;
                }
                Ok(n) => n,
                Err(e) => {
                    warn!("read error: {:?}", e);
                    break;
                }
            };
            let Ok(input) = from_utf8(&buf[..n]) else {
                warn!("Received non-UTF8 data");
                continue;
            };

            info!(" <- Received '{}'", input);
            let cmd = input.parse::<Command>();

            let mut output = SliceWriter::new(&mut buf);
            let mut quit = false;
            match cmd {
                Ok(Command::Blink(n)) => {
                    for _ in 0..=n {
                        led.set_high();
                        Timer::after_millis(100).await;
                        led.set_low();
                        Timer::after_millis(100).await;
                    }
                    let _ = write!(output, "-> Blinked {} times", n);
                }
                Ok(Command::Bye) => {
                    let _ = write!(output, "-> Goodbye!");
                    quit = true;
                }
                Err(e) => {
                    let _ = write!(output, "-> ERROR: {}", e);
                }
            }

            match socket.write_all(output.as_bytes()).await {
                Ok(()) => {
                    info!("Answered '{}'", output.as_str());
                }
                Err(e) => {
                    warn!("write error: {:?}", e);
                    break;
                }
            };
            if quit {
                info!("Closing connection");
                break;
            }
        }
        Timer::after_millis(10).await;
    }
}

#![no_std]
#![no_main]

mod command;
mod slice_writer;
use {
    command::*,
    slice_writer::*,
};

use {
    core::{
        str::from_utf8,
        fmt::Write as _,
    },
    cyw43::{
        aligned_bytes,
        JoinOptions,
    },
    cyw43_pio::{
        PioSpi,
        DEFAULT_CLOCK_DIVIDER,
    },
    embassy_executor::Spawner,
    embassy_net::{
        tcp::TcpSocket,
        Config,
        StackResources,
    },
    embassy_rp::{
        bind_interrupts,
        clocks::RoscRng,
        dma,
        gpio::{
            Level,
            Output,
        },
        peripherals::{
            DMA_CH0,
            DMA_CH1,
            PIO0,
            USB,
        },
        pio::{
            InterruptHandler,
            Pio,
        },
        usb::{
            Driver as UsbDriver,
            InterruptHandler as UsbIrqHandler,
        },
    },
    embassy_time::{
        Duration,
        Timer,
    },
    embedded_io_async::Write,
    log::{
        info,
        warn,
    },
    panic_halt as _,
    rp2040_hal::rom_data::reset_to_usb_boot,
    static_cell::StaticCell,
};

bind_interrupts!(struct Irqs {
    PIO0_IRQ_0 => InterruptHandler<PIO0>;
    DMA_IRQ_0 => dma::InterruptHandler<DMA_CH0>, dma::InterruptHandler<DMA_CH1>;
    USBCTRL_IRQ => UsbIrqHandler<USB>;
});

const WIFI_NETWORK: &str = env!("WIFI_SSID");
const WIFI_PASSWORD: &str = env!("WIFI_PASSWORD");

#[embassy_executor::task]
async fn cyw43_task(
    runner: cyw43::Runner<
        'static,
        cyw43::SpiBus<Output<'static>, PioSpi<'static, PIO0, 0>>,
        cyw43::Cyw43439,
    >,
) -> ! {
    runner.run().await
}

#[embassy_executor::task]
async fn net_task(mut runner: embassy_net::Runner<'static, cyw43::NetDriver<'static>>) -> ! {
    runner.run().await
}

#[embassy_executor::task]
async fn logger_task(driver: UsbDriver<'static, USB>) {
    embassy_usb_logger::run!(1024, log::LevelFilter::Info, driver);
}

#[embassy_executor::task]
async fn tick() {
    let mut i = 0;
    loop {
        Timer::after_millis(10_000).await;
        i += 1;
        log::info!("tick {i}");
    }
}

async fn expect<T, E: core::fmt::Display>(res: Result<T, E>) -> T {
    match res {
        Ok(v) => v,
        Err(e) => {
            log::error!("Fatal error: {}", e);
            Timer::after_millis(100).await;
            panic!("Fatal error: {}", e);
        }
    }
}

#[embassy_executor::main]
async fn main(spawner: Spawner) {
    let p = embassy_rp::init(Default::default());

    // Logger USB — must be before everything else
    let usb_driver = UsbDriver::new(p.USB, Irqs);
    spawner.spawn(expect(logger_task(usb_driver)).await);

    info!("Starting...");

    // LED on pin GP27 of the Kitronik 5331, active=high
    let mut led = Output::new(p.PIN_27, Level::Low);

    spawner.spawn(expect(tick()).await);

    let mut rng = RoscRng;

    let fw = aligned_bytes!("../../cyw43-firmware/43439A0.bin");
    let clm = aligned_bytes!("../../cyw43-firmware/43439A0_clm.bin");
    let nvram = aligned_bytes!("../../cyw43-firmware/nvram_rp2040.bin");

    let pwr = Output::new(p.PIN_23, Level::Low);
    let cs = Output::new(p.PIN_25, Level::High);
    let mut pio = Pio::new(p.PIO0, Irqs);
    let spi = PioSpi::new(
        &mut pio.common,
        pio.sm0,
        DEFAULT_CLOCK_DIVIDER,
        pio.irq0,
        cs,
        p.PIN_24,
        p.PIN_29,
        dma::Channel::new(p.DMA_CH0, Irqs),
        dma::Channel::new(p.DMA_CH1, Irqs),
    );

    static STATE: StaticCell<cyw43::State> = StaticCell::new();
    let state = STATE.init(cyw43::State::new());
    let (net_device, mut control, runner) = cyw43::new(state, pwr, spi, fw, nvram).await;
    spawner.spawn(expect(cyw43_task(runner)).await);

    control.init(clm).await;
    control
        .set_power_management(cyw43::PowerManagementMode::PowerSave)
        .await;

    let config = Config::dhcpv4(Default::default());

    // Generate random seed
    let seed = rng.next_u64();

    // Init network stack
    static RESOURCES: StaticCell<StackResources<3>> = StaticCell::new();
    let (stack, runner) = embassy_net::new(
        net_device,
        config,
        RESOURCES.init(StackResources::new()),
        seed,
    );

    spawner.spawn(expect(net_task(runner)).await);

    while let Err(err) = control
        .join(WIFI_NETWORK, JoinOptions::new(WIFI_PASSWORD.as_bytes()))
        .await
    {
        info!("join failed: {:?}", err);
    }

    info!("waiting for link...");
    stack.wait_link_up().await;

    info!("waiting for DHCP...");
    stack.wait_config_up().await;

    // And now we can use it!
    info!("Stack is up!");

    let mut rx_buffer = [0; 4096];
    let mut tx_buffer = [0; 4096];
    let mut buf = [0; 4096];

    loop {
        let mut socket = TcpSocket::new(stack, &mut rx_buffer, &mut tx_buffer);
        socket.set_timeout(Some(Duration::from_secs(10)));

        control.gpio_set(0, false).await;
        info!("Listening on TCP:1234...");
        led.set_high();
        Timer::after_millis(100).await;
        led.set_low();
        if let Err(e) = socket.accept(1234).await {
            warn!("accept error: {:?}", e);
            continue;
        }

        info!("Received connection from {:?}", socket.remote_endpoint());
        control.gpio_set(0, true).await;

        loop {
            let _ = socket.write_all(b"\n> ").await;
            let n = match socket.read(&mut buf).await {
                Ok(0) => {
                    warn!("read EOF");
                    break;
                }
                Ok(n) => n,
                Err(e) => {
                    warn!("read error: {:?}", e);
                    break;
                }
            };
            let Ok(input) = from_utf8(&buf[..n]) else {
                warn!("Received non-UTF8 data");
                continue;
            };

            info!(" <- Received '{}'", input);
            let cmd = input.parse::<Command>();

            let mut output = SliceWriter::new(&mut buf);
            let mut close = false;
            match cmd {
                Ok(Command::Blink(n)) => {
                    for _ in 0..=n {
                        led.set_high();
                        Timer::after_millis(100).await;
                        led.set_low();
                        Timer::after_millis(100).await;
                    }
                    let _ = write!(output, "-> Blinked {} times", n);
                }
                Ok(Command::Bye) => {
                    let _ = write!(output, "-> Goodbye!");
                    close = true;
                }
                Ok(Command::Quit) => {
                    let _ = write!(output, "-> Shutting down...");
                    Timer::after_millis(100).await;
                    reset_to_usb_boot(0, 0);
                }
                Err(e) => {
                    let _ = write!(output, "-> ERROR: {}", e);
                }
            }

            match socket.write_all(output.as_bytes()).await {
                Ok(()) => {
                    info!("Answered '{}'", output.as_str());
                }
                Err(e) => {
                    warn!("write error: {:?}", e);
                    break;
                }
            };
            if close {
                info!("Closing connection");
                break;
            }
        }
        Timer::after_millis(10).await;
    }
}

use core::str::FromStr;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Command {
    /// Blink a LED a number of times
    Blink(u8),
    /// Disconnect the client, but keep the server running
    Bye,
    /// Disconnect the client and shut down the Pico (back to the bootloader)
    Quit,
}

impl FromStr for Command {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut parts = s.split_whitespace();
        let command = parts.next().ok_or("empty command")?;
        match command {
            "blink" | "b" => {
                let count_str = parts.next()
                    .ok_or("Blink command missing a count")?;
                let count = count_str
                    .parse::<u8>()
                    .map_err(|_| "count not parsable as u8")?;
                Ok(Command::Blink(count))
            }
            "bye" => Ok(Command::Bye),
            "q" | "quit" => Ok(Command::Quit),
            _ => Err("unknown command"),
        }
    }
}

use core::fmt::Write;

/// Wraps a mutable byte slice and implements the `Write` trait for it,
/// allowing formatted strings to be written into the slice.
pub struct SliceWriter<'a> {
    buf: &'a mut [u8],
    cursor: usize,
}

impl SliceWriter<'_> {
    pub fn new(buf: &mut [u8]) -> SliceWriter<'_> {
        SliceWriter { buf, cursor: 0 }
    }
    pub fn reset(&mut self) {
        self.cursor = 0;
    }
    pub fn as_bytes(&self) -> &[u8] {
        &self.buf[..self.cursor]
    }
}

impl<'a> Write for SliceWriter<'a> {
    fn write_str(&mut self, s: &str) -> core::fmt::Result {
        let bytes = s.as_bytes();
        if self.cursor + bytes.len() > self.buf.len() {
            return Err(core::fmt::Error);
        }
        self.buf[self.cursor..self.cursor + bytes.len()]
            .copy_from_slice(bytes);
        self.cursor += bytes.len();
        Ok(())
    }
}

use {
    crate::SteeringDirection,
    core::str::FromStr,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Command {
    /// Blink a LED a number of times
    Blink(u8),
    /// Disconnect the client, but keep the server running
    Bye,
    /// Disconnect the client and shut down the Pico
    /// (back to the bootloader)
    Quit,
    /// Go forward (positive) or backward (negative)
    /// for a number of milliseconds
    GoMillis(i32),
    /// Turn left, right, or center the steering
    /// with a certain level of intensity (0-7)
    Steer {
        direction: SteeringDirection,
        level: u8, // in [1-7]
    },
}

impl FromStr for Command {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut parts = s.split_whitespace();
        let command = parts.next().ok_or("empty command")?;
        match command {
            "blink" | "b" => {
                let count_str = parts.next()
                    .ok_or("Blink command missing a count")?;
                let count = count_str
                    .parse::<u8>()
                    .map_err(|_| "count not parsable as u8")?;
                Ok(Command::Blink(count))
            }
            "go" | "g" => {
                let count_str = parts.next()
                    .ok_or("Blink command missing a count")?;
                let count = count_str
                    .parse::<i32>()
                    .map_err(|_| "count not parsable as i32")?;
                Ok(Command::GoMillis(count))
            }
            "l" | "left" => {
                let level_str = parts.next()
                    .ok_or("Steer left command missing a level")?;
                let level = level_str
                    .parse::<u8>()
                    .map_err(|_| "level not parsable as u8")?;
                if level > 7 {
                    return Err("level must be between 1 and 7");
                }
                Ok(Command::Steer {
                    direction: SteeringDirection::Left,
                    level,
                })
            }
            "r" |"right" => {
                let level_str = parts.next()
                    .ok_or("Steer right command missing a level")?;
                let level = level_str
                    .parse::<u8>()
                    .map_err(|_| "level not parsable as u8")?;
                if level > 7 {
                    return Err("level must be between 1 and 7");
                }
                Ok(Command::Steer {
                    direction: SteeringDirection::Right,
                    level,
                })
            }
            "c" | "center" => Ok(Command::Steer {
                direction: SteeringDirection::Center,
                level: 0,
            }),
            "bye" => Ok(Command::Bye),
            "q" | "quit" => Ok(Command::Quit),
            _ => Err("unknown command"),
        }
    }
}

#![no_std]
#![no_main]

mod command;
mod motor;
mod servo;
mod slice_writer;
mod utils;

pub use {
    command::*,
    motor::*,
    servo::*,
    slice_writer::*,
    utils::*,
};

use {
    core::{
        fmt::Write as _,
        str::from_utf8,
    },
    cyw43::{
        aligned_bytes,
        JoinOptions,
    },
    cyw43_pio::{
        PioSpi,
        DEFAULT_CLOCK_DIVIDER,
    },
    embassy_executor::Spawner,
    embassy_net::{
        tcp::TcpSocket,
        Config,
        StackResources,
    },
    embassy_rp::{
        bind_interrupts,
        clocks::RoscRng,
        dma,
        gpio::{
            Level,
            Output,
        },
        peripherals::{
            DMA_CH0,
            DMA_CH1,
            PIO0,
            USB,
        },
        pio::{
            InterruptHandler,
            Pio,
        },
        pwm::{
            Pwm,
        },
        usb::{
            Driver as UsbDriver,
            InterruptHandler as UsbIrqHandler,
        },
    },
    embassy_time::{
        Duration,
        Timer,
    },
    embedded_io_async::Write,
    log::{
        info,
        warn,
    },
    panic_halt as _,
    rp2040_hal::rom_data::reset_to_usb_boot,
    static_cell::StaticCell,
};

bind_interrupts!(struct Irqs {
    PIO0_IRQ_0 => InterruptHandler<PIO0>;
    DMA_IRQ_0 => dma::InterruptHandler<DMA_CH0>, dma::InterruptHandler<DMA_CH1>;
    USBCTRL_IRQ => UsbIrqHandler<USB>;
});

const WIFI_NETWORK: &str = env!("WIFI_SSID");
const WIFI_PASSWORD: &str = env!("WIFI_PASSWORD");

#[embassy_executor::task]
async fn cyw43_task(
    runner: cyw43::Runner<
        'static,
        cyw43::SpiBus<Output<'static>, PioSpi<'static, PIO0, 0>>,
        cyw43::Cyw43439,
    >,
) -> ! {
    runner.run().await
}

#[embassy_executor::task]
async fn net_task(mut runner: embassy_net::Runner<'static, cyw43::NetDriver<'static>>) -> ! {
    runner.run().await
}

#[embassy_executor::task]
async fn logger_task(driver: UsbDriver<'static, USB>) {
    embassy_usb_logger::run!(1024, log::LevelFilter::Info, driver);
}

#[embassy_executor::task]
async fn tick() {
    let mut i = 0;
    loop {
        Timer::after_millis(10_000).await;
        i += 1;
        log::info!("tick {i}");
    }
}

#[embassy_executor::main]
async fn main(spawner: Spawner) {
    let mut p = embassy_rp::init(Default::default());

    // Logger USB — must be before everything else
    let usb_driver = UsbDriver::new(p.USB, Irqs);
    spawner.spawn(expect(logger_task(usb_driver)).await);

    info!("Starting...");

    // LED on pin 27 of the Kitronik 5331, active=high
    let mut led = Output::new(p.PIN_27, Level::Low);

    // Motor on pins Motor1 of the Kitronik 5331, which are mapped
    // to pins 2 and 3 of the RP2040
    let mut motor = Motor::new(
        Output::new(p.PIN_2, Level::Low),
        Output::new(p.PIN_3, Level::Low),
    );

    // Servo on pins Motor2 of the Kitronik 5331,
    // which are mapped to pins 6 and 7 of the RP2040,
    // pins which are handled by the PWM slice 3
    // (On the RP2040, not every GPIO can be attached to every
    //  PWM generator, so the pin choice is partly dictated by
    //  hardware constraints)
    let pwm = Pwm::new_output_ab(
        p.PWM_SLICE3.reborrow(),
        p.PIN_6.reborrow(),
        p.PIN_7.reborrow(),
        LegoServo::pwm_config(),
    );
    let (channel_a, channel_b) = pwm.split();
    let mut servo = LegoServo {
        channel_a: channel_a.unwrap(), // safe ? probably...
        channel_b: channel_b.unwrap(),
    };

    spawner.spawn(expect(tick()).await);

    let mut rng = RoscRng;

    let fw = aligned_bytes!("../../cyw43-firmware/43439A0.bin");
    let clm = aligned_bytes!("../../cyw43-firmware/43439A0_clm.bin");
    let nvram = aligned_bytes!("../../cyw43-firmware/nvram_rp2040.bin");

    let pwr = Output::new(p.PIN_23, Level::Low);
    let cs = Output::new(p.PIN_25, Level::High);
    let mut pio = Pio::new(p.PIO0, Irqs);
    let spi = PioSpi::new(
        &mut pio.common,
        pio.sm0,
        DEFAULT_CLOCK_DIVIDER,
        pio.irq0,
        cs,
        p.PIN_24,
        p.PIN_29,
        dma::Channel::new(p.DMA_CH0, Irqs),
        dma::Channel::new(p.DMA_CH1, Irqs),
    );

    static STATE: StaticCell<cyw43::State> = StaticCell::new();
    let state = STATE.init(cyw43::State::new());
    let (net_device, mut control, runner) = cyw43::new(state, pwr, spi, fw, nvram)
        .await;
    spawner.spawn(expect(cyw43_task(runner)).await);

    control.init(clm).await;
    control
        .set_power_management(cyw43::PowerManagementMode::PowerSave)
        .await;

    let config = Config::dhcpv4(Default::default());

    // Generate random seed
    let seed = rng.next_u64();

    // Init network stack
    static RESOURCES: StaticCell<StackResources<3>> = StaticCell::new();
    let (stack, runner) = embassy_net::new(
        net_device,
        config,
        RESOURCES.init(StackResources::new()),
        seed,
    );

    spawner.spawn(expect(net_task(runner)).await);

    while let Err(err) = control
        .join(WIFI_NETWORK, JoinOptions::new(WIFI_PASSWORD.as_bytes()))
        .await
    {
        info!("join failed: {:?}", err);
    }

    info!("waiting for link...");
    stack.wait_link_up().await;

    info!("waiting for DHCP...");
    stack.wait_config_up().await;

    // And now we can use it!
    info!("Stack is up!");

    let mut rx_buffer = [0; 4096];
    let mut tx_buffer = [0; 4096];
    let mut buf = [0; 4096];

    loop {
        let mut socket = TcpSocket::new(stack, &mut rx_buffer, &mut tx_buffer);
        socket.set_timeout(Some(Duration::from_secs(10)));

        control.gpio_set(0, false).await;
        info!("Listening on TCP:1234...");
        led.set_high();
        Timer::after_millis(100).await;
        led.set_low();
        if let Err(e) = socket.accept(1234).await {
            warn!("accept error: {:?}", e);
            continue;
        }

        info!("Received connection from {:?}", socket.remote_endpoint());
        control.gpio_set(0, true).await;

        loop {
            let _ = socket.write_all(b"\n> ").await;
            let n = match socket.read(&mut buf).await {
                Ok(0) => {
                    warn!("read EOF");
                    break;
                }
                Ok(n) => n,
                Err(e) => {
                    warn!("read error: {:?}", e);
                    break;
                }
            };
            let Ok(input) = from_utf8(&buf[..n]) else {
                warn!("Received non-UTF8 data");
                continue;
            };

            info!(" <- Received '{}'", input);
            let cmd = input.parse::<Command>();

            let mut output = SliceWriter::new(&mut buf);
            let mut close = false;
            match cmd {
                Ok(Command::Blink(n)) => {
                    for _ in 0..=n {
                        led.set_high();
                        Timer::after_millis(100).await;
                        led.set_low();
                        Timer::after_millis(100).await;
                    }
                    let _ = write!(output, "-> Blinked {} times", n);
                }
                Ok(Command::GoMillis(m)) => {
                    motor.go_millis(m).await;
                    if m > 0 {
                        let _ = write!(output, "-> Went forward for {} milliseconds", m);
                    } else if m < 0 {
                        let _ = write!(output, "-> Went backward for {} milliseconds", -m);
                    } else {
                        let _ = write!(output, "-> Didn't move");
                    }
                }
                Ok(Command::Steer { direction, level }) => {
                    servo.set_position(direction, level).await;
                    let dir_str = match direction {
                        SteeringDirection::Left => "left",
                        SteeringDirection::Right => "right",
                        SteeringDirection::Center => "center",
                    };
                    if direction == SteeringDirection::Center {
                        let _ = write!(output, "-> Centered steering");
                    } else {
                        let _ = write!(
                            output,
                            "-> Turned {} with intensity level {}",
                            dir_str, level
                        );
                    }
                }
                Ok(Command::Bye) => {
                    let _ = write!(output, "-> Goodbye!");
                    close = true;
                }
                Ok(Command::Quit) => {
                    let _ = write!(output, "-> Shutting down...");
                    Timer::after_millis(100).await;
                    reset_to_usb_boot(0, 0);
                }
                Err(e) => {
                    let _ = write!(output, "-> ERROR: {}", e);
                }
            }

            match socket.write_all(output.as_bytes()).await {
                Ok(()) => {
                    info!("Answered '{}'", output.as_str());
                }
                Err(e) => {
                    warn!("write error: {:?}", e);
                    break;
                }
            };
            if close {
                info!("Closing connection");
                break;
            }
        }
        Timer::after_millis(10).await;
    }
}

use {
    embassy_rp::gpio::Output,
    embassy_time::{
        Duration,
        Timer,
    },
};

/// A simple motor driver that can be used to control
/// a motor using two GPIO pins.
pub struct Motor<'d> {
    pub p1: Output<'d>,
    pub p2: Output<'d>,
}
impl<'d> Motor<'d> {
    pub fn new(p1: Output<'d>, p2: Output<'d>) -> Self {
        Self { p1, p2 }
    }

    // If you have reversed the polarity of the motor, you can swap
    // the forward and backward methods to make it work without
    // changing the wiring.
    pub async fn forward(&mut self) {
        self.p1.set_low();
        self.p2.set_high();
    }

    pub async fn backward(&mut self) {
        self.p1.set_high();
        self.p2.set_low();
    }

    pub async fn stop(&mut self) {
        self.p1.set_low();
        self.p2.set_low();
    }

    pub async fn go_millis(&mut self, ms: i32) {
        if ms > 0 {
            self.forward().await;
            Timer::after(Duration::from_millis(ms as u64)).await;
        } else if ms < 0 {
            self.backward().await;
            Timer::after(Duration::from_millis((-ms) as u64)).await;
        }
        self.stop().await;
    }
}

use {
    crate::expect,
    embassy_rp::pwm::{
        Config,
        PwmOutput,
        SetDutyCycle,
    },
    log::{
        info,
    },
    fixed::types::U12F4,
};

/// Represents the steering direction.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SteeringDirection {
    Center,
    Left,
    Right,
}

/// Allow controlling a servo using the Motor2 pins with PWM.
pub struct LegoServo<'d> {
    pub channel_a: PwmOutput<'d>,
    pub channel_b: PwmOutput<'d>,
}
impl<'d> LegoServo<'d> {
    /// Build a configuration for the PWM slice 7 that allows to control the servo.
    ///
    /// The frequency of the servo is 1200 Hz while the clock of the Pico (sys_clk) is at 125 MHz.
    /// The formula is
    ///     frequency = sys_clk / (divider * top)
    ///  -> 125_000_000 / (4 * 26041) = 1200 Hz/
    ///     and divider * top must fit in 16 bits
    /// (I didn't come with this divider of 4, it seems to be the usual practice)
    pub fn pwm_config() -> Config {
        let mut pwm_config = Config::default();
        pwm_config.divider = U12F4::from_num(4);
        pwm_config.top = 26041;
        pwm_config
    }

    /// Brings the servo back to its center position (0% duty cycle on both channels).
    pub fn center(&mut self) {
        let _ = self.channel_a.set_duty_cycle_fully_off();
        let _ = self.channel_b.set_duty_cycle_fully_off();
        info!("Set steering to CENTER with duty cycle 0% on both channels");
    }

    /// Sets the servo position.
    ///
    /// * `dir`: The direction to turn (`Left` or `Right`).
    /// * `level`: The intensity of the turn from `0` (center) to `7` (maximum lock).
    pub async fn set_position(&mut self, dir: SteeringDirection, level: u8) {
        // Clamp the input level to the maximum supported Lego steps
        let level = level.min(7);

        if level == 0 {
            self.center();
            return;
        }

        // Levels which seem to correspond to the Lego steering steps
        let duty_num = match level {
            1 => 266,  // 26.6%
            2 => 389,  // 38.9%
            3 => 512,  // 51.2%
            4 => 630,  // 63.0%
            5 => 750,  // 75.0%
            6 => 872,  // 87.2%
            _ => 1000, // 100.0%
        };

        // Apply PWM to one channel while forcing the other to ground (0)
        // Note: If Left/Right logic is inverted relative to your chassis build,
        // simply swap the channel_a and channel_b assignments below.
        match dir {
            SteeringDirection::Center => self.center(),
            SteeringDirection::Left => {
                expect(self.channel_a.set_duty_cycle_fully_off())
                    .await;
                expect(self.channel_b.set_duty_cycle_fraction(duty_num, 1000))
                    .await;
            }
            SteeringDirection::Right => {
                expect(self.channel_b.set_duty_cycle_fully_off())
                    .await;
                expect(self.channel_a.set_duty_cycle_fraction(duty_num, 1000))
                    .await;
            }
        }
    }
}

use core::fmt::Write;

/// Wraps a mutable byte slice and implements the `Write` trait for it,
/// allowing formatted strings to be written into the slice.
pub struct SliceWriter<'a> {
    buf: &'a mut [u8],
    cursor: usize,
}

impl SliceWriter<'_> {
    pub fn new(buf: &mut [u8]) -> SliceWriter<'_> {
        SliceWriter { buf, cursor: 0 }
    }
    pub fn reset(&mut self) {
        self.cursor = 0;
    }
    pub fn as_bytes(&self) -> &[u8] {
        &self.buf[..self.cursor]
    }
}

impl<'a> Write for SliceWriter<'a> {
    fn write_str(&mut self, s: &str) -> core::fmt::Result {
        let bytes = s.as_bytes();
        if self.cursor + bytes.len() > self.buf.len() {
            return Err(core::fmt::Error);
        }
        self.buf[self.cursor..self.cursor + bytes.len()]
            .copy_from_slice(bytes);
        self.cursor += bytes.len();
        Ok(())
    }
}

use {
    embassy_time::Timer,
    panic_halt as _,
};

pub async fn expect<T, E: core::fmt::Debug>(res: Result<T, E>) -> T {
    match res {
        Ok(v) => v,
        Err(e) => {
            log::error!("Fatal error: {:?}", e);
            Timer::after_millis(100).await;
            panic!("Fatal error: {:?}", e);
        }
    }
}

use {
    crate::*,
    core::str::FromStr,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Command {
    /// Servo or Motor driving command
    Driving(DrivingCommand),
    /// Toggle the LED
    ToggleLed,
    /// Disconnect the client, but keep the server running
    Bye,
    /// Disconnect the client and shut down the Pico (back to the bootloader)
    Quit,
}

impl From<DrivingCommand> for Command {
    fn from(driving_command: DrivingCommand) -> Self {
        Command::Driving(driving_command)
    }
}

const DEFAULT_STEERING_LEVEL: u8 = 2;

impl FromStr for Command {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut parts = s.split_whitespace();
        let command = parts.next().ok_or("empty command")?;
        match command {
            "led" | "toggle-led" => {
                Ok(Command::ToggleLed)
            }
            "go" | "forward" | "go-forward" => {
                Ok(DrivingCommand::GoForward.into())
            }
            "back" | "backward" | "go-backward" => {
                Ok(DrivingCommand::GoBackward.into())
            }
            "l" | "left" => {
                let level = match parts.next() {
                    Some(level_str) => {
                        let level = level_str
                            .parse::<u8>()
                            .map_err(|_| "level not parsable as u8")?;
                        if level > 7 {
                            return Err("level must be between 1 and 7");
                        }
                        level
                    }
                    None => DEFAULT_STEERING_LEVEL,
                };
                Ok(DrivingCommand::Steer {
                    direction: SteeringDirection::Left,
                    level,
                }.into())
            }
            "r" |"right" => {
                let level = match parts.next() {
                    Some(level_str) => {
                        let level = level_str
                            .parse::<u8>()
                            .map_err(|_| "level not parsable as u8")?;
                        if level > 7 {
                            return Err("level must be between 1 and 7");
                        }
                        level
                    }
                    None => DEFAULT_STEERING_LEVEL,
                };
                Ok(DrivingCommand::Steer {
                    direction: SteeringDirection::Right,
                    level,
                }.into())
            }
            "c" | "center" => Ok(DrivingCommand::Steer {
                direction: SteeringDirection::Center,
                level: 0,
            }.into()),
            "bye" => Ok(Command::Bye),
            "stop" => Ok(DrivingCommand::Stop.into()),
            "q" | "quit" => Ok(Command::Quit),
            _ => Err("unknown command"),
        }
    }
}

use {
    crate::*,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DrivingCommand {
    /// Go forward
    GoForward,
    /// Go backward
    GoBackward,
    /// Stop moving
    Stop,
    /// Turn left, right, or center the steering
    /// with a certain level of intensity (0-7)
    Steer {
        direction: SteeringDirection,
        level: u8, // in [1-7]
    },
}

pub async fn apply_driving_command(
    command: DrivingCommand,
    motor: &mut Motor<'_>,
    servo: &mut LegoServo<'static>,
) {
    info!("Applying driving command: {:?}", command);
    match command {
        DrivingCommand::GoForward => {
            motor.forward().await;
        }
        DrivingCommand::GoBackward => {
            motor.backward().await;
        }
        DrivingCommand::Stop => {
            motor.stop().await;
        }
        DrivingCommand::Steer { direction, level } => {
            servo.set_position(direction, level).await;
        }
    }
}

use {
    embassy_rp::gpio::Output,
    embassy_time::{
        Duration,
        Timer,
    },
};

/// A simple motor driver that can be used to control a motor using two GPIO pins.
pub struct Motor<'d> {
    pub p1: Output<'d>,
    pub p2: Output<'d>,
}
impl<'d> Motor<'d> {
    pub fn new(p1: Output<'d>, p2: Output<'d>) -> Self {
        Self { p1, p2 }
    }

    // If you have reversed the polarity of the motor,
    // you can swap the forward and backward methods
    // to make it work without changing the wiring.
    pub async fn forward(&mut self) {
        self.p1.set_low();
        self.p2.set_high();
    }

    pub async fn backward(&mut self) {
        self.p1.set_high();
        self.p2.set_low();
    }

    pub async fn stop(&mut self) {
        self.p1.set_low();
        self.p2.set_low();
    }

    pub async fn go_millis(&mut self, ms: i32) {
        if ms > 0 {
            self.forward().await;
            Timer::after(Duration::from_millis(ms as u64)).await;
        } else if ms < 0 {
            self.backward().await;
            Timer::after(Duration::from_millis((-ms) as u64)).await;
        }
        self.stop().await;
    }
}

use {
    crate::expect,
    embassy_rp::pwm::{
        Config,
        PwmOutput,
        SetDutyCycle,
    },
    fixed::types::U12F4,
    log::info,
};

/// Represents the steering direction.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SteeringDirection {
    Center,
    Left,
    Right,
}

/// Allow controlling a servo using the Motor2 pins with PWM.
pub struct LegoServo<'d> {
    pub channel_a: PwmOutput<'d>,
    pub channel_b: PwmOutput<'d>,
}
impl<'d> LegoServo<'d> {
    /// Build a configuration for the PWM slice 7 that allows to control the servo.
    ///
    /// The frequency of the servo is 1200 Hz while the clock of the Pico (sys_clk) is at 125 MHz.
    /// The formula is
    ///     frequency = sys_clk / (divider * top)
    ///  -> 125_000_000 / (4 * 26041) = 1200 Hz/
    ///     and divider * top must fit in 16 bits
    /// (I didn't come with this divider of 4, it seems to be the usual practice)
    pub fn pwm_config() -> Config {
        let mut pwm_config = Config::default();
        pwm_config.divider = U12F4::from_num(4);
        pwm_config.top = 26041;
        pwm_config
    }

    /// Brings the servo back to its center position (0% duty cycle on both channels).
    pub fn center(&mut self) {
        let _ = self.channel_a.set_duty_cycle_fully_off();
        let _ = self.channel_b.set_duty_cycle_fully_off();
        info!("Set steering to CENTER with duty cycle 0% on both channels");
    }

    /// Sets the servo position.
    ///
    /// * `dir`: The direction to turn (`Left` or `Right`).
    /// * `level`: The intensity of the turn from `0` (center) to `7` (maximum lock).
    pub async fn set_position(&mut self, dir: SteeringDirection, level: u8) {
        // Clamp the input level to the maximum supported Lego steps
        let level = level.min(7);

        if level == 0 {
            self.center();
            return;
        }

        // Levels which seem to correspond to the Lego steering steps
        let duty_num = match level {
            1 => 266,  // 26.6%
            2 => 389,  // 38.9%
            3 => 512,  // 51.2%
            4 => 630,  // 63.0%
            5 => 750,  // 75.0%
            6 => 872,  // 87.2%
            _ => 1000, // 100.0%
        };

        // Apply PWM to one channel while forcing the other to ground (0)
        // Note: If Left/Right logic is inverted relative to your chassis build,
        // simply swap the channel_a and channel_b assignments below.
        match dir {
            SteeringDirection::Center => self.center(),
            SteeringDirection::Left => {
                expect(self.channel_a.set_duty_cycle_fully_off())
                    .await;
                expect(self.channel_b.set_duty_cycle_fraction(duty_num, 1000))
                    .await;
            }
            SteeringDirection::Right => {
                expect(self.channel_b.set_duty_cycle_fully_off()).await;
                expect(self.channel_a.set_duty_cycle_fraction(duty_num, 1000)).await;
            }
        }
    }
}

use core::fmt::Write;

/// Wraps a mutable byte slice and implements the `Write` trait for it,
/// allowing formatted strings to be written into the slice.
pub struct SliceWriter<'a> {
    buf: &'a mut [u8],
    cursor: usize,
}

impl SliceWriter<'_> {
    pub fn new(buf: &mut [u8]) -> SliceWriter<'_> {
        SliceWriter { buf, cursor: 0 }
    }
    pub fn reset(&mut self) {
        self.cursor = 0;
    }
    pub fn as_bytes(&self) -> &[u8] {
        &self.buf[..self.cursor]
    }
}

impl<'a> Write for SliceWriter<'a> {
    fn write_str(&mut self, s: &str) -> core::fmt::Result {
        let bytes = s.as_bytes();
        if self.cursor + bytes.len() > self.buf.len() {
            return Err(core::fmt::Error);
        }
        self.buf[self.cursor..self.cursor + bytes.len()]
            .copy_from_slice(bytes);
        self.cursor += bytes.len();
        Ok(())
    }
}

use {
    embassy_time::Timer,
    panic_halt as _,
};

pub async fn expect<T, E: core::fmt::Debug>(res: Result<T, E>) -> T {
    match res {
        Ok(v) => v,
        Err(e) => {
            log::error!("Fatal error: {:?}", e);
            Timer::after_millis(100).await;
            panic!("Fatal error: {:?}", e);
        }
    }
}

[package]
name = "picomobile-control"
version = "0.1.0"
edition = "2024"

[dependencies]
axum = "0.8"
tokio = { version = "1", features = ["full"] }
tower-http = { version = "0.6", features = ["fs"] }
serde = { version = "1", features = ["derive"] }
clap = { version = "4", features = ["derive"] }

use std::{net::SocketAddr, sync::Arc, time::Duration};

use axum::{
    extract::State,
    http::StatusCode,
    response::IntoResponse,
    routing::post,
    Json, Router,
};
use clap::Parser;
use serde::Deserialize;
use tokio::{
    io::AsyncWriteExt,
    net::TcpStream,
    sync::mpsc,
    time::sleep,
};
use tower_http::services::ServeDir;

#[derive(Parser)]
struct Args {
    /// Car IP address
    #[arg(long)]
    car_ip: String,

    /// Car TCP port
    #[arg(long)]
    car_port: u16,

    /// Web GUI port
    #[arg(long)]
    gui_port: u16,
}

#[derive(Clone)]
struct AppState {
    command_tx: mpsc::Sender<String>,
}

#[derive(Deserialize)]
struct CommandRequest {
    command: String,
}

#[tokio::main]
async fn main() {
    let args = Args::parse();

    let car_addr = Arc::new(format!("{}:{}", args.car_ip, args.car_port));
    let (command_tx, command_rx) = mpsc::channel::<String>(8);

    tokio::spawn(writer_task(
        Arc::clone(&car_addr),
        command_rx,
    ));

    let state = AppState {
        command_tx,
    };

    let app = Router::new()
        .route("/api/command", post(send_command))
        .fallback_service(ServeDir::new("static"))
        .with_state(state);

    let addr = SocketAddr::from(([0, 0, 0, 0], args.gui_port));

    println!("GUI available on http://{}", addr);

    let listener = tokio::net::TcpListener::bind(addr)
        .await
        .expect("bind failed");

    axum::serve(listener, app)
        .await
        .expect("server failed");
}

async fn send_command(
    State(state): State<AppState>,
    Json(req): Json<CommandRequest>,
) -> impl IntoResponse {
    match state.command_tx.try_send(req.command) {
        Ok(_) => StatusCode::NO_CONTENT,

        Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
            // Queue already contains newer commands waiting to be sent.
            // Drop this one.
            StatusCode::NO_CONTENT
        }

        Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
            StatusCode::SERVICE_UNAVAILABLE
        }
    }
}

async fn writer_task(
    car_addr: Arc<String>,
    mut rx: mpsc::Receiver<String>,
) {
    let mut stream: Option<TcpStream> = None;

    while let Some(command) = rx.recv().await {
        let payload = format!("{command}\n");

        loop {
            if stream.is_none() {
                match TcpStream::connect(car_addr.as_str()).await {
                    Ok(s) => {
                        println!("Connected to car");
                        stream = Some(s);
                    }
                    Err(e) => {
                        eprintln!("Failed to connect to car: {e}");
                        break;
                    }
                }
            }

            let write_result = match stream.as_mut() {
                Some(s) => s.write_all(payload.as_bytes()).await,
                None => unreachable!(),
            };

            match write_result {
                Ok(_) => {
                    break;
                }
                Err(e) => {
                    eprintln!(
                        "Connection lost while sending '{}': {}",
                        command,
                        e
                    );

                    stream = None;

                    sleep(Duration::from_millis(500)).await;
                }
            }
        }
    }

    println!("Writer task exiting");
}

const pressed = {
    forward: false,
    backward: false,
    left: false,
    right: false,
};

async function sendCommand(command) {
    try {
        const response = await fetch('/api/command', {
            method: 'POST',
            headers: {
                'Content-Type': 'application/json',
            },
            body: JSON.stringify({ command }),
        });

        if (!response.ok) {
            console.error(`Command failed: ${command}`);
        }
    } catch (error) {
        console.error(`Error sending "${command}":`, error);
    }
}

function press(action) {
    if (pressed[action]) {
        return;
    }

    pressed[action] = true;

    // Immediate transmission for low latency.
    sendCommand(action);
}

function release(action) {
    if (!pressed[action]) {
        return;
    }

    pressed[action] = false;

    if (
        (action === 'forward' || action === 'backward') &&
        !pressed.forward &&
        !pressed.backward
    ) {
        sendCommand('stop');
    }

    if (
        (action === 'left' || action === 'right') &&
        !pressed.left &&
        !pressed.right
    ) {
        sendCommand('center');
    }
}

async function repeatLoop() {
    while (true) {
        if (pressed.forward) {
            sendCommand('forward');
        } else if (pressed.backward) {
            sendCommand('backward');
        }
        if (pressed.left) {
            sendCommand('left');
        } else if (pressed.right) {
            sendCommand('right');
        }
        await new Promise(resolve => setTimeout(resolve, 50));
    }
}

function initializeButtons() {
    document.querySelectorAll('[data-action]').forEach(button => {
        const action = button.dataset.action;

        button.addEventListener('mousedown', () => press(action));
        button.addEventListener('mouseup', () => release(action));

        button.addEventListener('mouseleave', () => {
            if (pressed[action]) {
                release(action);
            }
        });

        // Touch support for phones/tablets.
        button.addEventListener('touchstart', event => {
            event.preventDefault();
            press(action);
        });

        button.addEventListener('touchend', event => {
            event.preventDefault();
            release(action);
        });

        button.addEventListener('touchcancel', event => {
            event.preventDefault();
            release(action);
        });
    });
}

function initializeKeyboard() {
    const keyMap = {
        ArrowUp: 'forward',
        ArrowDown: 'backward',
        ArrowLeft: 'left',
        ArrowRight: 'right',
    };

    document.addEventListener('keydown', event => {
        const action = keyMap[event.key];

        if (!action) {
            return;
        }

        event.preventDefault();

        if (!event.repeat) {
            press(action);
        }
    });

    document.addEventListener('keyup', event => {
        const action = keyMap[event.key];

        if (!action) {
            return;
        }

        event.preventDefault();
        release(action);
    });

    // Prevent "stuck" keys if focus is lost.
    window.addEventListener('blur', () => {
        for (const action of Object.keys(pressed)) {
            release(action);
        }
    });
}

document.addEventListener('DOMContentLoaded', () => {
    initializeButtons();
    initializeKeyboard();
    repeatLoop();
});

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="utf-8">
    <title>La PicoMobile</title>
    <link rel="stylesheet" href="style.css">
</head>
<body>

<h1>La PicoMobile</h1>

<div class="controller">
    <button data-action="forward">Forward ↑</button>

    <div class="steering">
        <button data-action="left">← Left</button>
        <button data-action="right">Right →</button>
    </div>

    <button data-action="backward">Backward ↓</button>
    <button data-action="toggle-led">Toggle LED</button>
</div>

<script src="app.js"></script>

</body>
</html>

body {
    font-family: sans-serif;
    text-align: center;
    margin-top: 2rem;
}

.controller {
    display: flex;
    flex-direction: column;
    gap: 1rem;
    align-items: center;
}

.steering {
    display: flex;
    gap: 0.5rem;
}

button {
    min-width: 120px;
    min-height: 60px;
    font-size: 1.1rem;
    cursor: pointer;
}