macltok

use std::{collections::HashSet, sync::Arc, time::Duration};

use futures_util::StreamExt;

use matrix_sdk::{

    Client, Room, SessionChange,

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

    send_queue::{LocalEchoContent, RoomSendQueueUpdate, SendQueueUpdate},

};

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

use ruma::{OwnedDeviceId, OwnedRoomId, OwnedTransactionId, time::Instant};

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

use super::executor::TaskSet;

use crate::{

    ConnectionState, Event, Termination, UploadOutcome, encryption::encryption_state,

    error::classify_sdk_error, 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;

    }

}

pub(super) async fn session(

    mut changes: tokio::sync::broadcast::Receiver<SessionChange>,

    events: EventSender,

    device_id: OwnedDeviceId,

    stop: mpsc::UnboundedSender<Termination>,

) {

    loop {

        let event = match changes.recv().await {

            Ok(SessionChange::TokensRefreshed) => Event::SessionTokensChanged,

            Ok(SessionChange::UnknownToken(data)) => {

                let _ = events

                    .send(Event::SessionAccessRejected {

                        soft_logout: data.soft_logout,

                    })

                    .await;

                let _ = stop.send(if data.soft_logout {

                    Termination::SoftLoggedOut { device_id }

                } else {

                    Termination::AccessRevoked

                });

                return;

            }

            // a missed refresh costs one needless save, a missed rejection comes round again

            Err(RecvError::Lagged(_)) => Event::SessionTokensChanged,

            Err(RecvError::Closed) => return,

        };

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

            return;

        }

    }

}

/// `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,

        }

    }

}

/// upload progress only goes out on the room's own channel, never the client-wide one

pub(super) async fn room_uploads(room: Room, events: EventSender, ready: watch::Sender<bool>) {

    let room_id = room.room_id().to_owned();

    let mut updates = match room.send_queue().subscribe().await {

        Ok((_, updates)) => updates,

        Err(error) => {

            tracing::warn!(%room_id, %error, "reading the media send queue failed");

            return;

        }

    };

    ready.send_replace(true);

    loop {

        let update = match updates.recv().await {

            Ok(update) => update,

            // a missed tick is fine, the next one says where it's got to

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

            Err(RecvError::Closed) => return,

        };

        let RoomSendQueueUpdate::MediaUpload {

            related_to,

            file,

            progress,

            ..

        } = update

        else {

            continue;

        };

        let event = Event::MediaUploadProgress {

            room_id: room_id.clone(),

            transaction_id: related_to,

            uploaded_bytes: progress.current as u64,

            total_bytes: progress.total as u64,

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

        };

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

            return;

        }

    }

}

/// progress only goes out on each room's own channel, everything else reaches this

/// one, which exists before any upload does

pub(super) async fn uploads(

    client: Client,

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

    events: EventSender,

) {

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

    let mut uploads = attachment_echoes(&client).await;

    loop {

        let finished = tokio::select! {

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

                Ok(update) => match upload_update(&mut uploads, update) {

                    Some(event) => vec![event],

                    None => continue,

                },

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

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

                    let current = attachment_echoes(&client).await;

                    let missed = uploads.difference(&current).cloned().collect::<Vec<_>>();

                    uploads = current;

                    finished_uploads(missed, UploadOutcome::Missed)

                }

                Err(RecvError::Closed) => return,

            },

            update = rooms.next() => {

                if update.is_none() {

                    return;

                }

                let forgotten = uploads

                    .iter()

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

                    .cloned()

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

                for upload in &forgotten {

                    uploads.remove(upload);

                }

                finished_uploads(forgotten, UploadOutcome::RoomForgotten)

            }

        };

        for event in finished {

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

                return;

            }

        }

    }

}

type Upload = (OwnedRoomId, OwnedTransactionId);

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 upload_update(uploads: &mut HashSet<Upload>, update: SendQueueUpdate) -> Option<Event> {

    let room_id = update.room_id;

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

        RoomSendQueueUpdate::NewLocalEvent(echo) => {

            if is_attachment_echo(&echo.content) {

                uploads.insert((room_id, echo.transaction_id));

            }

            return None;

        }

        RoomSendQueueUpdate::SendError {

            transaction_id,

            error,

            is_recoverable,

        } => {

            let upload = (room_id, transaction_id);

            if !uploads.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,

    };

    let upload = (room_id, transaction_id);

    if !uploads.remove(&upload) {

        return None;

    }

    let (room_id, transaction_id) = upload;

    Some(Event::MediaUploadFinished {

        room_id,

        transaction_id,

        outcome,

    })

}

fn finished_uploads(uploads: Vec<Upload>, outcome: UploadOutcome) -> Vec<Event> {

    uploads

        .into_iter()

        .map(|(room_id, transaction_id)| Event::MediaUploadFinished {

            room_id,

            transaction_id,

            outcome: outcome.clone(),

        })

        .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) {

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

    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 state = encryption_state(

            device_id.clone(),

            recovery_state,

            verification_state,

            backup_state,

        );

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

            return;

        }

        tokio::select! {

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

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

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

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

            }

            else => return,

        }

    }

}

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

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

    let mut info = client.room_info_notable_update_receiver();

    loop {

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

            return;

        }

        tokio::select! {

            update = rooms.next() => if update.is_none() { return },

            update = info.recv() => if matches!(update, Err(RecvError::Closed)) { return },

        }

        // a sync touching fifty rooms is one invalidation, not fifty

        while let Ok(_) | Err(tokio::sync::broadcast::error::TryRecvError::Lagged(_)) =

            info.try_recv()

        {}

    }

}

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

    let settings = client.notification_settings().await;

    let mut changes = settings.subscribe_to_changes();

    while let Ok(()) | Err(RecvError::Lagged(_)) = changes.recv().await {

        // room summaries carry the effective mode, and push rule changes never touch room info

        for event in [Event::NotificationSettingsChanged, Event::RoomsChanged] {

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

                return;

            }

        }

    }

}