aves-tyto 0.5.0

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, Mutex as StdMutex};
use std::time::Duration;

use tokio::sync::{Mutex, mpsc, oneshot};
use tokio::task::{AbortHandle, JoinHandle};
use tokio::{spawn, time::sleep};
use uuid::Uuid;

use aves::bus::BusRpcContext;
use aves::flow::{FlowTriggerRecord, FlowTriggerSourceService, LifecycleCorrelation, NodeRegistrationSupply};
use aves::runtime::ServiceReader;
use aves::utils::uuid::SpectaUuid;

use super::channels::{
    CHANNEL_MESSAGE_BUFFER_SIZE, InboundEvent, MessagingPlugin, MessagingPluginHandle, MessagingTask,
};
use super::database::MessageStore;
use super::messages::{MessagingRequest, MessagingResponse};
use super::registrations::{NodeCommand, TytoNodeHandle, message_node_registration, permission_node_registration};
use super::types::{
    ConnectionTransport, MessagingCredentials, MessagingError, MessagingInteraction, MessagingSendReceipt,
    NormalizedMessage, PersistencePolicy, SubscriptionTriggerConfig,
};

#[cfg(not(test))]
use super::channels::for_kind;

#[path = "connections.rs"]
mod connections;
#[path = "delivery.rs"]
mod delivery;
#[path = "flows.rs"]
mod flows;
#[path = "listeners.rs"]
mod listeners;

type Subscribers = Arc<Mutex<HashMap<Uuid, mpsc::Sender<NormalizedMessage>>>>;

struct ActorInbound {
    connection_id: SpectaUuid,
    event: InboundEvent,
}

struct ActorInteraction {
    connection_id: SpectaUuid,
    interaction: MessagingInteraction,
}

struct SendCompletion {
    reply_to: oneshot::Sender<Result<MessagingResponse, MessagingError>>,
    connection_id: SpectaUuid,
    body: String,
    metadata: Option<String>,
    retention_seconds: i64,
    result: Result<MessagingSendReceipt, MessagingError>,
    persistence: PersistencePolicy,
}

struct PermSendCompleted {
    request_id: String,
    connection_id: SpectaUuid,
    retention_seconds: i64,
    body: String,
    metadata: Option<String>,
    result: Result<String, MessagingError>,
}

struct EditCompleted {
    request_id: String,
    result: Result<(), MessagingError>,
}

struct NodePending {
    connection_id: SpectaUuid,
    retention_seconds: i64,
    body: String,
    metadata: Option<String>,
    result: Option<Result<MessagingSendReceipt, MessagingError>>,
}

struct Listener {
    trigger: FlowTriggerRecord,
    configuration: SubscriptionTriggerConfig,
    task: JoinHandle<()>,
    filter_task: MessagingTask,
}

impl Drop for Listener {
    fn drop(&mut self) {
        self.task.abort();
        self.filter_task.abort();
    }
}

struct Inflow {
    connection_id: SpectaUuid,
    retention_seconds: i64,
    plugin: MessagingPluginHandle,
    transport: ConnectionTransport,
    credentials: MessagingCredentials,
    credential_generation: u64,
    subscribers: Subscribers,
    task: Option<JoinHandle<()>>,
    producer: Arc<StdMutex<ProducerState>>,
}

impl Drop for Inflow {
    fn drop(&mut self) {
        if let Some(task) = &self.task {
            task.abort();
        }
        let mut producer = self.producer.lock().expect("inflow task state");
        producer.stopping = true;
        if let Some(handle) = &producer.handle {
            handle.abort();
        }
    }
}

#[derive(Default)]
struct ProducerState {
    handle: Option<AbortHandle>,
    stopping: bool,
}

