macltok

mod commands;

mod executor;

mod media;

mod mutations;

mod room;

mod timeline;

mod watchers;

use std::{

    collections::{HashMap, VecDeque},

    future::Future,

    panic::AssertUnwindSafe,

    sync::Arc,

};

use futures_util::FutureExt;

use matrix_sdk::{Client, SendOutsideWasm, SessionChange};

use matrix_sdk_ui::sync_service::SyncService;

use ruma::OwnedRoomId;

use thiserror::Error as ThisError;

use tokio::sync::{broadcast, mpsc, watch};

pub(crate) use media::MediaService;

use self::{

    commands::Job,

    executor::{EngineBoxFuture, TaskSet},

    mutations::{Mutation, MutationJob},

    room::{ActorNote, ActorRequest, RoomActor},

    timeline::{SessionContext, Supervisor},

};

use crate::{

    AccountCommand, ClientHandle, Command, ConnectionState, EncryptionCommand, EngineErrorKind,

    EngineFailure, Error, ErrorKind, Event, EventStream, LifecycleCommand, MediaCommand,

    NotificationCommand, RoomCommand, Termination, TimelineCommand,

    transport::{DriverPort, Request, Responder, transport_pair},

};

#[derive(Clone, Debug, Eq, PartialEq)]

#[non_exhaustive]

pub struct EngineConfig {

    pub command_capacity: usize,

    pub event_capacity: usize,

    pub initial_history_limit: u16,

    pub max_media_bytes: usize,

    pub max_media_sources: usize,

}

impl EngineConfig {

    #[must_use]

    pub const fn with_command_capacity(mut self, capacity: usize) -> Self {

        self.command_capacity = capacity;

        self

    }

    #[must_use]

    pub const fn with_event_capacity(mut self, capacity: usize) -> Self {

        self.event_capacity = capacity;

        self

    }

    #[must_use]

    pub const fn with_initial_history_limit(mut self, limit: u16) -> Self {

        self.initial_history_limit = limit;

        self

    }

    #[must_use]

    pub const fn with_max_media_bytes(mut self, bytes: usize) -> Self {

        self.max_media_bytes = bytes;

        self

    }

    #[must_use]

    pub const fn with_max_media_sources(mut self, sources: usize) -> Self {

        self.max_media_sources = sources;

        self

    }

}

impl Default for EngineConfig {

    fn default() -> Self {

        Self {

            command_capacity: 64,

            event_capacity: 128,

            initial_history_limit: 30,

            max_media_bytes: 20 * 1024 * 1024,

            max_media_sources: 2_048,

        }

    }

}

#[derive(Debug, ThisError)]

pub enum EngineError {

    #[error("the Matrix client is not authenticated")]

    Unauthenticated,

    #[error("command, event, and media source capacities must be nonzero")]

    ZeroCapacity,

    #[error("could not build the sync service")]

    SyncService(#[source] matrix_sdk_ui::sync_service::Error),

    #[error("the event receiver has closed")]

    EventReceiverClosed,

    #[error("engine task {task} panicked")]

    TaskPanicked { task: String },

}

impl EngineError {

    fn failure(&self) -> EngineFailure {

        EngineFailure {

            kind: match self {

                Self::EventReceiverClosed => EngineErrorKind::EventReceiverClosed,

                Self::TaskPanicked { .. } => EngineErrorKind::Watcher,

                Self::Unauthenticated | Self::ZeroCapacity | Self::SyncService(_) => {

                    EngineErrorKind::Other

                }

            },

            message: self.to_string(),

        }

    }

}

#[derive(Clone)]

enum TaskContext {

    /// a panic here costs the caller an `Interrupted` answer and nothing else

