elektromail/

store.rs

use std::collections::{HashMap, HashSet};
use std::io;
use std::sync::OnceLock;

use time::{OffsetDateTime, format_description};

use crate::sqlite_store::SqliteStore;

/// Storage backend configuration.
#[derive(Clone, Debug, Default)]
pub enum StorageBackend {
    /// In-memory storage (default, non-persistent).
    #[default]
    InMemory,
    /// SQLite storage with persistence to the given file path.
    Sqlite(String),
}

/// Unified store that can be either in-memory or SQLite-backed.
pub enum Store {
    /// In-memory storage (non-persistent).
    InMemory(InMemoryStore),
    /// SQLite-backed storage (persistent).
    Sqlite(SqliteStore),
}

#[derive(Clone)]
pub(crate) struct StoreSnapshot {
    pub users: Vec<String>,
    pub mailboxes: HashMap<String, Vec<String>>,
    pub subscriptions: HashMap<String, Vec<String>>,
    pub messages: HashMap<(String, String), Vec<Message>>,
}

impl Store {
    /// Create a new store from the given backend configuration.
    pub fn new(backend: StorageBackend) -> io::Result<Self> {
        match backend {
            StorageBackend::InMemory => Ok(Self::InMemory(InMemoryStore::new())),
            StorageBackend::Sqlite(path) => {
                let store = SqliteStore::new(&path).map_err(io::Error::other)?;
                Ok(Self::Sqlite(store))
            }
        }
    }

    #[allow(dead_code)]
    pub fn ensure_mailbox(&mut self, user: &str, mailbox: &str) {
        match self {
            Self::InMemory(s) => s.ensure_mailbox(user, mailbox),
            Self::Sqlite(s) => {
                let _ = s.ensure_mailbox_for_user(user, mailbox);
            }
        }
    }

    pub fn list_mailboxes(&mut self, user: &str) -> Vec<String> {
        match self {
            Self::InMemory(s) => s.list_mailboxes(user),
            Self::Sqlite(s) => s.list_mailboxes(user).unwrap_or_default(),
        }
    }

    pub fn create_mailbox(&mut self, user: &str, mailbox: &str) -> bool {
        match self {
            Self::InMemory(s) => s.create_mailbox(user, mailbox),
            Self::Sqlite(s) => s.create_mailbox(user, mailbox).unwrap_or(false),
        }
    }

    pub fn rename_mailbox(&mut self, user: &str, old: &str, new: &str) -> bool {
        match self {
            Self::InMemory(s) => s.rename_mailbox(user, old, new),
            Self::Sqlite(s) => s.rename_mailbox(user, old, new).unwrap_or(false),
        }
    }

    pub fn delete_mailbox(&mut self, user: &str, mailbox: &str) -> bool {
        match self {
            Self::InMemory(s) => s.delete_mailbox(user, mailbox),
            Self::Sqlite(s) => s.delete_mailbox(user, mailbox).unwrap_or(false),
        }
    }

    pub fn subscribe(&mut self, user: &str, mailbox: &str) -> bool {
        match self {
            Self::InMemory(s) => s.subscribe(user, mailbox),
            Self::Sqlite(s) => s.subscribe(user, mailbox).unwrap_or(false),
        }
    }

    pub fn unsubscribe(&mut self, user: &str, mailbox: &str) -> bool {
        match self {
            Self::InMemory(s) => s.unsubscribe(user, mailbox),
            Self::Sqlite(s) => s.unsubscribe(user, mailbox).unwrap_or(false),
        }
    }

    pub fn list_subscriptions(&mut self, user: &str) -> Vec<String> {
        match self {
            Self::InMemory(s) => s.list_subscriptions(user),
            Self::Sqlite(s) => s.list_subscriptions(user).unwrap_or_default(),
        }
    }

    pub fn append(
        &mut self,
        user: &str,
        mailbox: &str,
        data: Vec<u8>,
        internal_date: String,
    ) -> u32 {
        match self {
            Self::InMemory(s) => s.append(user, mailbox, data, internal_date),
            Self::Sqlite(s) => s.append(user, mailbox, data, internal_date).unwrap_or(0),
        }
    }

