macltok

use std::{

    collections::{HashMap, HashSet, VecDeque},

    sync::Arc,

    time::Duration,

};

use futures_util::StreamExt;

use matrix_sdk::{

    Client, Room,

    event_handler::EventHandlerHandle,

    ruma::events::{AnyMessageLikeEventContent, room::message::MessageType},

    send_queue::{LocalEchoContent, RoomSendQueueUpdate, SendQueueUpdate},

    serde_helpers::extract_thread_root,

};

use matrix_sdk_base::RoomInfoNotableUpdateReasons;

use matrix_sdk_ui::sync_service::{State as SyncState, SyncService};

use ruma::{

    OwnedEventId, OwnedRoomId, OwnedTransactionId,

    events::{

        AnyGlobalAccountDataEvent, AnyRoomAccountDataEvent, AnySyncStateEvent,

        AnySyncTimelineEvent, receipt::SyncReceiptEvent,

    },

    serde::Raw,

    time::Instant,

};

use tokio::sync::{

    Notify,

    broadcast::error::{RecvError, TryRecvError},

    mpsc, watch,

};

use super::executor::TaskSet;

use crate::{

    ConnectionState, Event, UploadOutcome,

    encryption::encryption_state,

    error::classify_sdk_error,

    normalise::{ENABLED_ROOMS, ROOM_PACKS, USER_PACK, recount_thread_unreads},

    space_order,

    transport::EventSender,

};

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);

const SYNC_STABLE: Duration = Duration::from_secs(30);

pub(super) async fn room_subscriptions(

    sync: Arc<SyncService>,

    mut rooms: watch::Receiver<Vec<ruma::OwnedRoomId>>,

) {

    while rooms.changed().await.is_ok() {

        let room_ids = rooms.borrow_and_update().clone();

        let room_ids = room_ids.iter().map(AsRef::as_ref).collect::<Vec<_>>();

        sync.room_list_service().subscribe_to_rooms(&room_ids).await;

    }

}

/// `states` must be taken before sync starts, a late subscriber never hears the transition it missed

pub(super) async fn connection(

    mut states: impl futures_util::Stream<Item = SyncState> + Unpin,

    sync: Arc<SyncService>,

    client: Client,

    events: EventSender,

    connection: watch::Sender<ConnectionState>,

) {

    let mut backoff = SYNC_RESTART_MIN;

    let mut online_since: Option<Instant> = None;

    while let Some(state) = states.next().await {

        // only a run that lasted counts as recovered, or a server that answers /versions

        // but refuses sync would have us bounce between offline and online forever

        let stable = online_since

            .take()

            .map(|since| since.elapsed() >= SYNC_STABLE);

        if stable == Some(true) {

            backoff = SYNC_RESTART_MIN;

        }

        let (state, restart) = match state {

            // only ever us stopping it for a restart, nothing a client needs to hear

            SyncState::Idle => continue,

            SyncState::Running => {

                online_since = Some(Instant::now());

                // the sdk parks a room queue on any send error and never unparks it

                client.send_queue().set_enabled(true).await;

                (ConnectionState::Online, false)

            }

            SyncState::Offline => (ConnectionState::Offline, stable == Some(false)),

            SyncState::Terminated => (ConnectionState::Offline, true),

            SyncState::Error(error) => {

                tracing::warn!(%error, "sync service reported an error");

                (ConnectionState::Offline, true)

            }

        };

        if *connection.borrow() != state {

            connection.send_replace(state);

            if events.send(Event::ConnectionChanged(state)).await.is_err() {

                return;

            }

        }

        if restart {

            sync.stop().await;

            matrix_sdk::sleep::sleep(backoff).await;

            backoff = (backoff * 2).min(SYNC_RESTART_MAX);

            sync.start().await;

        }

    }

}

/// a recoverable failure while we're online is a hiccup, so give the room another go

