// src/uart/driver.rs
use defmt::unwrap;
use embassy_executor::Spawner;
use embassy_futures::select::{select, Either};
use embassy_stm32::usart::BufferedUart;
use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
use embassy_sync::pipe::Pipe;
use embassy_time::Duration;
use embedded_io_async::{Read, Write};

use crate::uart::safety::{RX_PIPE_CAP, TX_PIPE_CAP};

pub struct UartHandle {
    pub tx: &'static Pipe<CriticalSectionRawMutex, TX_PIPE_CAP>,
    pub rx: &'static Pipe<CriticalSectionRawMutex, RX_PIPE_CAP>,
    pub yield_period: Duration,
}

#[embassy_executor::task]
pub async fn uart_task(
    mut uart: BufferedUart<'static>,
    tx_pipe: &'static Pipe<CriticalSectionRawMutex, TX_PIPE_CAP>,
    rx_pipe: &'static Pipe<CriticalSectionRawMutex, RX_PIPE_CAP>,
) {
    let mut rx_byte = [0u8; 1];
    let mut tx_buf = [0u8; 64];

    loop {
        let rx_fut = uart.read(&mut rx_byte);
        let tx_fut = async {
            let n = tx_pipe.read(&mut tx_buf).await;
            n
        };

        match select(rx_fut, tx_fut).await {
            // Incoming data from UART hardware
            Either::First(res) => {
                if let Ok(_) = res {
                    // Forward to RX pipe and echo to TX pipe (same behavior as before)
                    let _ = rx_pipe.write(&rx_byte).await;
                    let _ = tx_pipe.try_write(&rx_byte);
                }
            }
            // Outgoing data waiting in TX pipe
            Either::Second(n) => {
                unwrap!(uart.write(&tx_buf[..n]).await);
            }
        }
    }
}

pub fn spawn_for(
    spawner: &Spawner,
    uart: BufferedUart<'static>,
    tx_pipe: &'static Pipe<CriticalSectionRawMutex, TX_PIPE_CAP>,
    rx_pipe: &'static Pipe<CriticalSectionRawMutex, RX_PIPE_CAP>,
    yield_period: Duration,
) -> UartHandle {
    spawner.spawn(uart_task(uart, tx_pipe, rx_pipe)).unwrap();
    UartHandle {
        tx: tx_pipe,
        rx: rx_pipe,
        yield_period,
    }
}