    pub fn append_with_flags(
        &mut self,
        user: &str,
        mailbox: &str,
        data: Vec<u8>,
        internal_date: String,
        flags: &FlagSet,
    ) -> u32 {
        match self {
            Self::InMemory(s) => s.append_with_flags(user, mailbox, data, internal_date, flags),
            Self::Sqlite(s) => s
                .append_with_flags(user, mailbox, data, internal_date, flags)
                .unwrap_or(0),
        }
    }

    pub fn list(&mut self, user: &str, mailbox: &str) -> Vec<Message> {
        match self {
            Self::InMemory(s) => s.list(user, mailbox),
            Self::Sqlite(s) => s.list(user, mailbox).unwrap_or_default(),
        }
    }

    pub fn apply_flags_by_seq(
        &mut self,
        user: &str,
        mailbox: &str,
        seq: u32,
        op: FlagOp,
        flags: &FlagSet,
    ) {
        match self {
            Self::InMemory(s) => s.apply_flags_by_seq(user, mailbox, seq, op, flags),
            Self::Sqlite(s) => {
                let _ = s.apply_flags_by_seq(user, mailbox, seq, op, flags);
            }
        }
    }

    pub fn apply_flags_by_uid(
        &mut self,
        user: &str,
        mailbox: &str,
        uid: u32,
        op: FlagOp,
        flags: &FlagSet,
    ) {
        match self {
            Self::InMemory(s) => s.apply_flags_by_uid(user, mailbox, uid, op, flags),
            Self::Sqlite(s) => {
                let _ = s.apply_flags_by_uid(user, mailbox, uid, op, flags);
            }
        }
    }

    pub fn expunge_deleted(&mut self, user: &str, mailbox: &str) -> Vec<u32> {
        match self {
            Self::InMemory(s) => s.expunge_deleted(user, mailbox),
            Self::Sqlite(s) => s.expunge_deleted(user, mailbox).unwrap_or_default(),
        }
    }

    pub fn expunge_deleted_by_uid(&mut self, user: &str, mailbox: &str, uids: &[u32]) -> Vec<u32> {
        match self {
            Self::InMemory(s) => s.expunge_deleted_by_uid(user, mailbox, uids),
            Self::Sqlite(s) => s
                .expunge_deleted_by_uid(user, mailbox, uids)
                .unwrap_or_default(),
        }
    }

    pub fn copy_by_seq_set(&mut self, user: &str, src: &str, seqs: &[u32], dest: &str) -> usize {
        match self {
            Self::InMemory(s) => s.copy_by_seq_set(user, src, seqs, dest),
            Self::Sqlite(s) => s.copy_by_seq_set(user, src, seqs, dest).unwrap_or(0),
        }
    }

    pub fn move_by_seq_set(&mut self, user: &str, src: &str, seqs: &[u32], dest: &str) -> Vec<u32> {
        match self {
            Self::InMemory(s) => s.move_by_seq_set(user, src, seqs, dest),
            Self::Sqlite(s) => s.move_by_seq_set(user, src, seqs, dest).unwrap_or_default(),
        }
    }

    pub fn seqs_from_uids(&self, user: &str, mailbox: &str, uids: &[u32]) -> Vec<u32> {
        match self {
            Self::InMemory(s) => s.seqs_from_uids(user, mailbox, uids),
            Self::Sqlite(s) => s.seqs_from_uids(user, mailbox, uids).unwrap_or_default(),
        }
    }

    /// Reset all state (users, mailboxes, messages).
    pub fn reset(&mut self) {
        match self {
            Self::InMemory(s) => s.reset(),
            Self::Sqlite(s) => {
                let _ = s.reset();
            }
        }
    }

