prns-runtime-embassy 0.3.4

use core::cell::RefCell;

use embassy_sync::blocking_mutex::raw::RawMutex;
use embassy_sync::blocking_mutex::Mutex as BlockingMutex;
use embassy_sync::channel::Sender;
use embassy_sync::signal::Signal;
use portable_atomic::{AtomicU64, Ordering};

use crate::engine::{
    CloseLink, CommandId, IssuedCommand, PacketReceiptDelivered, PrnsCommand, Respond, RespondData,
    RespondPayload, SendSinglePacket, SendSinglePacketFailure, SendSinglePacketPayload, Settlement,
};
use crate::routing::links::LinkId;
use crate::wire::DestinationHash;

use super::super::request_endpoints::RespondToken;
use super::super::{PrnsNodeApi, SendError};

const NO_AWAITER: u64 = u64::MAX;

/// Fixed completion storage for at most `N` concurrently awaited commands.
pub struct CompletionPool<M: RawMutex, const N: usize> {
    next_id: AtomicU64,
    awaited: BlockingMutex<M, RefCell<[u64; N]>>,
    slots: [Signal<M, Settlement>; N],
}

impl<M: RawMutex, const N: usize> Default for CompletionPool<M, N> {
    fn default() -> Self {
        Self::new()
    }
}

impl<M: RawMutex, const N: usize> CompletionPool<M, N> {
    #[must_use]
    pub const fn new() -> Self {
        Self {
            next_id: AtomicU64::new(0),
            awaited: BlockingMutex::new(RefCell::new([NO_AWAITER; N])),
            slots: [const { Signal::new() }; N],
        }
    }

    fn mint(&self) -> CommandId {
        loop {
            let id = self.next_id.fetch_add(1, Ordering::Relaxed);
            if id != NO_AWAITER {
                return CommandId(id);
            }
        }
    }

    fn claim(&self, id: CommandId) -> Option<usize> {
        self.awaited.lock(|cell| {
            let mut awaited = cell.borrow_mut();
            let slot = awaited.iter().position(|&a| a == NO_AWAITER)?;
            self.slots[slot].reset();
            awaited[slot] = id.0;
            Some(slot)
        })
    }

    fn release(&self, slot: usize, id: CommandId) {
        self.awaited.lock(|cell| {
            let mut awaited = cell.borrow_mut();
            if awaited[slot] == id.0 {
                awaited[slot] = NO_AWAITER;
                self.slots[slot].reset();
            }
        });
    }

    fn settle(&self, id: CommandId, settlement: Settlement) -> bool {
        self.awaited.lock(|cell| {
            let mut awaited = cell.borrow_mut();
            match awaited.iter().position(|&a| a == id.0) {
                Some(slot) => {
                    awaited[slot] = NO_AWAITER;
                    self.slots[slot].signal(settlement);
                    true
                }
                None => false,
            }
        })
    }

    async fn parked(&self, slot: usize) -> Settlement {
        self.slots[slot].wait().await
    }
}

pub struct PrnsNodeHandle<'a, M: RawMutex, const COMMANDS: usize, const N: usize> {
    commands: Sender<'a, M, IssuedCommand, COMMANDS>,
    pool: &'a CompletionPool<M, N>,
}

impl<M: RawMutex, const COMMANDS: usize, const N: usize> Clone
    for PrnsNodeHandle<'_, M, COMMANDS, N>
{
    fn clone(&self) -> Self {
        *self
    }
}

impl<M: RawMutex, const COMMANDS: usize, const N: usize> Copy
    for PrnsNodeHandle<'_, M, COMMANDS, N>
{
}

impl<'a, M: RawMutex, const COMMANDS: usize, const N: usize> PrnsNodeHandle<'a, M, COMMANDS, N> {
    #[must_use]
    pub fn new(
        commands: Sender<'a, M, IssuedCommand, COMMANDS>,
        pool: &'a CompletionPool<M, N>,
    ) -> Self {
        Self { commands, pool }
    }

    /// Queues a command without awaiting settlement and returns its ID, or `None` when the command lane is full.
    pub fn issue(&self, command: PrnsCommand) -> Option<CommandId> {
        let id = self.pool.mint();
        self.commands.try_send(IssuedCommand { id, command }).ok()?;
        Some(id)
    }

