continuwuity/src/api/client/sync/v3/joined.rs

use std::collections::{BTreeMap, HashMap};

use conduwuit::{
	Result, at, err, extract_variant, is_equal_to,
	matrix::{
		Event,
		pdu::{PduCount, PduEvent},
	},
	ref_at,
	result::FlatOk,
	utils::{
		BoolExt, IterStream, ReadyExt, TryFutureExtExt,
		future::OptionStream,
		math::ruma_from_u64,
		stream::{BroadbandExt, Tools, WidebandExt},
	},
};
use conduwuit_service::{
	Services,
	rooms::{
		lazy_loading,
		lazy_loading::{Options, Witness},
		short::ShortStateHash,
	},
};
use futures::{
	FutureExt, StreamExt, TryFutureExt,
	future::{OptionFuture, join, join3, join4, try_join4},
};
use ruma::{
	OwnedEventId, OwnedRoomId, OwnedUserId, RoomId, UserId,
	api::client::sync::sync_events::{
		UnreadNotificationsCount,
		v3::{Ephemeral, JoinedRoom, RoomAccountData, RoomSummary, State as RoomState, Timeline},
	},
	events::{
		AnyRawAccountDataEvent, AnySyncEphemeralRoomEvent, StateEventType,
		TimelineEventType::*,
		room::member::{MembershipState, RoomMemberEventContent},
	},
	serde::Raw,
	uint,
};
use service::rooms::short::{ShortEventId, ShortStateKey};
use tracing::trace;

use super::{load_timeline, share_encrypted_room};
use crate::client::{
	TimelinePdus, ignored_filter,
	sync::v3::{DeviceListUpdates, SyncContext},
};