    pub(crate) fn snapshot(&mut self) -> StoreSnapshot {
        let users = self.list_users();
        let mut mailboxes = HashMap::new();
        let mut subscriptions = HashMap::new();
        let mut messages = HashMap::new();

        for user in &users {
            let boxes = self.list_mailboxes(user);
            mailboxes.insert(user.clone(), boxes.clone());
            let subs = self.list_subscriptions(user);
            subscriptions.insert(user.clone(), subs);
            for mailbox in boxes {
                let list = self.list(user, &mailbox);
                messages.insert((user.clone(), mailbox), list);
            }
        }

        StoreSnapshot {
            users,
            mailboxes,
            subscriptions,
            messages,
        }
    }

    pub(crate) fn restore_snapshot(&mut self, snapshot: &StoreSnapshot) {
        self.reset();

        for user in &snapshot.users {
            let mailboxes = snapshot.mailboxes.get(user).cloned().unwrap_or_default();
            for mailbox in mailboxes {
                self.ensure_mailbox(user, &mailbox);
                if let Some(items) = snapshot.messages.get(&(user.clone(), mailbox.clone())) {
                    for message in items {
                        let flags = FlagSet {
                            seen: message.seen,
                            flagged: message.flagged,
                            deleted: message.deleted,
                            answered: message.answered,
                            draft: message.draft,
                        };
                        self.append_with_flags(
                            user,
                            &mailbox,
                            message.data.clone(),
                            message.internal_date.clone(),
                            &flags,
                        );
                    }
                }
            }
            let subs = snapshot
                .subscriptions
                .get(user)
                .cloned()
                .unwrap_or_default();
            for mailbox in subs {
                let _ = self.subscribe(user, &mailbox);
            }
        }
    }

    /// Purge all messages but keep users and mailboxes.
    pub fn purge_messages(&mut self) {
        match self {
            Self::InMemory(s) => s.purge_messages(),
            Self::Sqlite(s) => {
                let _ = s.purge_messages();
            }
        }
    }

    /// List all users.
    pub fn list_users(&self) -> Vec<String> {
        match self {
            Self::InMemory(s) => s.list_users(),
            Self::Sqlite(s) => s.list_users().unwrap_or_default(),
        }
    }

    /// Delete a message by UID.
    pub fn delete_by_uid(&mut self, user: &str, mailbox: &str, uid: u32) {
        match self {
            Self::InMemory(s) => s.delete_by_uid(user, mailbox, uid),
            Self::Sqlite(s) => {
                let _ = s.delete_by_uid(user, mailbox, uid);
            }
        }
    }
}

#[derive(Clone, Debug)]
pub struct Message {
    pub uid: u32,
    pub data: Vec<u8>,
    pub internal_date: String,
    pub seen: bool,
    pub flagged: bool,
    pub deleted: bool,
    pub answered: bool,
    pub draft: bool,
}

pub const DEFAULT_INTERNAL_DATE: &str = "01-Jan-2020 00:00:00 +0000";

/// Return the current timestamp formatted for IMAP INTERNALDATE.
pub fn current_internal_date() -> String {
    format_internal_date(OffsetDateTime::now_utc())
}

fn format_internal_date(date: OffsetDateTime) -> String {
    static FORMAT: OnceLock<Vec<format_description::FormatItem<'static>>> = OnceLock::new();
    let format = FORMAT.get_or_init(|| {
        format_description::parse(
            "[day]-[month repr:short]-[year] [hour]:[minute]:[second] [offset_hour sign:mandatory][offset_minute]",
        )
        .expect("internal date format should be valid")
    });
    date.format(format)
        .unwrap_or_else(|_| DEFAULT_INTERNAL_DATE.to_string())
}

#[derive(Clone, Copy, Debug)]
pub enum FlagOp {
    Add,
    Remove,
    Replace,
}

#[derive(Default, Clone, Debug)]
pub struct FlagSet {
    pub seen: bool,
    pub flagged: bool,
    pub deleted: bool,
    pub answered: bool,
    pub draft: bool,
}