    /// Sends one packet and awaits proof, returning `Busy` when all `N` completion slots are claimed.
    pub async fn send_single_packet(
        &self,
        destination: DestinationHash,
        data: &[u8],
    ) -> Result<PacketReceiptDelivered, SendError<SendSinglePacketFailure>> {
        let payload =
            SendSinglePacketPayload::from_slice(data).map_err(|()| SendError::PayloadTooLarge)?;
        let id = self.pool.mint();
        let slot = self.pool.claim(id).ok_or(SendError::Busy)?;
        let _guard = SlotGuard {
            pool: self.pool,
            slot,
            id,
        };
        self.commands
            .try_send(IssuedCommand {
                id,
                command: PrnsCommand::SendSinglePacket(SendSinglePacket {
                    destination,
                    payload,
                }),
            })
            .map_err(|_| SendError::NodeStopped)?;
        match self.pool.parked(slot).await {
            Settlement::SendSinglePacket(result) => result.map_err(SendError::Failed),
            _ => Err(SendError::NodeStopped),
        }
    }

    /// Responds inline; returns `false` when the body exceeds the link MDU or the command lane is full.
    pub fn respond_packed(&self, responder: RespondToken, packed: &[u8]) -> bool {
        match RespondData::from_slice(packed) {
            Ok(data) => self.respond_owned_packed(responder, data),
            Err(_) => false,
        }
    }

    /// Moves a prebuilt response into the command lane, returning `false` when full.
    pub fn respond_owned_packed(&self, responder: RespondToken, data: RespondData) -> bool {
        self.issue(PrnsCommand::Respond(Respond {
            link_id: responder.link_id,
            request_id: responder.request_id,
            payload: RespondPayload::Packed(data),
        }))
        .is_some()
    }

    pub fn respond_static_bytes(&self, responder: RespondToken, data: &'static [u8]) -> bool {
        self.issue(PrnsCommand::Respond(Respond {
            link_id: responder.link_id,
            request_id: responder.request_id,
            payload: RespondPayload::StaticBytes(data),
        }))
        .is_some()
    }

    #[cfg(feature = "large-static-responses")]
    pub fn respond_static_file(
        &self,
        responder: RespondToken,
        name: &'static str,
        bytes: &'static [u8],
    ) -> bool {
        self.issue(PrnsCommand::Respond(Respond {
            link_id: responder.link_id,
            request_id: responder.request_id,
            payload: RespondPayload::StaticFile { name, bytes },
        }))
        .is_some()
    }

    /// Sever an active link. Returns `false` if the command lane is full.
    pub fn close_link(&self, link_id: LinkId) -> bool {
        self.issue(PrnsCommand::CloseLink(CloseLink { link_id }))
            .is_some()
    }

    pub(super) fn settle(&self, id: CommandId, settlement: Settlement) -> bool {
        self.pool.settle(id, settlement)
    }
}

struct SlotGuard<'a, M: RawMutex, const N: usize> {
    pool: &'a CompletionPool<M, N>,
    slot: usize,
    id: CommandId,
}

impl<M: RawMutex, const N: usize> Drop for SlotGuard<'_, M, N> {
    fn drop(&mut self) {
        self.pool.release(self.slot, self.id);
    }
}

impl<M: RawMutex, const COMMANDS: usize, const N: usize> PrnsNodeApi
    for PrnsNodeHandle<'_, M, COMMANDS, N>
{
    fn issue(&self, command: PrnsCommand) -> Option<CommandId> {
        self.issue(command)
    }

    async fn send_single_packet(
        &self,
        destination: DestinationHash,
        data: &[u8],
    ) -> Result<PacketReceiptDelivered, SendError<SendSinglePacketFailure>> {
        self.send_single_packet(destination, data).await
    }

    fn respond_packed(&self, responder: RespondToken, packed: &[u8]) -> bool {
        self.respond_packed(responder, packed)
    }

    fn close_link(&self, link_id: LinkId) -> bool {
        self.close_link(link_id)
    }
}

#[cfg(test)]
mod tests;