#[tracing::instrument(
	name = "joined",
	level = "debug",
	skip_all,
	fields(
		room_id = ?room_id,
	),
)]
#[allow(clippy::too_many_arguments)]
pub(super) async fn load_joined_room(
	services: &Services,
	SyncContext {
		sender_user,
		sender_device,
		since,
		next_batch,
		full_state,
		filter,
	}: SyncContext<'_>,
	ref room_id: OwnedRoomId,
) -> Result<(JoinedRoom, DeviceListUpdates)> {
	let mut device_list_updates = DeviceListUpdates::new();
	let sincecount = since.map(PduCount::Normal);
	let next_batchcount = PduCount::Normal(next_batch);

	// the shortstatehash of the room's state right now
	let current_shortstatehash = services
		.rooms
		.state
		.get_room_shortstatehash(room_id)
		.map_err(|_| err!(Database(error!("Room {room_id} has no state"))));

	// the shortstatehash of what the room's state was when the `since` token was
	// issued
	let since_shortstatehash = OptionFuture::from(since.map(|since| {
		services
			.rooms
			.user
			.get_token_shortstatehash(room_id, since)
			.ok()
	}))
	.map(|v| Ok(v.flatten()));

	let timeline = load_timeline(
		services,
		sender_user,
		room_id,
		sincecount,
		Some(next_batchcount),
		10_usize,
	);

	let receipt_events = services
		.rooms
		.read_receipt
		.readreceipts_since(room_id, since)
		.filter_map(|(read_user, _, edu)| async move {
			services
				.users
				.user_is_ignored(read_user, sender_user)
				.await
				.or_some((read_user.to_owned(), edu))
		})
		.collect::<HashMap<OwnedUserId, Raw<AnySyncEphemeralRoomEvent>>>()
		.map(Ok);

	let (current_shortstatehash, since_shortstatehash, timeline, receipt_events) =
		try_join4(current_shortstatehash, since_shortstatehash, timeline, receipt_events)
			.boxed()
			.await?;

	let TimelinePdus { pdus: timeline_pdus, limited } = timeline;
	let is_initial_sync = since_shortstatehash.is_none();

	let timeline_start_shortstatehash = async {
		if let Some((_, pdu)) = timeline_pdus.first() {
			if let Ok(shortstatehash) = services
				.rooms
				.state_accessor
				.pdu_shortstatehash(&pdu.event_id)
				.await
			{
				return shortstatehash;
			}
		}

		current_shortstatehash
	};

	let last_notification_read: OptionFuture<_> = timeline_pdus
		.is_empty()
		.then(|| {
			services
				.rooms
				.user
				.last_notification_read(sender_user, room_id)
		})
		.into();

	let since_sender_member: OptionFuture<_> = since_shortstatehash
		.map(|short| {
			services
				.rooms
				.state_accessor
				.state_get_content(short, &StateEventType::RoomMember, sender_user.as_str())
				.ok()
		})
		.into();

	let is_encrypted_room = services
		.rooms
		.state_accessor
		.state_get(current_shortstatehash, &StateEventType::RoomEncryption, "")
		.is_ok();

	let (
		last_notification_read,
		since_sender_member,
		timeline_start_shortstatehash,
		is_encrypted_room,
	) = join4(
		last_notification_read,
		since_sender_member,
		timeline_start_shortstatehash,
		is_encrypted_room,
	)
	.await;

	let joined_since_last_sync =
		since_sender_member
			.flatten()
			.is_none_or(|content: RoomMemberEventContent| {
				content.membership != MembershipState::Join
			});

	let lazy_loading_enabled = (filter.room.state.lazy_load_options.is_enabled()
		|| filter.room.timeline.lazy_load_options.is_enabled())
		&& !full_state;

	let lazy_loading_context = &lazy_loading::Context {
		user_id: sender_user,
		device_id: Some(sender_device),
		room_id,
		token: since,
		options: Some(&filter.room.state.lazy_load_options),
	};

	let lazy_loading_witness = OptionFuture::from(lazy_loading_enabled.then(|| {
		let witness: Witness = timeline_pdus
			.iter()
			.map(ref_at!(1))
			.map(Event::sender)
			.map(Into::into)
			.chain(receipt_events.keys().map(Into::into))
			.collect();

		services
			.rooms
			.lazy_loading
			.witness_retain(witness, lazy_loading_context)
	}))
	.await;

	/* replace with multiple steps
		glossary for my own sanity:
		- full state: every state event from the start of the room to the start of the timeline
		- incremental state: state events from `since` to the start of the timeline
		- state_events: the `state` key on the JSON object we return

		if initial sync or full_state:
			get full state
			use full state as state_events
		else if TL is limited:
			get incremental state
			use incremental state as state_events
			if encryption is enabled:
				use incremental state to extend device list
		else:
			state_events is empty
		compute counts and heroes from state_events
	*/
	let mut state_events = if is_initial_sync || full_state {
		// reset lazy loading state on initial sync
		if is_initial_sync {
			services
				.rooms
				.lazy_loading
				.reset(lazy_loading_context)
				.await;
		}

		calculate_state_initial(
			services,
			sender_user,
			timeline_start_shortstatehash,
			lazy_loading_witness.as_ref(),
		)
		.boxed()
		.await?
	} else if limited {
		let state_incremental = calculate_state_incremental(
			services,
			sender_user,
			since_shortstatehash,
			timeline_start_shortstatehash,
			lazy_loading_witness.as_ref(),
		)
		.boxed()
		.await?;

		// calculate device list updates for E2EE
		if is_encrypted_room {
			// add users with changed keys to the `changed` list
			services
				.users
				.room_keys_changed(room_id, since, Some(next_batch))
				.map(at!(0))
				.map(ToOwned::to_owned)
				.ready_for_each(|user_id| {
					device_list_updates.changed.insert(user_id);
				})
				.await;

			// add users who now share encrypted rooms to `changed` and
			// users who no longer share encrypted rooms to `left`
			for state_event in &state_incremental {
				if state_event.kind == RoomMember {
					let Some(content): Option<RoomMemberEventContent> =
						state_event.get_content().ok()
					else {
						continue;
					};

					let Some(user_id): Option<OwnedUserId> = state_event
						.state_key
						.as_ref()
						.and_then(|key| key.parse().ok())
					else {
						continue;
					};

					{
						use MembershipState::*;

						if matches!(content.membership, Leave | Join) {
							let shares_encrypted_room = share_encrypted_room(
								services,
								sender_user,
								&user_id,
								Some(room_id),
							)
							.await;
							match content.membership {
								| Leave if !shares_encrypted_room => {
									device_list_updates.left.insert(user_id);
								},
								| Join if joined_since_last_sync || shares_encrypted_room => {
									device_list_updates.changed.insert(user_id);
								},
								| _ => (),
							}
						}
					}
				}
			}
		}

		state_incremental
	} else {
		vec![]
	};

	// only compute room counts and heroes (aka the summary) if the room's members
	// changed since the last sync
	let (joined_member_count, invited_member_count, heroes) =
		if state_events.iter().any(|event| event.kind == RoomMember) {
			calculate_counts(services, room_id, sender_user).await?
		} else {
			(None, None, None)
		};

	let is_sender_membership = |pdu: &PduEvent| {
		pdu.kind == StateEventType::RoomMember.into()
			&& pdu
				.state_key
				.as_deref()
				.is_some_and(is_equal_to!(sender_user.as_str()))
	};

	let joined_sender_member: Option<_> = (joined_since_last_sync && timeline_pdus.is_empty())
		.then(|| {
			state_events
				.iter()
				.position(is_sender_membership)
				.map(|pos| state_events.swap_remove(pos))
		})
		.flatten();

	let prev_batch = timeline_pdus
		.first()
		.map(at!(0))
		.or_else(|| joined_sender_member.is_some().and(since).map(Into::into));

	let timeline_pdus = timeline_pdus
		.into_iter()
		.stream()
		.wide_filter_map(|item| ignored_filter(services, item, sender_user))
		.map(at!(1))
		.chain(joined_sender_member.into_iter().stream())
		.map(Event::into_format)
		.collect::<Vec<_>>();

	let account_data_events = services
		.account_data
		.changes_since(Some(room_id), sender_user, since, Some(next_batch))
		.ready_filter_map(|e| extract_variant!(e, AnyRawAccountDataEvent::Room))
		.collect();

	let send_notification_counts =
		last_notification_read.is_none_or(|count| since.is_none_or(|since| count > since));

	let notification_count: OptionFuture<_> = send_notification_counts
		.then(|| {
			services
				.rooms
				.user
				.notification_count(sender_user, room_id)
				.map(TryInto::try_into)
				.unwrap_or(uint!(0))
		})
		.into();

	let highlight_count: OptionFuture<_> = send_notification_counts
		.then(|| {
			services
				.rooms
				.user
				.highlight_count(sender_user, room_id)
				.map(TryInto::try_into)
				.unwrap_or(uint!(0))
		})
		.into();

	let typing_events = services
		.rooms
		.typing
		.last_typing_update(room_id)
		.and_then(|count| async move {
			if since.is_some_and(|since| count <= since) {
				return Ok(Vec::<Raw<AnySyncEphemeralRoomEvent>>::new());
			}

			let typings = services
				.rooms
				.typing
				.typings_all(room_id, sender_user)
				.await?;

			Ok(vec![serde_json::from_str(&serde_json::to_string(&typings)?)?])
		})
		.unwrap_or(Vec::new());

	let unread_notifications = join(notification_count, highlight_count);
	let events = join3(timeline_pdus, account_data_events, typing_events);
	let (unread_notifications, events) = join(unread_notifications, events).boxed().await;

	let (room_events, account_data_events, typing_events) = events;
	let (notification_count, highlight_count) = unread_notifications;

	let last_privateread_update = if let Some(since) = since {
		services
			.rooms
			.read_receipt
			.last_privateread_update(sender_user, room_id)
			.await > since
	} else {
		true
	};

	let private_read_event = if last_privateread_update {
		services
			.rooms
			.read_receipt
			.private_read_get(room_id, sender_user)
			.await
			.ok()
	} else {
		None
	};

	let edus: Vec<Raw<AnySyncEphemeralRoomEvent>> = receipt_events
		.into_values()
		.chain(typing_events.into_iter())
		.chain(private_read_event.into_iter())
		.collect();

	// Save the state after this sync so we can send the correct state diff next
	// sync
	services
		.rooms
		.user
		.associate_token_shortstatehash(room_id, next_batch, current_shortstatehash)
		.await;

	let joined_room = JoinedRoom {
		account_data: RoomAccountData { events: account_data_events },
		summary: RoomSummary {
			joined_member_count: joined_member_count.map(ruma_from_u64),
			invited_member_count: invited_member_count.map(ruma_from_u64),
			heroes: heroes
				.into_iter()
				.flatten()
				.map(TryInto::try_into)
				.filter_map(Result::ok)
				.collect(),
		},
		unread_notifications: UnreadNotificationsCount { highlight_count, notification_count },
		timeline: Timeline {
			limited,
			prev_batch: prev_batch.as_ref().map(ToString::to_string),
			events: room_events,
		},
		state: RoomState {
			events: state_events.into_iter().map(Event::into_format).collect(),
		},
		ephemeral: Ephemeral { events: edus },
		unread_thread_notifications: BTreeMap::new(),
	};

	Ok((joined_room, device_list_updates))
}