#[derive(Default)]
pub struct InMemoryStore {
    pub users: HashMap<String, HashMap<String, Vec<Message>>>,
    pub subscriptions: HashMap<String, HashSet<String>>,
}

impl InMemoryStore {
    pub fn new() -> Self {
        Self {
            users: HashMap::new(),
            subscriptions: HashMap::new(),
        }
    }

    pub fn ensure_mailbox(&mut self, user: &str, mailbox: &str) {
        let mailboxes = self.users.entry(user.to_string()).or_default();
        mailboxes.entry(mailbox.to_string()).or_default();
    }

    pub fn list_mailboxes(&mut self, user: &str) -> Vec<String> {
        self.ensure_mailbox(user, "INBOX");
        let mailboxes = self.users.entry(user.to_string()).or_default();
        let mut names: Vec<String> = mailboxes.keys().cloned().collect();
        names.sort();
        names
    }

    pub fn create_mailbox(&mut self, user: &str, mailbox: &str) -> bool {
        if mailbox.eq_ignore_ascii_case("INBOX") {
            return false;
        }
        let mailboxes = self.users.entry(user.to_string()).or_default();
        if mailboxes.contains_key(mailbox) {
            return false;
        }
        mailboxes.insert(mailbox.to_string(), Vec::new());
        true
    }

    pub fn rename_mailbox(&mut self, user: &str, old: &str, new: &str) -> bool {
        let mailboxes = self.users.entry(user.to_string()).or_default();
        if !mailboxes.contains_key(old) || mailboxes.contains_key(new) {
            return false;
        }
        if let Some(messages) = mailboxes.remove(old) {
            mailboxes.insert(new.to_string(), messages);
            if let Some(subs) = self.subscriptions.get_mut(user)
                && subs.remove(old)
            {
                subs.insert(new.to_string());
            }
            return true;
        }
        false
    }

    pub fn delete_mailbox(&mut self, user: &str, mailbox: &str) -> bool {
        if mailbox.eq_ignore_ascii_case("INBOX") {
            return false;
        }
        let mailboxes = self.users.entry(user.to_string()).or_default();
        let removed = mailboxes.remove(mailbox).is_some();
        if removed && let Some(subs) = self.subscriptions.get_mut(user) {
            subs.remove(mailbox);
        }
        removed
    }

    pub fn subscribe(&mut self, user: &str, mailbox: &str) -> bool {
        let exists = self
            .users
            .get(user)
            .and_then(|boxes| boxes.get(mailbox))
            .is_some()
            || mailbox.eq_ignore_ascii_case("INBOX");
        if !exists && !mailbox.eq_ignore_ascii_case("INBOX") {
            return false;
        }
        let subs = self.subscriptions.entry(user.to_string()).or_default();
        subs.insert(mailbox.to_string());
        true
    }

    pub fn unsubscribe(&mut self, user: &str, mailbox: &str) -> bool {
        let subs = self.subscriptions.entry(user.to_string()).or_default();
        subs.remove(mailbox)
    }

    pub fn list_subscriptions(&mut self, user: &str) -> Vec<String> {
        let subs = self.subscriptions.entry(user.to_string()).or_default();
        let mut names: Vec<String> = subs.iter().cloned().collect();
        names.sort();
        names
    }

    pub fn append(
        &mut self,
        user: &str,
        mailbox: &str,
        data: Vec<u8>,
        internal_date: String,
    ) -> u32 {
        self.ensure_mailbox(user, mailbox);
        let mailboxes = self.users.entry(user.to_string()).or_default();
        let box_ref = mailboxes.entry(mailbox.to_string()).or_default();
        let uid = u32::try_from(box_ref.len())
            .unwrap_or(u32::MAX)
            .saturating_add(1);
        box_ref.push(Message {
            uid,
            data,
            internal_date,
            seen: false,
            flagged: false,
            deleted: false,
            answered: false,
            draft: false,
        });
        uid
    }