    Command(&'static str),

    /// a panic here takes the engine with it

    Watcher(&'static str),

    RoomActor {

        room_id: OwnedRoomId,

        generation: u64,

    },

}

impl TaskContext {

    fn label(&self) -> String {

        match self {

            Self::Command(name) | Self::Watcher(name) => (*name).to_owned(),

            Self::RoomActor {

                room_id,

                generation,

            } => format!("room actor {room_id} generation {generation}"),

        }

    }

}

/// tasks catch their own panics, since the sdk's wasm join error drops the payload

fn panic_message(payload: Box<dyn std::any::Any + Send>) -> String {

    payload

        .downcast_ref::<&str>()

        .map(ToString::to_string)

        .or_else(|| payload.downcast_ref::<String>().cloned())

        .unwrap_or_else(|| "task panicked with a non-string payload".to_owned())

}

struct TaskRegistry {

    tasks: TaskSet<(TaskContext, Option<String>)>,

}

impl TaskRegistry {

    fn spawn(

        &mut self,

        context: TaskContext,

        task: impl Future<Output = ()> + SendOutsideWasm + 'static,

    ) {

        self.tasks.spawn(async move {

            let panic = AssertUnwindSafe(task)

                .catch_unwind()

                .await

                .err()

                .map(panic_message);

            (context, panic)

        });

    }

    fn is_empty(&self) -> bool {

        self.tasks.is_empty()

    }

    async fn join_next(&mut self) -> Option<(TaskContext, Option<String>)> {

        loop {

            if let Ok(done) = self.tasks.join_next().await? {

                return Some(done);

            }

        }

    }

    async fn shutdown(&mut self) {

        self.tasks.shutdown().await;

    }

}

enum RoomSlot {

    Active {

        generation: u64,

        sender: mpsc::Sender<ActorRequest>,

    },

    /// opens and closes that arrive while the actor winds down, replayed in order once it's gone

    Closing {

        generation: u64,

        pending: VecDeque<Deferred>,

    },

}

impl RoomSlot {

    const fn generation(&self) -> u64 {

        match self {

            Self::Active { generation, .. } | Self::Closing { generation, .. } => *generation,

        }

    }

}

enum Deferred {

    Open(Responder),

    Close(Responder),

}

struct Queues {

    mutations: mpsc::Receiver<MutationJob>,

    notifications: mpsc::Receiver<Job<NotificationCommand>>,

    encryption: mpsc::Receiver<Job<EncryptionCommand>>,

    session_changes: broadcast::Receiver<SessionChange>,

}

/// requests other devices send us, handed from the sdk's event handler to the encryption worker

type IncomingVerifications =

    mpsc::UnboundedSender<matrix_sdk::encryption::verification::VerificationRequest>;

pub struct Engine {

    client: Client,

    sync: Arc<SyncService>,

    port: DriverPort,

    config: EngineConfig,

    media: Arc<MediaService>,

    tasks: TaskRegistry,

    rooms: HashMap<OwnedRoomId, RoomSlot>,

    next_generation: u64,

    notes: mpsc::UnboundedReceiver<ActorNote>,

    note_sender: mpsc::UnboundedSender<ActorNote>,

    /// the rooms sliding sync should follow, applied off the run loop since it touches the store

    subscriptions: watch::Sender<Vec<OwnedRoomId>>,

    mutations: mpsc::Sender<MutationJob>,

    notifications: mpsc::Sender<Job<NotificationCommand>>,

    encryption: mpsc::Sender<Job<EncryptionCommand>>,

    queues: Option<Queues>,

    /// one progress forwarder per room that has uploaded, and whether it's subscribed yet

    uploads: HashMap<OwnedRoomId, watch::Receiver<bool>>,

    /// the client can outlive us, so the handler has to come off it when we stop

    verification_handler: Option<matrix_sdk::event_handler::EventHandlerHandle>,

    stop: mpsc::UnboundedSender<Termination>,

    stops: mpsc::UnboundedReceiver<Termination>,

}

impl Engine {

    /// the escape hatch for a client you built yourself. `Session` is the usual way in

    pub async fn from_client(

        client: Client,

        config: EngineConfig,

    ) -> Result<(Self, ClientHandle, EventStream), EngineError> {

        if client.user_id().is_none() {

            return Err(EngineError::Unauthenticated);

        }

        if config.command_capacity == 0

            || config.event_capacity == 0

            || config.max_media_sources == 0

        {

            return Err(EngineError::ZeroCapacity);

        }

        let session_changes = client.subscribe_to_session_changes();

        let sync = Arc::new(

            SyncService::builder(client.clone())

                .with_offline_mode()

                .build()

                .await

                .map_err(EngineError::SyncService)?,

        );

        let media = Arc::new(MediaService::new(

            client.clone(),

            config.max_media_bytes,

            config.max_media_sources,

        ));

        let (handle, port, events) = transport_pair(

            client.clone(),

            sync.clone(),

            media.clone(),

            config.command_capacity,

            config.event_capacity,

        );

        client.send_queue().enable_upload_progress(true);

        let (note_sender, notes) = mpsc::unbounded_channel();

        let (stop, stops) = mpsc::unbounded_channel();

        let (mutations, mutation_queue) = mpsc::channel(config.command_capacity);

        let (notifications, notification_queue) = mpsc::channel(config.command_capacity);

        let (encryption, encryption_queue) = mpsc::channel(config.command_capacity);

        let engine = Self {

            client,

            sync,

            port,

            config,

            media,

            tasks: TaskRegistry {

                tasks: TaskSet::new(),

            },

            rooms: HashMap::new(),

            next_generation: 0,

            notes,

            note_sender,

            subscriptions: watch::Sender::new(Vec::new()),

            mutations,

            notifications,

            encryption,

            queues: Some(Queues {

                mutations: mutation_queue,

                notifications: notification_queue,

                encryption: encryption_queue,

                session_changes,

            }),

            uploads: HashMap::new(),

            verification_handler: None,

            stop,

            stops,

        };

        Ok((engine, handle, events))

    }

    /// boxed because the engine's future is deep enough to blow the compiler's query depth

    /// in every crate that spawns it, which isn't something a client should have to fix

    pub fn run(self) -> EngineBoxFuture<'static, Result<Termination, EngineError>> {

        Box::pin(self.run_inner())

    }

    async fn run_inner(mut self) -> Result<Termination, EngineError> {

        let events = self.port.event_sender();

        let user_id = self

            .client

            .user_id()

            .expect("authentication was checked during construction")

            .to_owned();

        self.port

            .connection_sender()

            .send_replace(ConnectionState::Connecting);

        for event in [

            Event::Authenticated { user_id },

            Event::ConnectionChanged(ConnectionState::Connecting),

        ] {

            if events.send(event).await.is_err() {

                self.shutdown().await;

                return Err(EngineError::EventReceiverClosed);

            }

        }

        self.start_tasks();

        self.resume_upload_forwarders().await;

        self.sync.start().await;

        let mut shutdown_reply = None;

        let ending = loop {

            tokio::select! {

                biased;

                () = events.closed() => break Err(EngineError::EventReceiverClosed),

                Some(termination) = self.stops.recv() => break Ok(termination),

                Some((context, panic)) = self.tasks.join_next(), if !self.tasks.is_empty() => {

                    let Some(message) = panic else { continue };

                    if let TaskContext::Command(label) = context {

                        tracing::error!(label, panic = %message, "command task panicked");

                        continue;

                    }

                    break Err(EngineError::TaskPanicked { task: context.label() });

                }

                Some(note) = self.notes.recv() => self.actor_note(note).await,

                request = self.port.recv() => {

                    let Some(request) = request else { break Ok(Termination::HandlesDropped) };

                    if let Some(reply) = self.dispatch(request).await {

                        shutdown_reply = Some(reply);

                        break Ok(Termination::ShutDown);

                    }

                }

            }

        };

        if let Err(error) = &ending {

            let _ = events.send(Event::Failed(error.failure())).await;

        }

        self.shutdown().await;

        let _ = events.send(Event::Stopped).await;

        if let Some(reply) = shutdown_reply {

            reply.done(Ok(()));

        }

        // everything else still queued answers `Interrupted` as it drops, but a shutdown did happen

        while let Some(Request { command, reply }) = self.port.try_recv() {

            if matches!(command, Command::Lifecycle(LifecycleCommand::Shutdown)) {

                reply.done(Ok(()));

            }

        }

        ending

    }

    fn start_tasks(&mut self) {

        let queues = self.queues.take().expect("engine tasks start once");

        let (incoming, incoming_verifications) = mpsc::unbounded_channel();

        self.verification_handler = Some(self.verification_handler(incoming));

        let client = &self.client;

        let events = || self.port.event_sender();

        let watchers: [(&'static str, EngineBoxFuture<'static, ()>); 11] = [

            (

                "upload watcher",

                Box::pin(watchers::uploads(

                    client.clone(),

                    client.send_queue().subscribe(),

                    events(),

                )),

            ),

            (

                "room subscription watcher",

                Box::pin(watchers::room_subscriptions(

                    self.sync.clone(),

                    self.subscriptions.subscribe(),

                )),

            ),

            (

                "session change watcher",

                Box::pin(watchers::session(

                    queues.session_changes,

                    events(),

                    client

                        .device_id()

                        .expect("authentication was checked during construction")

                        .to_owned(),

                    self.stop.clone(),

                )),

            ),

            (

                "connection state watcher",

                Box::pin(watchers::connection(

                    self.sync.state(),

                    self.sync.clone(),

                    client.clone(),

                    events(),

                    self.port.connection_sender(),

                )),

            ),

            (

                "send queue watcher",

                Box::pin(watchers::send_queue_errors(

                    client.clone(),

                    client.send_queue().subscribe_errors(),

                    self.port.connection_sender().subscribe(),

                )),

            ),

            (

                "encryption state watcher",

                Box::pin(watchers::encryption(client.clone(), events())),

            ),

            (

                "room list watcher",

                Box::pin(watchers::rooms(client.clone(), events())),

            ),

            (

                "notification settings watcher",

                Box::pin(watchers::notification_settings(client.clone(), events())),

            ),

            (

                "room command worker",

                Box::pin(mutations::run_scheduler(

                    client.clone(),

                    queues.mutations,

                    self.config.command_capacity,

                )),

            ),

            (

                "notification command worker",

                Box::pin(commands::run_notifications(

                    client.clone(),

                    queues.notifications,

                )),

            ),

            (

                "encryption command worker",

                Box::pin(commands::run_encryption(

                    client.clone(),

                    events(),

                    queues.encryption,

                    incoming_verifications,

                )),

            ),

        ];

        for (label, task) in watchers {

            self.tasks.spawn(TaskContext::Watcher(label), task);

        }

    }

    fn verification_handler(

        &self,

        incoming: IncomingVerifications,

    ) -> matrix_sdk::event_handler::EventHandlerHandle {

        use matrix_sdk::ruma::events::key::verification::request::ToDeviceKeyVerificationRequestEvent;

        self.client.add_event_handler(

            move |event: ToDeviceKeyVerificationRequestEvent, client: Client| {

                let incoming = incoming.clone();