/// Calculate the "initial state", or all events from the start of the room up
/// to (but not including) the `current_shortstatehash`. If
/// `lazy_loading_witness` is `None`, lazy loading will be disabled.
#[tracing::instrument(
	name = "initial",
	level = "trace",
	skip_all,
	fields(current_shortstatehash)
)]
#[allow(clippy::too_many_arguments)]
async fn calculate_state_initial(
	services: &Services,
	sender_user: &UserId,
	current_shortstatehash: ShortStateHash,
	lazy_loading_witness: Option<&Witness>,
) -> Result<Vec<PduEvent>> {
	let (shortstatekeys, event_ids): (Vec<_>, Vec<_>) = services
		.rooms
		.state_accessor
		.state_full_ids(current_shortstatehash)
		.unzip()
		.await;

	trace!("event ids for initial sync @ {:?}: {:?}", current_shortstatehash, event_ids);

	services
		.rooms
		.short
		.multi_get_statekey_from_short(shortstatekeys.into_iter().stream())
		.zip(event_ids.into_iter().stream())
		.ready_filter_map(|item| Some((item.0.ok()?, item.1)))
		.ready_filter_map(|((event_type, state_key), event_id)| {
			let lazy = lazy_loading_witness.is_some_and(|witness| {
				event_type == StateEventType::RoomMember
					&& state_key.as_str().try_into().is_ok_and(|user_id: &UserId| {
						sender_user != user_id && !witness.contains(user_id)
					})
			});

			lazy.or_some(event_id)
		})
		.broad_filter_map(|event_id: OwnedEventId| async move {
			services.rooms.timeline.get_pdu(&event_id).await.ok()
		})
		.collect()
		.map(Ok)
		.await
}

