macltok

use std::{

    collections::HashMap,

    fmt,

    pin::Pin,

    sync::{Arc, Mutex, MutexGuard, PoisonError},

    task::{Context, Poll},

    time::Duration,

};

use futures_util::{Stream, StreamExt};

use imbl::Vector;

use matrix_sdk::{

    SendOutsideWasm,

    executor::{AbortHandle, spawn},

};

use ruma::{OwnedDeviceId, OwnedRoomId, OwnedTransactionId, OwnedUserId};

#[cfg(feature = "serde")]

use serde::{Deserialize, Serialize};

use tokio::sync::broadcast;

use crate::{

    ClientHandle, Command, ConnectionState, EncryptionState, EngineFailure, ErrorKind, Event,

    HistoryState, RoomSummary, TimelineCommand, TimelineDiff, TimelineId, TimelineItem,

    TypingState, UploadOutcome, VerificationState, VerificationTarget,

    engine::executor::EngineBoxFuture,

};

#[cfg(not(target_family = "wasm"))]

type EngineBoxStream<T> = futures_util::stream::BoxStream<'static, T>;

#[cfg(target_family = "wasm")]

type EngineBoxStream<T> = futures_util::stream::LocalBoxStream<'static, T>;

const DEFAULT_CAPACITY: usize = 256;

const RESYNC_BACKOFF: Duration = Duration::from_millis(100);

const RESYNC_BACKOFF_CAP: Duration = Duration::from_secs(10);

/// everything the engine has said so far, folded together

#[derive(Clone, Debug, PartialEq)]

#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]

#[non_exhaustive]

pub struct State {

    pub user_id: Option<OwnedUserId>,

    pub session: SessionStatus,

    pub connection: ConnectionState,

    pub encryption: Option<EncryptionState>,

    pub verification: Option<VerificationState>,

    pub verification_target: Option<VerificationTarget>,

    pub rooms: Vector<RoomSummary>,

    pub timelines: imbl::HashMap<TimelineId, TimelineState>,

    pub typing: imbl::HashMap<OwnedRoomId, Vec<OwnedUserId>>,

    pub uploads: imbl::HashMap<OwnedTransactionId, UploadState>,

    /// what stopped the engine, if something went wrong

    pub engine_failure: Option<EngineFailure>,

    /// the last save of the session record that failed, until one succeeds

    pub session_save_failure: Option<String>,

}

#[derive(Clone, Debug, PartialEq)]

#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]

#[non_exhaustive]

pub struct TimelineState {

    pub items: Vector<TimelineItem>,

    pub history: HistoryState,

    /// a batch didn't fit, so `items` is the last good copy until the next reset

    pub diverged: bool,

}

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

#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]

#[non_exhaustive]

pub struct UploadState {

    pub room_id: OwnedRoomId,

    pub uploaded_bytes: u64,

    pub total_bytes: u64,

    pub complete: bool,

    pub failure: Option<UploadFailure>,

}

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

#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]

#[non_exhaustive]

pub struct UploadFailure {

    pub kind: ErrorKind,

    pub will_retry: bool,

}

/// the `Termination` it'll end in, as far as the events tell

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

#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]

#[non_exhaustive]

pub enum SessionStatus {

    Active,

    SoftLoggedOut {

        device_id: OwnedDeviceId,

    },

    AccessRevoked,

    LoggedOut,

    /// stopped with the session still good

    Stopped,

}

#[derive(Clone, Debug)]

#[non_exhaustive]

pub struct Update {

    /// one more than the update before it, and a `State` carries the one it's current as of

    pub sequence: u64,

    pub change: StateChange,

}

#[derive(Clone, Debug)]

#[non_exhaustive]

pub enum StateChange {

    /// start over from this. a subscription opens with one and gets another after lagging

    State(Arc<State>),

    Authenticated(OwnedUserId),

    Session(SessionStatus),

    /// the record on disk may now hold a spent refresh token

    SessionSaveFailed(String),

    /// saved again after `SessionSaveFailed`

    SessionSaved,

    Connection(ConnectionState),

    Encryption(EncryptionState),

    Verification(VerificationState),

    VerificationTarget(VerificationTarget),

    /// room modes arrive with the `Rooms` that follows

    NotificationSettings,

    Rooms(Vector<RoomSummary>),

    Timeline {

        timeline_id: TimelineId,

        change: TimelineChange,

    },

    History {

        timeline_id: TimelineId,

        state: HistoryState,

    },

    Typing(TypingState),

    Upload {

        transaction_id: OwnedTransactionId,

        upload: UploadState,

    },

    UploadFinished {

        room_id: OwnedRoomId,

        transaction_id: OwnedTransactionId,

        outcome: UploadOutcome,

    },

    Failed(EngineFailure),

}

#[derive(Clone, Debug)]

#[non_exhaustive]

pub enum TimelineChange {

    Reset {

        items: Vector<TimelineItem>,

    },

    /// `items` is what applying `diffs` left

    Changed {

        diffs: Arc<[TimelineDiff]>,

        items: Vector<TimelineItem>,

    },

    /// a batch didn't fit. a resync is on its way and nothing changes until its reset

