macltok

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

use ruma::OwnedUserId;

use serde_json::Value;

use crate::{

    Error, ErrorKind, EventContext, MentionCaptureState, MentionCoverage, MentionEntry, MentionKey,

    MentionsBoundary, MentionsPage, SavedMentionContext, TimelineItem,

    normalise::MediaRegistration,

};

#[derive(Clone)]

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

pub(super) struct SavedEvent {

    pub item: TimelineItem,

    pub media: Vec<MediaRegistration>,

    pub raw: Option<Value>,

}

#[derive(Clone)]

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

pub(super) struct Candidate {

    pub room_id: ruma::OwnedRoomId,

    pub raw: Value,

    pub attempts: u32,

    pub retry_at_ms: u64,

}

#[derive(Clone)]

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

pub(super) struct Window {

    pub before: Vec<String>,

    pub target: Option<String>,

    pub after: Vec<String>,

    pub start: Option<String>,

    pub end: Option<String>,

}

#[derive(Clone, Default)]

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

pub(super) struct Data {

    pub revision: u64,

    pub entries: BTreeMap<String, MentionEntry>,

    pub events: BTreeMap<String, SavedEvent>,

    pub windows: BTreeMap<String, Window>,

    pub capture_attempts: BTreeMap<String, u32>,

    pub candidates: BTreeMap<String, Candidate>,

    pub redacted: BTreeMap<String, u64>,

    pub edits: BTreeMap<String, Value>,

    pub coverage: MentionCoverage,

    pub checkpoints: BTreeMap<String, String>,

    pub observed_rooms: BTreeSet<ruma::OwnedRoomId>,

    pub history_cursor: Option<String>,

    pub history_exhausted: bool,

    pub latest_cursor: Option<String>,

    #[cfg_attr(feature = "serde", serde(default))]

    pub latest_anchor: Option<String>,

    #[cfg_attr(feature = "serde", serde(default))]

    pub latest_head: Option<String>,

}

pub(super) fn id(key: &MentionKey) -> String {

    format!("{}\n{}", key.room_id, key.event_id)

}

pub(super) fn now() -> u64 {

    matrix_sdk::ruma::MilliSecondsSinceUnixEpoch::now()

        .get()

        .into()

}

pub(super) fn invalid(message: &'static str) -> Error {

    Error::new(ErrorKind::Invalid, message)

}

impl Data {

    pub fn touch_contexts(&mut self, identity: &str) {

        for (key, window) in &self.windows {

            if window.target.as_deref() == Some(identity)

                || window

                    .before

                    .iter()

                    .chain(&window.after)

                    .any(|key| key == identity)

            {

                if let Some(entry) = self.entries.get_mut(key) {

                    entry.revision = self.revision + 1;

                }

            }

        }

    }

    pub fn page(

        &self,

        account: &OwnedUserId,

        cursor: Option<&str>,

        limit: u16,

    ) -> Result<MentionsPage, Error> {

        let (boundary, last) = if let Some(cursor) = cursor {

            let value: Value =

                serde_json::from_str(cursor).map_err(|_| invalid("invalid mentions cursor"))?;

            if value["account"].as_str() != Some(account.as_str()) {

                return Err(invalid("mentions cursor belongs to a different account"));

            }

            let boundary = value["revision"]

                .as_u64()

                .ok_or_else(|| invalid("invalid mentions cursor"))?;

            let timestamp = value["timestamp"]

                .as_u64()

                .ok_or_else(|| invalid("invalid mentions cursor"))?;

            let key = value["key"]

                .as_str()

                .ok_or_else(|| invalid("invalid mentions cursor"))?

                .to_owned();

            if boundary > self.revision {

                return Err(invalid("mentions cursor is newer than this store"));

            }

            (boundary, Some((timestamp, key)))

        } else {

            (self.revision, None)

        };

        let mut entries: Vec<_> = self

            .entries

            .iter()

            .filter(|(key, entry)| {

                entry.discovery_revision <= boundary

                    && last.as_ref().is_none_or(|last| {

                        (entry.timestamp_ms, key.as_str()) < (last.0, last.1.as_str())

                    })

            })

            .map(|(_, entry)| entry.clone())

            .collect();

        entries.sort_by(|a, b| (b.timestamp_ms, id(&b.key)).cmp(&(a.timestamp_ms, id(&a.key))));

        let limit = usize::from(limit.clamp(1, 50));

        let end_reached = entries.len() <= limit;

        entries.truncate(limit);

        let next_cursor = if end_reached {

            None

        } else {

            entries.last().map(|entry| serde_json::json!({

            "account": account.as_str(), "revision": boundary, "timestamp": entry.timestamp_ms, "key": id(&entry.key)

        }).to_string())

        };

        let mut coverage = self.coverage.clone();

        coverage.unresolved_candidates = self.candidates.len();

        Ok(MentionsPage {

            entries,

            next_cursor,

            end_reached,

            coverage,

            unread_count: self.entries.values().filter(|entry| !entry.read).count(),

            revision: self.revision,

            boundary: MentionsBoundary {

                account: account.clone(),

                revision: boundary,

            },

        })

    }