/// Calculate the "incremental state", or all events from the
/// `since_shortstatehash` up to (but not including)
/// the `current_shortstatehash`. If `lazy_loading_witness` is `None`, lazy
/// loading will be disabled.
#[tracing::instrument(name = "incremental", level = "trace", skip_all)]
#[allow(clippy::too_many_arguments)]
async fn calculate_state_incremental<'a>(
	services: &Services,
	sender_user: &'a UserId,
	since_shortstatehash: Option<ShortStateHash>,
	current_shortstatehash: ShortStateHash,
	lazy_loading_witness: Option<&'a Witness>,
) -> Result<Vec<PduEvent>> {
	let since_shortstatehash = since_shortstatehash.unwrap_or(current_shortstatehash);

	let state_get_shorteventid = |user_id: &'a UserId| {
		services
			.rooms
			.state_accessor
			.state_get_shortid(
				current_shortstatehash,
				&StateEventType::RoomMember,
				user_id.as_str(),
			)
			.ok()
	};

	let lazy_state_ids: OptionFuture<_> = lazy_loading_witness
		.map(|witness| {
			StreamExt::into_future(
				witness
					.iter()
					.stream()
					.broad_filter_map(|user_id| state_get_shorteventid(user_id)),
			)
		})
		.into();

	let state_diff_shortids = services
		.rooms
		.state_accessor
		.state_added((since_shortstatehash, current_shortstatehash))
		.boxed();

	state_diff_shortids
		.broad_filter_map(|(shortstatekey, shorteventid)| async move {
			if lazy_loading_witness.is_none() {
				return Some(shorteventid);
			}

			lazy_filter(services, sender_user, shortstatekey, shorteventid).await
		})
		.chain(lazy_state_ids.stream())
		.broad_filter_map(|shorteventid| {
			services
				.rooms
				.short
				.get_eventid_from_short(shorteventid)
				.ok()
		})
		.broad_filter_map(|event_id: OwnedEventId| async move {
			services.rooms.timeline.get_pdu(&event_id).await.ok()
		})
		.collect::<Vec<_>>()
		.map(Ok)
		.await
}