pub(super) async fn send_queue_errors(

    client: Client,

    mut errors: tokio::sync::broadcast::Receiver<matrix_sdk::send_queue::SendQueueRoomError>,

    online: watch::Receiver<ConnectionState>,

) {

    let mut retries = TaskSet::new();

    loop {

        let error = tokio::select! {

            error = errors.recv() => error,

            Some(_) = retries.join_next(), if !retries.is_empty() => continue,

        };

        match error {

            Ok(error) if error.is_recoverable => {

                let client = client.clone();

                let online = online.clone();

                retries.spawn(async move {

                    matrix_sdk::sleep::sleep(SEND_QUEUE_RETRY).await;

                    if *online.borrow() == ConnectionState::Online

                        && let Some(room) = client.get_room(&error.room_id)

                    {

                        room.send_queue().set_enabled(true);

                    }

                });

            }

            Ok(_) | Err(RecvError::Lagged(_)) => {}

            Err(RecvError::Closed) => return,

        }

    }

}

/// the one place upload events leave from, so an upload is never reported after it

/// ended. the client-wide channel exists before any upload does. it holds 32 updates and

/// progress ticks share it, so this keeps draining while the event stream is full

pub(super) async fn uploads(

    client: Client,

    mut updates: tokio::sync::broadcast::Receiver<SendQueueUpdate>,

    events: EventSender,

) {

    let (_, mut rooms) = client.rooms_stream();

    let mut tracker = UploadTracker::new(attachment_echoes(&client).await);

    let mut pending = VecDeque::new();

    loop {

        let reports: Vec<Event> = tokio::select! {

            biased;

            update = updates.recv() => match update {

                Ok(update) => tracker.update(update).into_iter().collect(),

                Err(RecvError::Lagged(skipped)) => {

                    tracing::warn!(skipped, "upload tracker lagged");

                    tracker.resnapshot(attachment_echoes(&client).await)

                }

                Err(RecvError::Closed) => return,

            },

            update = rooms.next() => {

                if update.is_none() {

                    return;

                }

                tracker.finish_where(

                    |(room_id, _)| client.get_room(room_id).is_none(),

                    &UploadOutcome::RoomForgotten,

                )

            }

            permit = events.reserve(), if !pending.is_empty() => {

                let Ok(permit) = permit else { return };

                permit.send(pending.pop_front().expect("pending was checked"));

                continue;

            }

        };

        for event in reports {

            queue_upload_event(&mut pending, event);

        }

    }

}

/// a progress tick replaces a waiting one for the same upload, unless something else about

/// that upload is queued after it

fn queue_upload_event(pending: &mut VecDeque<Event>, event: Event) {

    if matches!(event, Event::MediaUploadProgress { .. })

        && let Some(waiting) = pending

            .iter_mut()

            .rev()

            .find(|waiting| upload_of(waiting) == upload_of(&event))

        && matches!(waiting, Event::MediaUploadProgress { .. })

    {

        *waiting = event;

        return;

    }

    pending.push_back(event);

}

fn upload_of(event: &Event) -> Option<(&OwnedRoomId, &OwnedTransactionId)> {

    match event {

        Event::MediaUploadProgress {

            room_id,

            transaction_id,

            ..

        }

        | Event::MediaUploadFailed {

            room_id,

            transaction_id,

            ..

        }

        | Event::MediaUploadFinished {

            room_id,

            transaction_id,

            ..

        } => Some((room_id, transaction_id)),

        _ => None,

    }

}

type Upload = (OwnedRoomId, OwnedTransactionId);

/// the sdk sends progress ticks from a detached task, so a tick can land after its

/// upload's end. it trails by one task hop, so a short memory of ends covers it

const REMEMBERED_ENDS: usize = 256;

/// everything reported gets exactly one `MediaUploadFinished` and nothing after it

struct UploadTracker {

    live: HashSet<Upload>,

    ended: HashSet<Upload>,

    ended_order: VecDeque<Upload>,

}

impl UploadTracker {

    fn new(live: HashSet<Upload>) -> Self {

        Self {

            live,

            ended: HashSet::new(),

            ended_order: VecDeque::new(),

        }

    }

