macltok

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

use ruma::{OwnedRoomId, OwnedUserId};

use serde_json::Value;

use crate::{

    Error, ErrorKind, EventContext, MentionCaptureState, MentionCoverage, MentionEntry, MentionGap,

    MentionGapReason, MentionKey, MentionsBoundary, MentionsPage, SavedMentionContext,

    TimelineItem, normalise::MediaRegistration,

};

const CANDIDATE_LIFETIME_MS: u64 = 7 * 24 * 60 * 60 * 1000;

#[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>,

}

impl PartialEq for SavedEvent {

    fn eq(&self, other: &Self) -> bool {

        self.item == other.item

            && self.raw == other.raw

            && self

                .media

                .iter()

                .map(|media| &media.id)

                .eq(other.media.iter().map(|media| &media.id))

    }

}

#[derive(Clone, PartialEq)]

#[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,

    pub first_seen_ms: u64,

}

impl Candidate {

    pub fn new(room_id: ruma::OwnedRoomId, raw: Value, retry_at_ms: u64) -> Self {

        Self {

            room_id,

            raw,

            attempts: 0,

            retry_at_ms,

            first_seen_ms: now(),

        }

    }

    pub fn encrypted(&self) -> bool {

        self.raw["type"] == "m.room.encrypted"

    }

}

#[derive(Clone, PartialEq)]

#[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, PartialEq)]

#[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 history_cursor: Option<String>,

    pub history_exhausted: bool,

    pub latest_cursor: Option<String>,

    pub latest_anchor: Option<String>,

    pub latest_head: Option<String>,

}

#[derive(Clone, Default, PartialEq)]

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

pub(super) struct SyncState {

    pub inbox_revision: u64,

    pub checkpoints: BTreeMap<String, Checkpoint>,

}

/// what a set of rooms holds in the inbox, compared before and after a change so an

/// unchanged inbox costs only the rooms it touched

#[derive(PartialEq)]

pub(super) struct Scope {

    rooms: Vec<RoomScope>,

    coverage: MentionCoverage,

    cursors: (

        Option<String>,

        bool,

        Option<String>,

        Option<String>,

        Option<String>,

    ),

}

#[derive(PartialEq)]

struct RoomScope {

    entries: Vec<MentionEntry>,

    events: Vec<SavedEvent>,

    windows: Vec<Window>,

    capture_attempts: Vec<u32>,

    candidates: Vec<Candidate>,

    redacted: Vec<u64>,

    edits: Vec<Value>,

    keys: Vec<String>,

}

fn in_room<V: Clone>(

    map: &BTreeMap<String, V>,

    room: &ruma::RoomId,

    keys: &mut Vec<String>,

) -> Vec<V> {

    // no room id holds a newline, so a room's "room\nevent" keys sit below "room\x0b"

    map.range(format!("{room}\n")..format!("{room}\u{b}"))

        .map(|(key, value)| {

            keys.push(key.clone());

            value.clone()

        })

        .collect()

}

#[derive(Clone, Debug, PartialEq)]

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

pub(super) struct Checkpoint {

    pub pos: String,

    pub previous: 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)

}

#[cfg(feature = "serde")]

pub(super) fn json_len(value: &impl serde::Serialize) -> usize {

    struct Count(usize);

    impl std::io::Write for Count {

        fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {

            self.0 += bytes.len();

            Ok(bytes.len())

        }

        fn flush(&mut self) -> std::io::Result<()> {

            Ok(())

        }

    }

    let mut count = Count(0);

    let _ = serde_json::to_writer(&mut count, value);

    count.0

}

#[cfg(not(feature = "serde"))]

pub(super) fn json_len<T>(_: &T) -> usize {

    0

}

fn timestamp(value: &Value) -> u64 {

    value["origin_server_ts"].as_u64().unwrap_or(0)

}

impl Data {

