macltok

use std::{

    borrow::Cow,

    collections::{HashMap, HashSet, VecDeque},

    future::Future,

    io::Read as _,

    panic::AssertUnwindSafe,

    sync::Arc,

    time::Duration,

};

use futures_util::{FutureExt, StreamExt};

use matrix_sdk::{

    Client, Room, SessionChange,

    media::{MediaFormat, MediaRequestParameters, MediaThumbnailSettings},

    room::edit::EditedContent,

    room::reply::{EnforceThread, Reply as SdkReply},

    ruma::{

        OwnedDeviceId, OwnedEventId, OwnedRoomId, RoomId,

        api::client::receipt::create_receipt::v3::ReceiptType,

        events::{

            AnyGlobalAccountDataEventContent, AnyMessageLikeEventContent,

            GlobalAccountDataEventType, Mentions, RoomAccountDataEventType,

            receipt::ReceiptThread,

            room::message::{

                AddMentions, MessageType, Relation, ReplyWithinThread,

                RoomMessageEventContentWithoutRelation,

            },

        },

        serde::Raw,

    },

    send_queue::{LocalEchoContent, RoomSendQueueUpdate, SendHandle},

};

use matrix_sdk_ui::{

    sync_service::{State as SyncState, SyncService},

    timeline::{

        RoomExt, Timeline, TimelineEventItemId, TimelineFocus, TimelineReadReceiptTracking,

    },

};

use thiserror::Error;

use tokio::{

    sync::{Mutex, mpsc, oneshot, watch},

    task::{AbortHandle, Id as TaskId, JoinHandle, JoinSet},

};

use crate::{

    AccountCommand, AttachmentUpload, ClientHandle, Command, CommandErrorKind, CommandFailure,

    CommandId, CommandKind, ConnectionState, CreateRoomOptions, EncryptionCommand, EngineErrorKind,

    EngineFailure, Event, EventStream, HistoryState, InteractiveAuth, InteractiveAuthChallenge,

    InteractiveAuthOutcome, LifecycleCommand, MediaCommand, MediaId, MediaSize, MessageContent,

    MessageEditContent, MessageMentions, NotificationCommand, PowerLevelChanges, RichText,

    RoomCommand, RoomNotificationMode, RoomPreset, RoomVisibility, TimelineCommand, TimelineId,

    TypingState,

    encryption::{VerificationController, encryption_state},

    error::{CommandError, classify_sdk_error},

    normalise::{MediaRegistration, MediaScope, normalise_timeline, normalise_timeline_diffs},

    transport::{

        ClientQueryError, CommandEnvelope, CommandResultError, DriverPort, DriverRequest,

        EventSender, SendEventError, transport_pair,

    },

    upload::prepare_upload,

};

use matrix_sdk::ruma::api::{Metadata, OutgoingRequest, auth_scheme::SendAccessToken};

use matrix_sdk_crypto::{AttachmentDecryptor, MediaEncryptionInfo};

#[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, Error)]

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::Unauthenticated

                | Self::ZeroCapacity

                | Self::SyncService(_)

                | Self::TaskPanicked { .. } => EngineErrorKind::Other,

            },

            message: self.to_string(),

        }

    }

}

type UploadWatchers = HashMap<

    OwnedRoomId,

    (

        AbortHandle,

        watch::Receiver<Option<Result<(), CommandErrorKind>>>,

    ),

>;

const SYNC_RESTART_MIN: Duration = Duration::from_secs(1);

const SYNC_RESTART_MAX: Duration = Duration::from_secs(60);

const SEND_QUEUE_RETRY: Duration = Duration::from_secs(5);

type MediaFlights = Arc<Mutex<HashMap<(MediaId, MediaSize), Vec<CommandId>>>>;

pub struct Engine {

    client: Client,

    sync: Arc<SyncService>,

    port: DriverPort,

    rooms: HashMap<OwnedRoomId, RoomActorState>,

    next_room_generation: u64,

    room_actor_events: mpsc::UnboundedReceiver<RoomActorEvent>,

    room_actor_event_sender: mpsc::UnboundedSender<RoomActorEvent>,

    upload_watchers: UploadWatchers,

    tasks: TaskRegistry,

    session_changes: Option<tokio::sync::broadcast::Receiver<SessionChange>>,

    initial_history_limit: u16,

    room_command_capacity: usize,

    max_media_bytes: usize,

    media: Arc<Mutex<MediaRegistry>>,

    notification_commands: mpsc::Sender<CommandJob<NotificationCommand>>,

    notification_command_receiver: Option<mpsc::Receiver<CommandJob<NotificationCommand>>>,