    fn update(&mut self, update: SendQueueUpdate) -> Option<Event> {

        let room_id = update.room_id;

        let (transaction_id, outcome) = match update.update {

            RoomSendQueueUpdate::NewLocalEvent(echo) => {

                let upload = (room_id, echo.transaction_id);

                if is_attachment_echo(&echo.content) && !self.ended.contains(&upload) {

                    self.live.insert(upload);

                }

                return None;

            }

            RoomSendQueueUpdate::MediaUpload {

                related_to,

                file,

                progress,

                ..

            } => {

                let upload = (room_id, related_to);

                if self.ended.contains(&upload) {

                    return None;

                }

                self.live.insert(upload.clone());

                let (room_id, transaction_id) = upload;

                return Some(Event::MediaUploadProgress {

                    room_id,

                    transaction_id,

                    uploaded_bytes: progress.current as u64,

                    total_bytes: progress.total as u64,

                    complete: file.is_some() && progress.current == progress.total,

                });

            }

            RoomSendQueueUpdate::SendError {

                transaction_id,

                error,

                is_recoverable,

            } => {

                let upload = (room_id, transaction_id);

                if !self.live.contains(&upload) {

                    return None;

                }

                let (room_id, transaction_id) = upload;

                return Some(Event::MediaUploadFailed {

                    room_id,

                    transaction_id,

                    kind: classify_sdk_error(&error),

                    will_retry: is_recoverable,

                });

            }

            RoomSendQueueUpdate::SentEvent {

                transaction_id,

                event_id,

            } => (transaction_id, UploadOutcome::Sent { event_id }),

            RoomSendQueueUpdate::CancelledLocalEvent { transaction_id } => {

                (transaction_id, UploadOutcome::Cancelled)

            }

            _ => return None,

        };

        self.finish((room_id, transaction_id), outcome)

    }

    fn resnapshot(&mut self, current: HashSet<Upload>) -> Vec<Event> {

        let reports = self.finish_where(|upload| !current.contains(upload), &UploadOutcome::Missed);

        let ended = &self.ended;

        self.live

            .extend(current.into_iter().filter(|upload| !ended.contains(upload)));

        reports

    }

    fn finish_where(

        &mut self,

        ends: impl Fn(&Upload) -> bool,

        outcome: &UploadOutcome,

    ) -> Vec<Event> {

        let ending = self

            .live

            .iter()

            .filter(|upload| ends(upload))

            .cloned()

            .collect::<Vec<_>>();

        ending

            .into_iter()

            .filter_map(|upload| self.finish(upload, outcome.clone()))

            .collect()

    }

    /// an end for an upload we never saw is remembered too, its ticks may still be coming

    fn finish(&mut self, upload: Upload, outcome: UploadOutcome) -> Option<Event> {

        if self.ended.insert(upload.clone()) {

            self.ended_order.push_back(upload.clone());

            if self.ended_order.len() > REMEMBERED_ENDS

                && let Some(oldest) = self.ended_order.pop_front()

            {

                self.ended.remove(&oldest);

            }

        }

        if !self.live.remove(&upload) {

            return None;

        }

        let (room_id, transaction_id) = upload;

        Some(Event::MediaUploadFinished {

            room_id,

            transaction_id,

            outcome,

        })

    }

}

async fn attachment_echoes(client: &Client) -> HashSet<Upload> {

    let Ok(echoes) = client.send_queue().local_echoes().await else {

        return HashSet::new();

    };

    echoes

        .into_iter()

        .flat_map(|(room_id, echoes)| {

            echoes

                .into_iter()

                .filter(|echo| is_attachment_echo(&echo.content))

                .map(move |echo| (room_id.clone(), echo.transaction_id))

        })

        .collect()

}