async fn lazy_filter(
	services: &Services,
	sender_user: &UserId,
	shortstatekey: ShortStateKey,
	shorteventid: ShortEventId,
) -> Option<ShortEventId> {
	let (event_type, state_key) = services
		.rooms
		.short
		.get_statekey_from_short(shortstatekey)
		.await
		.ok()?;

	(event_type != StateEventType::RoomMember || state_key == sender_user.as_str())
		.then_some(shorteventid)
}

async fn calculate_counts(
	services: &Services,
	room_id: &RoomId,
	sender_user: &UserId,
) -> Result<(Option<u64>, Option<u64>, Option<Vec<OwnedUserId>>)> {
	let joined_member_count = services
		.rooms
		.state_cache
		.room_joined_count(room_id)
		.unwrap_or(0);

	let invited_member_count = services
		.rooms
		.state_cache
		.room_invited_count(room_id)
		.unwrap_or(0);

	let (joined_member_count, invited_member_count) =
		join(joined_member_count, invited_member_count).await;

	let small_room = joined_member_count.saturating_add(invited_member_count) <= 5;

	let heroes: OptionFuture<_> = small_room
		.then(|| calculate_heroes(services, room_id, sender_user))
		.into();

	Ok((Some(joined_member_count), Some(invited_member_count), heroes.await))
}

async fn calculate_heroes(
	services: &Services,
	room_id: &RoomId,
	sender_user: &UserId,
) -> Vec<OwnedUserId> {
	services
		.rooms
		.timeline
		.all_pdus(sender_user, room_id)
		.ready_filter(|(_, pdu)| pdu.kind == RoomMember)
		.fold_default(|heroes: Vec<_>, (_, pdu)| {
			fold_hero(heroes, services, room_id, sender_user, pdu)
		})
		.await
}

async fn fold_hero(
	mut heroes: Vec<OwnedUserId>,
	services: &Services,
	room_id: &RoomId,
	sender_user: &UserId,
	pdu: PduEvent,
) -> Vec<OwnedUserId> {
	let Some(user_id): Option<&UserId> =
		pdu.state_key.as_deref().map(TryInto::try_into).flat_ok()
	else {
		return heroes;
	};

	if user_id == sender_user {
		return heroes;
	}

	let Ok(content): Result<RoomMemberEventContent, _> = pdu.get_content() else {
		return heroes;
	};

	// The membership was and still is invite or join
	if !matches!(content.membership, MembershipState::Join | MembershipState::Invite) {
		return heroes;
	}

	if heroes.iter().any(is_equal_to!(user_id)) {
		return heroes;
	}

	let (is_invited, is_joined) = join(
		services.rooms.state_cache.is_invited(user_id, room_id),
		services.rooms.state_cache.is_joined(user_id, room_id),
	)
	.await;

	if !is_joined && is_invited {
		return heroes;
	}

	heroes.push(user_id.to_owned());
	heroes
}