    Diverged,

    Closed,

}

/// drains the event stream into a `State` that any number of `MirrorHandle`s watch

pub struct Mirror {

    handle: ClientHandle,

    events: EngineBoxStream<Event>,

    shared: Arc<Shared>,

    resyncs: HashMap<TimelineId, AbortHandle>,

    rooms_fetch: Option<AbortHandle>,

}

impl fmt::Debug for Mirror {

    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {

        formatter.debug_struct("Mirror").finish_non_exhaustive()

    }

}

impl Mirror {

    /// takes any stream of the engine's events, so a host can still look at them on the way past

    pub fn new(

        handle: ClientHandle,

        events: impl Stream<Item = Event> + SendOutsideWasm + 'static,

    ) -> (Self, MirrorHandle) {

        Self::with_capacity(handle, events, DEFAULT_CAPACITY)

    }

    /// `capacity` updates can queue for a subscriber before it lags and starts over. a zero

    /// capacity counts as one

    pub fn with_capacity(

        handle: ClientHandle,

        events: impl Stream<Item = Event> + SendOutsideWasm + 'static,

        capacity: usize,

    ) -> (Self, MirrorHandle) {

        let state = State {

            user_id: None,

            session: SessionStatus::Active,

            connection: *handle.connection_state().borrow(),

            encryption: handle.encryption_state(),

            verification: None,

            verification_target: None,

            rooms: Vector::new(),

            timelines: imbl::HashMap::new(),

            typing: imbl::HashMap::new(),

            uploads: imbl::HashMap::new(),

            engine_failure: None,

            session_save_failure: None,

        };

        let shared = Arc::new(Shared {

            published: Mutex::new(Published {

                state: Arc::new(state),

                sequence: 0,

                updates: Some(broadcast::channel(capacity.max(1)).0),

                rooms_fetch: RoomsFetch::Idle,

            }),

        });

        #[cfg(not(target_family = "wasm"))]

        let events = events.boxed();

        #[cfg(target_family = "wasm")]

        let events = events.boxed_local();

        let mirror = Self {

            handle,

            events,

            shared: shared.clone(),

            resyncs: HashMap::new(),

            rooms_fetch: None,

        };

        (mirror, MirrorHandle { shared })

    }

    /// resolves once the event stream ends, which is just after `Stopped`

    pub fn run(mut self) -> EngineBoxFuture<'static, ()> {