    pub fn context(

        &self,

        key: &MentionKey,

    ) -> Result<(SavedMentionContext, Vec<MediaRegistration>), Error> {

        let entry = self

            .entries

            .get(&id(key))

            .ok_or_else(|| Error::new(ErrorKind::EventNotFound, "mention was not found"))?;

        let mut media = Vec::new();

        let mut item = |key: &String| {

            self.events.get(key).map(|event| {

                media.extend(event.media.clone());

                event.item.clone()

            })

        };

        let context = self.windows.get(&id(key)).map(|window| EventContext {

            events_before: window.before.iter().filter_map(&mut item).collect(),

            event: window.target.as_ref().and_then(&mut item),

            events_after: window.after.iter().filter_map(&mut item).collect(),

            start_token: window.start.clone(),

            end_token: window.end.clone(),

        });

        Ok((

            SavedMentionContext {

                key: key.clone(),

                state: entry.capture.clone(),

                context,

                revision: self.revision,

            },

            media,

        ))

    }

    pub fn collect_events(&mut self) {

        let mut used: BTreeSet<String> = self.entries.keys().cloned().collect();

        for window in self.windows.values() {

            used.extend(window.before.iter().chain(&window.after).cloned());

            used.extend(window.target.iter().cloned());

        }

        self.events.retain(|key, _| used.contains(key));

        self.capture_attempts

            .retain(|key, _| self.entries.contains_key(key));

    }

    pub fn evict_one(&mut self) -> bool {

        let oldest = |keys: Vec<String>, entries: &BTreeMap<String, MentionEntry>| {

            keys.into_iter()

                .min_by_key(|key| entries.get(key).map_or(0, |entry| entry.timestamp_ms))

        };

        if let Some(key) = oldest(

            self.windows

                .keys()

                .filter(|key| self.entries.get(*key).is_some_and(|entry| entry.read))

                .cloned()

                .collect(),

            &self.entries,

        )

        .or_else(|| oldest(self.windows.keys().cloned().collect(), &self.entries))

        {

            self.windows.remove(&key);

            if let Some(entry) = self.entries.get_mut(&key) {

                entry.capture = MentionCaptureState::Unavailable;

                entry.revision = self.revision;

            }

        } else if let Some(key) = self

            .entries

            .iter()

            .min_by_key(|(_, entry)| (!entry.read, entry.timestamp_ms))

            .map(|(key, _)| key.clone())

        {

            self.entries.remove(&key);

        } else if let Some(key) = self.candidates.keys().next().cloned() {

            self.candidates.remove(&key);

        } else if let Some(key) = self.edits.keys().next().cloned() {

            self.edits.remove(&key);

        } else if let Some(key) = self.redacted.keys().next().cloned() {

            self.redacted.remove(&key);

        } else {

            return false;

        }

        self.coverage.truncated = true;

        self.collect_events();

        true

    }

    pub fn bound(&mut self, entries: usize, events: usize, candidates: usize) {

        while self.entries.len() > entries {

            if let Some(key) = self

                .entries

                .iter()

                .min_by_key(|(_, entry)| (!entry.read, entry.timestamp_ms))

                .map(|(key, _)| key.clone())

            {

                self.entries.remove(&key);

                self.windows.remove(&key);

                self.coverage.truncated = true;

            }

        }

        self.collect_events();

        while self.events.len() > events && self.evict_one() {}

        while self.candidates.len() > candidates {

            let key = self.candidates.keys().next().cloned().unwrap();

            self.candidates.remove(&key);

            self.coverage.truncated = true;

        }

        for map in [&mut self.redacted] {

            while map.len() > entries + candidates {

                let key = map

                    .iter()

                    .min_by_key(|(_, timestamp)| *timestamp)

                    .map(|(key, _)| key.clone())

                    .unwrap();

                map.remove(&key);

                self.coverage.truncated = true;

            }

        }

        while self.edits.len() > candidates {

            let key = self.edits.keys().next().cloned().unwrap();

            self.edits.remove(&key);

            self.coverage.truncated = true;

        }

        self.coverage.unresolved_candidates = self.candidates.len();

    }

}