pub(super) struct Actor {
    receiver: mpsc::Receiver<MessagingRequest>,
    broadcast_tx: mpsc::Sender<MessagingResponse>,
    config_dir: PathBuf,
    store: Option<MessageStore>,
    node_rx: Option<mpsc::Receiver<NodeCommand>>,
    node_done_tx: mpsc::Sender<(LifecycleCorrelation, Result<MessagingSendReceipt, MessagingError>)>,
    node_done_rx: Option<mpsc::Receiver<(LifecycleCorrelation, Result<MessagingSendReceipt, MessagingError>)>>,
    perm_tx: mpsc::Sender<PermSendCompleted>,
    perm_rx: Option<mpsc::Receiver<PermSendCompleted>>,
    edit_tx: mpsc::Sender<EditCompleted>,
    edit_rx: Option<mpsc::Receiver<EditCompleted>>,
    node_pending: HashMap<LifecycleCorrelation, NodePending>,
    inbound_tx: mpsc::Sender<ActorInbound>,
    inbound_rx: Option<mpsc::Receiver<ActorInbound>>,
    interaction_tx: mpsc::Sender<ActorInteraction>,
    interaction_rx: Option<mpsc::Receiver<ActorInteraction>>,
    send_tx: mpsc::Sender<SendCompletion>,
    send_rx: Option<mpsc::Receiver<SendCompletion>>,
    worker_tick_tx: mpsc::Sender<()>,
    worker_tick_rx: Option<mpsc::Receiver<()>>,
    worker_task: Option<JoinHandle<()>>,
    credentials: HashMap<SpectaUuid, MessagingCredentials>,
    credential_generation: HashMap<SpectaUuid, u64>,
    flows: ServiceReader<dyn FlowTriggerSourceService>,
    knowledge_context: BusRpcContext,
    listeners: HashMap<Uuid, Listener>,
    inflows: HashMap<SpectaUuid, Inflow>,
    running: bool,
}

impl Actor {
    /// Constructs the actor and registers its Flow node kinds.
    pub(super) async fn new(
        receiver: mpsc::Receiver<MessagingRequest>,
        broadcast_tx: mpsc::Sender<MessagingResponse>,
        config_dir: PathBuf,
        node_registration_supply: NodeRegistrationSupply,
        flows: ServiceReader<dyn FlowTriggerSourceService>,
        knowledge_context: BusRpcContext,
    ) -> Result<Self, MessagingError> {
        let store = Some(MessageStore::open(&config_dir).await?);
        let (inbound_tx, inbound_rx) = mpsc::channel(CHANNEL_MESSAGE_BUFFER_SIZE);
        let (interaction_tx, interaction_rx) = mpsc::channel(CHANNEL_MESSAGE_BUFFER_SIZE);
        let (send_tx, send_rx) = mpsc::channel(CHANNEL_MESSAGE_BUFFER_SIZE);
        let (worker_tick_tx, worker_tick_rx) = mpsc::channel(8);
        let (node_tx, node_rx) = mpsc::channel(64);
        let (node_done_tx, node_done_rx) = mpsc::channel(64);
        let (perm_tx, perm_rx) = mpsc::channel(64);
        let (edit_tx, edit_rx) = mpsc::channel(64);
        let node_handle = Arc::new(TytoNodeHandle { tx: node_tx });
        node_registration_supply.push(message_node_registration(node_handle.clone()));
        node_registration_supply.push(permission_node_registration(node_handle));
        Ok(Self {
            receiver,
            broadcast_tx,
            config_dir,
            store,
            node_rx: Some(node_rx),
            node_done_tx,
            node_done_rx: Some(node_done_rx),
            perm_tx,
            perm_rx: Some(perm_rx),
            edit_tx,
            edit_rx: Some(edit_rx),
            node_pending: HashMap::new(),
            inbound_tx,
            inbound_rx: Some(inbound_rx),
            interaction_tx,
            interaction_rx: Some(interaction_rx),
            send_tx,
            send_rx: Some(send_rx),
            worker_tick_tx,
            worker_tick_rx: Some(worker_tick_rx),
            worker_task: None,
            credentials: HashMap::new(),
            credential_generation: HashMap::new(),
            flows,
            knowledge_context,
            listeners: HashMap::new(),
            inflows: HashMap::new(),
            running: false,
        })
    }

    async fn store(&mut self) -> Result<&MessageStore, MessagingError> {
        if self.store.is_none() {
            self.store = Some(MessageStore::open(&self.config_dir).await?);
        }
        Ok(self.store.as_ref().expect("opened"))
    }