    encryption_commands: mpsc::Sender<EncryptionWork>,

    encryption_command_receiver: Option<mpsc::Receiver<EncryptionWork>>,

    verification: Arc<Mutex<VerificationController>>,

    room_commands: mpsc::Sender<RoomJob>,

    room_command_receiver: Option<mpsc::Receiver<RoomJob>>,

    global_account_gates: HashMap<String, oneshot::Receiver<()>>,

    room_account_gates: HashMap<(OwnedRoomId, String), oneshot::Receiver<()>>,

    media_flights: MediaFlights,

    stop: Arc<tokio::sync::Notify>,

}

struct CommandJob<T> {

    id: CommandId,

    command: T,

}

enum EncryptionWork {

    Command(CommandJob<EncryptionCommand>),

}

struct RoomJob {

    id: CommandId,

    command: RoomCommand,

    reply: Option<oneshot::Sender<Result<OwnedRoomId, CommandResultError>>>,

}

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

enum RoomMutationKey {

    Room(OwnedRoomId),

    Target(String),

    Direct(ruma::OwnedUserId),

    Independent(CommandId),

}

struct RoomJobCompletion {

    key: RoomMutationKey,

    id: CommandId,

    reply: Option<oneshot::Sender<Result<OwnedRoomId, CommandResultError>>>,

    result: Result<Option<OwnedRoomId>, CommandError>,

}

#[derive(Clone)]

enum TaskContext {

    Command {

        id: CommandId,

        kind: CommandKind,

    },

    Watcher(&'static str),

    RoomActor {

        room_id: OwnedRoomId,

        generation: u64,

    },

}

impl TaskContext {

    fn label(&self) -> String {

        match self {

            Self::Command { id, kind } => format!("{kind:?} command {}", id.get()),

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

            Self::RoomActor {

                room_id,

                generation,

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

        }

    }

}

struct TaskRegistry {

    tasks: JoinSet<()>,

    contexts: HashMap<TaskId, TaskContext>,

}

impl TaskRegistry {

    fn new() -> Self {

        Self {

            tasks: JoinSet::new(),

            contexts: HashMap::new(),

        }

    }

    fn is_empty(&self) -> bool {

        self.tasks.is_empty()

    }

    fn spawn(

        &mut self,

        context: TaskContext,

        future: impl Future<Output = ()> + Send + 'static,

    ) -> AbortHandle {

        let handle = self.tasks.spawn(future);

        self.contexts.insert(handle.id(), context);

        handle

    }

    async fn join_next(&mut self) -> Option<(TaskContext, Result<(), tokio::task::JoinError>)> {

        let result = self.tasks.join_next_with_id().await?;

        let id = match &result {

            Ok((id, ())) => *id,

            Err(error) => error.id(),

        };

        let context = self

            .contexts

            .remove(&id)

            .expect("every registered task has context");

        Some((context, result.map(|(_, ())| ())))

    }

    async fn shutdown(&mut self) {

        self.tasks.abort_all();

        while self.tasks.join_next().await.is_some() {}

        self.contexts.clear();

    }

}

enum RoomActorState {

    Opening {

        generation: u64,

        sender: mpsc::Sender<RoomActorRequest>,

    },

    Open {

        generation: u64,

        sender: mpsc::Sender<RoomActorRequest>,

    },

    Closing {

        generation: u64,

        pending: VecDeque<DeferredRoomLifecycle>,

    },

}

enum DeferredRoomLifecycle {

    Open(CommandId),

    Close(CommandId),

}

impl RoomActorState {

    const fn generation(&self) -> u64 {

        match self {

            Self::Opening { generation, .. }

            | Self::Open { generation, .. }

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

        }

    }

    fn sender(&self) -> Option<&mpsc::Sender<RoomActorRequest>> {

        match self {

            Self::Opening { sender, .. } | Self::Open { sender, .. } => Some(sender),

            Self::Closing { .. } => None,

        }

    }

}

enum RoomActorEvent {

    Ready {

        room_id: OwnedRoomId,

        generation: u64,

    },

    Exited {

        room_id: OwnedRoomId,

        generation: u64,

    },

}

enum EngineActivity {

    Request(Option<Box<DriverRequest>>),

    Task(Option<(TaskContext, Result<(), tokio::task::JoinError>)>),

    RoomActor(Option<RoomActorEvent>),

}

/// least recently seen goes first. stale order slots are skipped by stamp rather than hunted down

struct MediaRegistry {

    entries: HashMap<MediaId, (MediaRegistration, u64)>,