    pub fn append_with_flags(
        &mut self,
        user: &str,
        mailbox: &str,
        data: Vec<u8>,
        internal_date: String,
        flags: &FlagSet,
    ) -> u32 {
        self.ensure_mailbox(user, mailbox);
        let mailboxes = self.users.entry(user.to_string()).or_default();
        let box_ref = mailboxes.entry(mailbox.to_string()).or_default();
        let uid = u32::try_from(box_ref.len())
            .unwrap_or(u32::MAX)
            .saturating_add(1);
        box_ref.push(Message {
            uid,
            data,
            internal_date,
            seen: flags.seen,
            flagged: flags.flagged,
            deleted: flags.deleted,
            answered: flags.answered,
            draft: flags.draft,
        });
        uid
    }

    pub fn list(&mut self, user: &str, mailbox: &str) -> Vec<Message> {
        self.ensure_mailbox(user, mailbox);
        self.users
            .get(user)
            .and_then(|boxes| boxes.get(mailbox).cloned())
            .unwrap_or_default()
    }

    pub fn apply_flags_by_seq(
        &mut self,
        user: &str,
        mailbox: &str,
        seq: u32,
        op: FlagOp,
        flags: &FlagSet,
    ) {
        let index = seq.saturating_sub(1) as usize;
        if let Some(messages) = self
            .users
            .get_mut(user)
            .and_then(|boxes| boxes.get_mut(mailbox))
            && let Some(message) = messages.get_mut(index)
        {
            apply_flags(message, op, flags);
        }
    }

    pub fn apply_flags_by_uid(
        &mut self,
        user: &str,
        mailbox: &str,
        uid: u32,
        op: FlagOp,
        flags: &FlagSet,
    ) {
        if let Some(messages) = self
            .users
            .get_mut(user)
            .and_then(|boxes| boxes.get_mut(mailbox))
        {
            for message in messages {
                if message.uid == uid {
                    apply_flags(message, op, flags);
                }
            }
        }
    }

    pub fn expunge_deleted(&mut self, user: &str, mailbox: &str) -> Vec<u32> {
        let mut expunged = Vec::new();
        if let Some(messages) = self
            .users
            .get_mut(user)
            .and_then(|boxes| boxes.get_mut(mailbox))
        {
            let mut i = 0;
            while i < messages.len() {
                if messages[i].deleted {
                    if let Ok(seq) = u32::try_from(i + 1) {
                        expunged.push(seq);
                    }
                    messages.remove(i);
                } else {
                    i += 1;
                }
            }
        }
        expunged
    }

    pub fn expunge_deleted_by_uid(&mut self, user: &str, mailbox: &str, uids: &[u32]) -> Vec<u32> {
        let mut expunged = Vec::new();
        if let Some(messages) = self
            .users
            .get_mut(user)
            .and_then(|boxes| boxes.get_mut(mailbox))
        {
            let mut i = 0;
            while i < messages.len() {
                if messages[i].deleted && uids.contains(&messages[i].uid) {
                    if let Ok(seq) = u32::try_from(i + 1) {
                        expunged.push(seq);
                    }
                    messages.remove(i);
                } else {
                    i += 1;
                }
            }
        }
        expunged
    }

    pub fn copy_by_seq_set(&mut self, user: &str, src: &str, seqs: &[u32], dest: &str) -> usize {
        self.ensure_mailbox(user, dest);
        let source_messages = self
            .users
            .get(user)
            .and_then(|boxes| boxes.get(src).cloned())
            .unwrap_or_default();
        let mailboxes = self.users.entry(user.to_string()).or_default();
        let dest_box = mailboxes.entry(dest.to_string()).or_default();
        let mut copied = 0;
        for seq in seqs {
            let index = seq.saturating_sub(1) as usize;
            if let Some(message) = source_messages.get(index) {
                let uid = u32::try_from(dest_box.len())
                    .unwrap_or(u32::MAX)
                    .saturating_add(1);
                dest_box.push(Message {
                    uid,
                    data: message.data.clone(),
                    internal_date: message.internal_date.clone(),
                    seen: message.seen,
                    flagged: message.flagged,
                    deleted: false,
                    answered: message.answered,
                    draft: message.draft,
                });
                copied += 1;
            }
        }
        copied
    }