    /// Reconciles listeners before starting the periodic maintenance worker.
    async fn start_server(&mut self) -> Result<MessagingResponse, MessagingError> {
        self.running = true;
        let result = self.reconcile().await;
        if self.worker_task.is_none() {
            self.worker_task = Some(spawn(Self::worker_loop(self.worker_tick_tx.clone())));
        }
        result
    }

    /// Stops owned tasks and releases the store after maintenance.
    async fn stop_server(&mut self) -> Result<MessagingResponse, MessagingError> {
        self.running = false;
        if let Some(task) = self.worker_task.take() {
            task.abort();
        }
        self.listeners.clear();
        self.inflows.clear();
        if let Some(store) = self.store.take() {
            store.maintain().await;
        }
        Ok(MessagingResponse::Success)
    }

    /// Selects the plugin used only while reconciling inbound listeners.
    fn reconciliation_plugin(transport: &ConnectionTransport) -> Box<dyn MessagingPlugin> {
        #[cfg(test)]
        {
            let _ = transport;
            tests::pending_channel()
        }
        #[cfg(not(test))]
        {
            for_kind(transport.provider())
        }
    }

    /// Sends periodic maintenance ticks, waiting one full period before the first.
    async fn worker_loop(tick: mpsc::Sender<()>) {
        loop {
            sleep(Duration::from_secs(300)).await;
            if tick.send(()).await.is_err() {
                break;
            }
        }
    }

    /// Runs the sequential actor loop; channel closure terminates the actor.
    pub(super) async fn run(mut self) {
        let Some(mut inbound_rx) = self.inbound_rx.take() else {
            return;
        };
        let Some(mut send_rx) = self.send_rx.take() else {
            return;
        };
        let Some(mut node_rx) = self.node_rx.take() else {
            return;
        };
        let Some(mut node_done_rx) = self.node_done_rx.take() else {
            return;
        };
        let Some(mut perm_rx) = self.perm_rx.take() else {
            return;
        };
        let Some(mut edit_rx) = self.edit_rx.take() else {
            return;
        };
        let Some(mut interaction_rx) = self.interaction_rx.take() else {
            return;
        };
        let Some(mut worker_tick_rx) = self.worker_tick_rx.take() else {
            return;
        };
        loop {
            tokio::select! {
                request = self.receiver.recv() => {
                    let Some(request) = request else { break; };
                    self.handle_message(request).await;
                }
                tick = worker_tick_rx.recv() => {
                    let Some(()) = tick else { break; };
                    self.handle_worker_tick().await;
                }
                inbound = inbound_rx.recv() => {
                    let Some(event) = inbound else { break; };
                    self.handle_inbound_delivery(event).await;
                }
                interaction = interaction_rx.recv() => {
                    let Some(event) = interaction else { break; };
                    self.handle_interaction(event).await;
                }
                completion = send_rx.recv() => {
                    let Some(completion) = completion else { break; };
                    self.handle_send_completion(completion).await;
                }
                completion = perm_rx.recv() => {
                    let Some(completion) = completion else { break; };
                    self.handle_permission_send_completion(completion).await;
                }
                completion = edit_rx.recv() => {
                    let Some(completion) = completion else { break; };
                    self.handle_edit_completion(completion).await;
                }
                command = node_rx.recv() => {
                    let Some(command) = command else { break; };
                    self.handle_node_command(command).await;
                }
                completion = node_done_rx.recv() => {
                    let Some((correlation, result)) = completion else { break; };
                    if let Some(pending) = self.node_pending.get_mut(&correlation) {
                        pending.result = Some(result);
                    }
                }
            }
        }
    }
}

impl Actor {
    fn connection_uuid(id: &str) -> Result<SpectaUuid, MessagingError> {
        Uuid::parse_str(id)
            .map(SpectaUuid)
            .map_err(|error| MessagingError::ConnectionNotFound {
                connection_id: error.to_string(),
            })
    }
}

#[rustfmt::skip]
#[cfg(test)]
mod tests;