    order: VecDeque<(MediaId, u64)>,

    next_stamp: u64,

    limit: usize,

}

impl MediaRegistry {

    fn new(limit: usize) -> Self {

        Self {

            entries: HashMap::new(),

            order: VecDeque::new(),

            next_stamp: 0,

            limit,

        }

    }

    fn get(&self, id: &MediaId) -> Option<&MediaRegistration> {

        self.entries.get(id).map(|(media, _)| media)

    }

    fn insert_all(&mut self, media: Vec<MediaRegistration>) {

        for media in media {

            let stamp = self.next_stamp;

            self.next_stamp += 1;

            self.order.push_back((media.id.clone(), stamp));

            self.entries.insert(media.id.clone(), (media, stamp));

        }

        while self.entries.len() > self.limit {

            let Some((id, stamp)) = self.order.pop_front() else {

                break;

            };

            if self

                .entries

                .get(&id)

                .is_some_and(|(_, current)| *current == stamp)

            {

                self.entries.remove(&id);

            }

        }

        if self.order.len() > self.limit.saturating_mul(4) {

            let entries = &self.entries;

            self.order

                .retain(|(id, stamp)| entries.get(id).is_some_and(|(_, current)| current == stamp));

        }

    }

    fn release(&mut self, scope: &MediaScope) {

        self.entries.retain(|id, _| !scope.owns(id));

        let entries = &self.entries;

        self.order.retain(|(id, _)| entries.contains_key(id));

    }

}

enum RoomActorRequest {

    Opened {

        id: CommandId,

    },

    Execute {

        id: CommandId,

        command: TimelineCommand,

    },

    Close {

        id: CommandId,

    },

}

struct RoomActor {

    room_id: OwnedRoomId,

    generation: u64,

    room: Room,

    receiver: mpsc::Receiver<RoomActorRequest>,

    events: EventSender,

    engine_events: mpsc::UnboundedSender<RoomActorEvent>,

    media: Arc<Mutex<MediaRegistry>>,

    initial_history_limit: u16,

    open_ids: Vec<CommandId>,

    background: JoinSet<(CommandId, Result<(), CommandError>)>,

    background_ids: HashMap<TaskId, CommandId>,

}

type BackgroundResult =

    Result<(TaskId, (CommandId, Result<(), CommandError>)), tokio::task::JoinError>;

type BackgroundWork = std::pin::Pin<Box<dyn Future<Output = Result<(), CommandError>> + Send>>;

/// network-bound work runs beside the actor so one slow redaction doesn't queue the room behind it

enum TimelineOutcome {

    Done(Result<(), CommandError>),

    Background(BackgroundWork),

}

struct TimelineSession {

    id: TimelineId,

    timeline: Arc<Timeline>,

    events: EventSender,

    watcher: JoinHandle<()>,

    media: Arc<Mutex<MediaRegistry>>,

    current_user_id: Option<ruma::OwnedUserId>,

    scope: MediaScope,

    pagination_ids: Arc<Mutex<HashSet<CommandId>>>,

    pagination_sender: mpsc::Sender<PaginationRequest>,

    pagination_worker: JoinHandle<()>,

    pagination_results: Option<mpsc::UnboundedSender<(CommandId, Result<(), CommandError>)>>,

    pagination_result_task: JoinHandle<()>,

}

struct PaginationRequest {

    timeline: Arc<Timeline>,

    events: EventSender,

    timeline_id: TimelineId,

    limit: u16,

    command_id: Option<CommandId>,

    results: mpsc::UnboundedSender<(CommandId, Result<(), CommandError>)>,

}

impl TimelineSession {

    async fn start(

        id: TimelineId,

        timeline: Timeline,

        events: EventSender,

        media: Arc<Mutex<MediaRegistry>>,

        current_user_id: Option<ruma::OwnedUserId>,

        initial_history_limit: u16,

    ) -> Result<Self, SendEventError> {

        let timeline = Arc::new(timeline);

        let scope = MediaScope::new();

        let watcher = start_timeline_watcher(

            id.clone(),

            timeline.clone(),

            events.clone(),

            media.clone(),

            current_user_id.clone(),

            scope.clone(),

        )

        .await?;

        let (pagination_sender, mut pagination_receiver) = mpsc::channel::<PaginationRequest>(32);

        let pagination_worker = tokio::spawn(async move {

            while let Some(request) = pagination_receiver.recv().await {

                let command_id = request.command_id;

                let results = request.results.clone();

                if AssertUnwindSafe(run_pagination(request))

                    .catch_unwind()

                    .await

                    .is_err()

                {