    pub fn move_by_seq_set(&mut self, user: &str, src: &str, seqs: &[u32], dest: &str) -> Vec<u32> {
        let mut expunged = Vec::new();
        self.copy_by_seq_set(user, src, seqs, dest);
        if let Some(messages) = self
            .users
            .get_mut(user)
            .and_then(|boxes| boxes.get_mut(src))
        {
            let mut indices: Vec<usize> = seqs
                .iter()
                .filter_map(|seq| seq.checked_sub(1).map(|v| v as usize))
                .collect();
            indices.sort_by(|a, b| b.cmp(a));
            for idx in indices {
                if idx < messages.len() {
                    if let Ok(seq) = u32::try_from(idx + 1) {
                        expunged.push(seq);
                    }
                    messages.remove(idx);
                }
            }
        }
        expunged
    }

    pub fn seqs_from_uids(&self, user: &str, mailbox: &str, uids: &[u32]) -> Vec<u32> {
        let mut seqs = Vec::new();
        if let Some(messages) = self.users.get(user).and_then(|boxes| boxes.get(mailbox)) {
            for uid in uids {
                if let Some(pos) = messages.iter().position(|m| m.uid == *uid)
                    && let Ok(seq) = u32::try_from(pos + 1)
                {
                    seqs.push(seq);
                }
            }
        }
        seqs
    }

    /// Reset all state.
    pub fn reset(&mut self) {
        self.users.clear();
        self.subscriptions.clear();
    }

    /// Purge all messages but keep users and mailboxes.
    pub fn purge_messages(&mut self) {
        for mailboxes in self.users.values_mut() {
            for messages in mailboxes.values_mut() {
                messages.clear();
            }
        }
    }

    /// List all users.
    pub fn list_users(&self) -> Vec<String> {
        self.users.keys().cloned().collect()
    }

    /// Delete a message by UID.
    pub fn delete_by_uid(&mut self, user: &str, mailbox: &str, uid: u32) {
        if let Some(messages) = self
            .users
            .get_mut(user)
            .and_then(|boxes| boxes.get_mut(mailbox))
        {
            messages.retain(|m| m.uid != uid);
        }
    }
}

/// Convert a UID to a sequence number for the given mailbox.
pub fn uid_to_seq(store: &Store, user: &str, mailbox: &str, uid: u32) -> Option<u32> {
    match store {
        Store::InMemory(s) => s
            .users
            .get(user)
            .and_then(|boxes| boxes.get(mailbox))
            .and_then(|messages| {
                messages
                    .iter()
                    .position(|m| m.uid == uid)
                    .and_then(|idx| u32::try_from(idx + 1).ok())
            }),
        Store::Sqlite(s) => s.list(user, mailbox).ok().and_then(|messages| {
            messages
                .iter()
                .position(|m| m.uid == uid)
                .and_then(|idx| u32::try_from(idx + 1).ok())
        }),
    }
}

fn apply_flags(message: &mut Message, op: FlagOp, flags: &FlagSet) {
    match op {
        FlagOp::Add => {
            if flags.seen {
                message.seen = true;
            }
            if flags.flagged {
                message.flagged = true;
            }
            if flags.deleted {
                message.deleted = true;
            }
            if flags.answered {
                message.answered = true;
            }
            if flags.draft {
                message.draft = true;
            }
        }
        FlagOp::Remove => {
            if flags.seen {
                message.seen = false;
            }
            if flags.flagged {
                message.flagged = false;
            }
            if flags.deleted {
                message.deleted = false;
            }
            if flags.answered {
                message.answered = false;
            }
            if flags.draft {
                message.draft = false;
            }
        }
        FlagOp::Replace => {
            message.seen = flags.seen;
            message.flagged = flags.flagged;
            message.deleted = flags.deleted;
            message.answered = flags.answered;
            message.draft = flags.draft;
        }
    }
}