    pub fn scope<'a>(&self, rooms: impl IntoIterator<Item = &'a ruma::OwnedRoomId>) -> Scope {

        Scope {

            rooms: rooms

                .into_iter()

                .map(|room| {

                    let mut keys = Vec::new();

                    RoomScope {

                        entries: in_room(&self.entries, room, &mut keys),

                        events: in_room(&self.events, room, &mut keys),

                        windows: in_room(&self.windows, room, &mut keys),

                        capture_attempts: in_room(&self.capture_attempts, room, &mut keys),

                        candidates: in_room(&self.candidates, room, &mut keys),

                        redacted: in_room(&self.redacted, room, &mut keys),

                        edits: in_room(&self.edits, room, &mut keys),

                        keys,

                    }

                })

                .collect(),

            coverage: self.coverage.clone(),

            cursors: (

                self.history_cursor.clone(),

                self.history_exhausted,

                self.latest_cursor.clone(),

                self.latest_anchor.clone(),

                self.latest_head.clone(),

            ),

        }

    }

    /// folds an inbox built in memory while the store was unreachable into the one the

    /// store held, keeping the newer side where both have a key

    pub fn absorb(&mut self, newer: Data) {

        self.revision = self.revision.max(newer.revision);

        self.entries.extend(newer.entries);

        self.events.extend(newer.events);

        self.windows.extend(newer.windows);

        self.capture_attempts.extend(newer.capture_attempts);

        self.candidates.extend(newer.candidates);

        self.redacted.extend(newer.redacted);

        self.edits.extend(newer.edits);

        for gap in newer.coverage.gaps {

            self.gap(gap.room_id, gap.reason, gap.token);

        }

        self.coverage.earliest_timestamp_ms = self

            .coverage

            .earliest_timestamp_ms

            .into_iter()

            .chain(newer.coverage.earliest_timestamp_ms)

            .min();

        self.coverage.truncated |= newer.coverage.truncated;

        self.history_cursor = newer.history_cursor.or(self.history_cursor.take());

        self.latest_cursor = newer.latest_cursor.or(self.latest_cursor.take());

        self.latest_anchor = newer.latest_anchor.or(self.latest_anchor.take());

        self.latest_head = newer.latest_head.or(self.latest_head.take());

    }

    pub fn gap(

        &mut self,

        room_id: Option<ruma::OwnedRoomId>,

        reason: MentionGapReason,

        token: Option<String>,

    ) {

        if self

            .coverage

            .gaps

            .iter()

            .any(|gap| gap.room_id == room_id && gap.reason == reason && gap.token == token)

        {

            return;

        }

        self.coverage.gaps.push(MentionGap {

            room_id,

            reason,

            token,

        });

        if self.coverage.gaps.len() > 128 {

            self.coverage.gaps.remove(0);

            self.coverage.truncated = true;

        }

    }

    pub fn expire_candidates(&mut self, now_ms: u64) -> bool {

        let expired: Vec<_> = self

            .candidates

            .iter()

            .filter(|(_, candidate)| {

                now_ms.saturating_sub(candidate.first_seen_ms) > CANDIDATE_LIFETIME_MS

            })

            .map(|(key, candidate)| (key.clone(), candidate.room_id.clone()))

            .collect();

        for (key, room_id) in &expired {

            self.candidates.remove(key);

            self.gap(

                Some(room_id.clone()),

                MentionGapReason::CandidateExpired,

                None,

            );

        }

        !expired.is_empty()

    }

    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,

        hidden: &BTreeSet<OwnedRoomId>,

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

                !hidden.contains(&entry.key.room_id)

                    && 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.gaps.retain(|gap| {

            gap.room_id

                .as_ref()

                .is_none_or(|room_id| !hidden.contains(room_id))

        });

        coverage.unresolved_candidates = self

            .candidates

            .values()

            .filter(|candidate| !hidden.contains(&candidate.room_id))

            .count();

        Ok(MentionsPage {

            entries,

            next_cursor,

            end_reached,

            coverage,

            unread_count: self

                .entries

                .values()

                .filter(|entry| !entry.read && !hidden.contains(&entry.key.room_id))

                .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)).and_then(|window| {

            let event = window.target.as_ref().and_then(&mut item)?;

            Some(EventContext {

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

                event,

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

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

        }

        let mut freed = 0;

        self.events.retain(|key, event| {

            let keep = used.contains(key);

            if !keep {

                freed += key.len() + json_len(event);

            }

            keep

        });

        self.capture_attempts

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

        freed

    }

    /// returns roughly how many plaintext JSON bytes the eviction freed

    pub fn evict_one(&mut self) -> Option<usize> {

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

        };

        let freed = 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))

        {

            let window = self.windows.remove(&key);

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

                entry.capture = MentionCaptureState::Unavailable;

                entry.revision = self.revision;

            }

            key.len() + json_len(&window)

        } else if let Some(key) = self