        Box::pin(async move {

            self.refetch_rooms();

            while let Some(event) = self.events.next().await {

                self.apply(event);

            }

        })

    }

    fn apply(&mut self, event: Event) {

        match event {

            Event::ConnectionChanged(connection) => self.publish(|state| {

                state.connection = connection;

                Some(StateChange::Connection(connection))

            }),

            Event::Authenticated { user_id } => self.publish(|state| {

                state.user_id = Some(user_id.clone());

                Some(StateChange::Authenticated(user_id))

            }),

            Event::SessionSaveFailed { message } => self.publish(|state| {

                state.session_save_failure = Some(message.clone());

                Some(StateChange::SessionSaveFailed(message))

            }),

            Event::SessionTokensChanged => self.publish(|state| {

                state.session_save_failure.take()?;

                Some(StateChange::SessionSaved)

            }),

            Event::SessionAccessRejected { soft_logout } => self.publish(|state| {

                state.session = match soft_logout {

                    Some(device_id) => SessionStatus::SoftLoggedOut { device_id },

                    None => SessionStatus::AccessRevoked,

                };

                Some(StateChange::Session(state.session.clone()))

            }),

            Event::LoggedOut => self.publish(|state| {

                state.session = SessionStatus::LoggedOut;

                Some(StateChange::Session(SessionStatus::LoggedOut))

            }),

            Event::Stopped => {

                self.cancel_resyncs();

                self.publish(|state| {

                    if state.session != SessionStatus::Active {

                        return None;

                    }

                    state.session = SessionStatus::Stopped;

                    Some(StateChange::Session(SessionStatus::Stopped))

                });

            }

            Event::Failed(failure) => self.publish(|state| {

                state.engine_failure = Some(failure.clone());

                Some(StateChange::Failed(failure))

            }),

            Event::NotificationSettingsChanged => {

                self.publish(|_| Some(StateChange::NotificationSettings));

            }

            Event::RoomsChanged => self.refetch_rooms(),

            Event::EncryptionChanged(encryption) => self.publish(|state| {

                state.encryption = Some(encryption.clone());

                Some(StateChange::Encryption(encryption))

            }),

            Event::VerificationChanged(verification) => self.publish(|state| {

                state.verification = Some(verification.clone());

                Some(StateChange::Verification(verification))

            }),

            Event::VerificationTargetChanged(target) => self.publish(|state| {

                state.verification_target = Some(target.clone());

                Some(StateChange::VerificationTarget(target))

            }),

            Event::TimelineReset { timeline_id, items } => {

                self.cancel_resync(&timeline_id);

                self.publish(|state| {

                    let history = state

                        .timelines

                        .get(&timeline_id)

                        .map_or(HistoryState::Idle, |timeline| timeline.history.clone());

                    let timeline = TimelineState {

                        items: items.clone(),

                        history,

                        diverged: false,

                    };

                    state.timelines.insert(timeline_id.clone(), timeline);

                    Some(StateChange::Timeline {

                        timeline_id,

                        change: TimelineChange::Reset { items },

                    })

                });

            }

            Event::TimelineChanged { timeline_id, diffs } => self.apply_diffs(timeline_id, diffs),

            Event::HistoryChanged {

                timeline_id,

                state: history,

            } => self.publish(|state| {

                state.timelines.get_mut(&timeline_id)?.history = history.clone();

                Some(StateChange::History {

                    timeline_id,

                    state: history,

                })

            }),

            Event::TimelineClosed { timeline_id } => {

                self.cancel_resync(&timeline_id);

                self.publish(|state| {

                    state.timelines.remove(&timeline_id)?;

                    Some(StateChange::Timeline {

                        timeline_id,

                        change: TimelineChange::Closed,

                    })

                });

            }

            Event::TypingChanged(typing) => self.publish(|state| {

                if typing.users.is_empty() {

                    state.typing.remove(&typing.room_id);

                } else {

                    state

                        .typing

                        .insert(typing.room_id.clone(), typing.users.clone());

                }

                Some(StateChange::Typing(typing))

            }),

            Event::MediaUploadProgress {

                room_id,

                transaction_id,

                uploaded_bytes,

                total_bytes,

                complete,

            } => self.publish(|state| {

                let upload = UploadState {

                    room_id,

                    uploaded_bytes,

                    total_bytes,

                    complete,

                    failure: None,

                };

                state.uploads.insert(transaction_id.clone(), upload.clone());

                Some(StateChange::Upload {

                    transaction_id,

                    upload,

                })

            }),

            Event::MediaUploadFailed {

                room_id,

                transaction_id,

                kind,

                will_retry,

            } => self.publish(|state| {

                let upload = state

                    .uploads

                    .entry(transaction_id.clone())

                    .or_insert_with(|| UploadState {

                        room_id,

                        uploaded_bytes: 0,

                        total_bytes: 0,

                        complete: false,

                        failure: None,

                    });

                upload.failure = Some(UploadFailure { kind, will_retry });

                Some(StateChange::Upload {

                    transaction_id,

                    upload: upload.clone(),

                })

            }),

            Event::MediaUploadFinished {

                room_id,

                transaction_id,

                outcome,

            } => self.publish(|state| {

                state.uploads.remove(&transaction_id);

                Some(StateChange::UploadFinished {

                    room_id,

                    transaction_id,

                    outcome,

                })

            }),

        }

    }

    fn apply_diffs(&mut self, timeline_id: TimelineId, diffs: Vec<TimelineDiff>) {

        let mut diverged = false;

        self.publish(|state| {

            let Some(timeline) = state.timelines.get_mut(&timeline_id) else {

                diverged = true;

                state.timelines.insert(

                    timeline_id.clone(),

                    TimelineState {

                        items: Vector::new(),

                        history: HistoryState::Idle,

                        diverged: true,

                    },

                );

                return Some(StateChange::Timeline {

                    timeline_id: timeline_id.clone(),

                    change: TimelineChange::Diverged,

                });

            };

            if timeline.diverged {

                return None;

            }

            let diffs: Arc<[TimelineDiff]> = diffs.into();

            let change = match TimelineDiff::try_apply_all(diffs.to_vec(), &mut timeline.items) {

                Ok(()) => TimelineChange::Changed {

                    diffs,

                    items: timeline.items.clone(),

                },

                Err(error) => {

                    tracing::debug!(?timeline_id, length = error.length, "timeline diverged");

                    timeline.diverged = true;

                    diverged = true;

                    TimelineChange::Diverged

                }

            };

            Some(StateChange::Timeline {

                timeline_id: timeline_id.clone(),

                change,

            })

        });

        if diverged {

            let task = spawn(resync(self.handle.clone(), timeline_id.clone()));

            self.resyncs.insert(timeline_id, task.abort_handle());

        }

    }

    fn refetch_rooms(&mut self) {

        let mut published = self.shared.lock();

        match published.rooms_fetch {

            RoomsFetch::Idle => {

                published.rooms_fetch = RoomsFetch::Running;

                drop(published);

                let task = spawn(fetch_rooms(self.handle.clone(), self.shared.clone()));

                self.rooms_fetch = Some(task.abort_handle());

            }

            RoomsFetch::Running => published.rooms_fetch = RoomsFetch::Again,

            RoomsFetch::Again => {}

        }

    }

    fn cancel_resync(&mut self, timeline_id: &TimelineId) {

        if let Some(task) = self.resyncs.remove(timeline_id) {

            task.abort();

        }

    }

    fn cancel_resyncs(&mut self) {

        for (_, task) in self.resyncs.drain() {

            task.abort();

        }

    }

    fn publish(&self, change: impl FnOnce(&mut State) -> Option<StateChange>) {

        self.shared.publish(change);

    }

}

impl Drop for Mirror {

    fn drop(&mut self) {