fn is_attachment_echo(content: &LocalEchoContent) -> bool {

    let LocalEchoContent::Event {

        serialized_event, ..

    } = content

    else {

        return false;

    };

    matches!(

        serialized_event.deserialize(),

        Ok(AnyMessageLikeEventContent::RoomMessage(content))

            if matches!(

                content.msgtype,

                MessageType::Audio(_) | MessageType::File(_) | MessageType::Image(_) | MessageType::Video(_)

            )

    )

}

pub(super) async fn encryption(

    client: Client,

    events: EventSender,

    mut enabling: watch::Receiver<bool>,

) {

    let Some(device_id) = client.device_id().map(ToOwned::to_owned) else {

        tracing::warn!("encryption watcher stopped because the client has no device id");

        return;

    };

    let recovery = client.encryption().recovery();

    let backups = client.encryption().backups();

    // until this is cached, the timeline labels undecryptable events from a local guess

    matrix_sdk::executor::spawn({

        let backups = backups.clone();

        async move {

            let _ = backups.exists_on_server().await;

        }

    });

    let mut recovery_stream = recovery.state_stream();

    let mut backup_stream = backups.state_stream();

    let mut verification_stream = client.encryption().verification_state();

    let mut recovery_state = recovery.state();

    let mut backup_state = backups.state();

    let mut verification_state = verification_stream.get();

    loop {

        let mut state = encryption_state(

            device_id.clone(),

            recovery_state,

            verification_state,

            backup_state,

        );

        if *enabling.borrow_and_update() {

            state.recovery = crate::RecoveryState::Enabling;

        }

        if events.send(Event::EncryptionChanged(state)).await.is_err() {

            return;

        }

        tokio::select! {

            Some(update) = recovery_stream.next() => recovery_state = update,

            Ok(()) = enabling.changed() => {}

            Some(update) = verification_stream.next() => verification_state = update,

            Some(update) = backup_stream.next() => {

                backup_state = update.unwrap_or_else(|_| backups.state());

            }

            else => return,

        }

    }

}

/// space orders, pins and power levels all feed summaries without raising a notable room info update

pub(super) async fn rooms(client: Client, events: EventSender) {

    let (_, mut rooms) = client.rooms_stream();

    let mut info = client.room_info_notable_update_receiver();

    let changed = Arc::new(Notify::new());

    let handlers = [

        client.add_event_handler({

            let changed = changed.clone();

            move |event: Raw<AnyRoomAccountDataEvent>| {

                let changed = changed.clone();

                async move {

                    if event_type(&event).as_deref() == Some(space_order::EVENT_TYPE) {

                        changed.notify_one();

                    }

                }

            }

        }),

        client.add_event_handler({

            let changed = changed.clone();

            move |event: Raw<AnySyncStateEvent>| {

                let changed = changed.clone();

                async move {

                    if matches!(

                        event_type(&event).as_deref(),

                        Some("m.room.pinned_events" | "m.room.power_levels")

                    ) {

                        changed.notify_one();

                    }

                }

            }

        }),

    ];

    let (touches, mut touched) = mpsc::unbounded_channel();

    let thread_handlers = thread_handlers(&client, &touches);

    let mut timelines = client.event_cache().subscribe_to_room_generic_updates();

    'invalidate: loop {

        if events.send(Event::RoomsChanged).await.is_err() {

            break;

        }

        // thread events and threaded receipts leave room info alone, so neither reaches `info`

        let mut pending = Touched::default();

        loop {

            let mut invalidate = tokio::select! {

                update = rooms.next() => if update.is_none() { break 'invalidate } else { true },

                update = info.recv() => {

                    if matches!(update, Err(RecvError::Closed)) { break 'invalidate }

                    true

                }

                () = changed.notified() => true,

                Some((room_id, root)) = touched.recv() => {

                    pending.add(room_id, root);

                    false

                }

                update = timelines.recv() => match update {

                    Ok(update) => {

                        pending.add(update.room_id, None);

                        false

                    }

                    Err(RecvError::Lagged(_)) => true,

                    Err(RecvError::Closed) => break 'invalidate,

                },

            };