unifier/

store.rs

//! In-memory hot store with dirty tracking; flush writes only changed data to disk.

use std::collections::{BTreeMap, BTreeSet, VecDeque};
use std::fs;
use std::path::{Path, PathBuf};

use uuid::Uuid;

use crate::constants::{CRON, KEYS, MAILBOX};
use crate::cron::CronSchedule;
use crate::error::{Error, Result};
use crate::fs_text::{read_text, write_text, write_text_atomic};
use crate::home::UnifierHome;
use crate::paths::{cron_dir, key_path, mailbox_dir, message_path, parse_message_id};
use crate::postbox::Message;
use crate::scope::resolve_under_root;

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum KeyState {
    Present { value: String, dirty: bool },
    Deleted { dirty: bool },
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct MsgState {
    body: String,
    dirty: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
enum MessageKind {
    Mailbox,
    Cron,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct EventState {
    pub(crate) body: String,
    pub(crate) dirty: bool,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum StagingValue {
    Present(String),
    Deleted,
}

/// Active tick: reads frozen snapshot, writes go to staging only.
#[derive(Debug, Clone)]
pub(crate) struct ActiveTick {
    pub number: u64,
    pub(crate) read_snapshot: BTreeMap<String, String>,
    pub(crate) staging: BTreeMap<String, StagingValue>,
    pub(crate) locks: BTreeSet<String>,
}

/// In-memory mirror of the postbox; disk is touched only on [`HotStore::flush`].
#[derive(Debug, Default)]
pub struct HotStore {
    pub(crate) keys: BTreeMap<String, KeyState>,
    mailboxes: BTreeMap<String, BTreeMap<Uuid, MsgState>>,
    cron: BTreeMap<String, BTreeMap<Uuid, MsgState>>,
    pub(crate) events: BTreeMap<Uuid, EventState>,
    removed_messages: BTreeSet<(MessageKind, String, Uuid)>,
    pub(crate) committed_tick: u64,
    pub(crate) active_tick: Option<ActiveTick>,
    pub(crate) tick_queue: VecDeque<String>,
}

impl HotStore {
    pub fn load(home: &UnifierHome) -> Result<Self> {
        let mut store = Self::default();
        store.load_keys(home)?;
        store.mailboxes = load_message_buckets(home, MAILBOX)?;
        store.cron = load_message_buckets(home, CRON)?;
        store.events = crate::tick::load_events(home)?;
        store.committed_tick = crate::tick::load_committed_tick(home)?;
        Ok(store)
    }

    pub fn is_dirty(&self) -> bool {
        self.keys.values().any(|k| match k {
            KeyState::Present { dirty, .. } | KeyState::Deleted { dirty } => *dirty,
        }) || self.has_dirty_messages(&self.mailboxes)
            || self.has_dirty_messages(&self.cron)
            || self.events.values().any(|e| e.dirty)
            || !self.removed_messages.is_empty()
            || self.active_tick.is_some()
    }

    pub fn put_key(&mut self, key: &str, value: &str) -> Result<()> {
        validate_key(key)?;
        if let Some(tick) = &self.active_tick {
            if tick.locks.contains(key) {
                return Err(Error::msg(format!(
                    "key locked during tick {}: {key}",
                    tick.number
                )));
            }
        }
        if let Some(tick) = &mut self.active_tick {
            tick.staging
                .insert(key.to_string(), StagingValue::Present(value.to_string()));
            return Ok(());
        }
        self.keys.insert(
            key.to_string(),
            KeyState::Present {
                value: value.to_string(),
                dirty: true,
            },
        );
        Ok(())
    }

    pub fn get_key(&self, key: &str) -> Result<Option<String>> {
        validate_key(key)?;
        if let Some(tick) = &self.active_tick {
            if let Some(staged) = tick.staging.get(key) {
                return Ok(match staged {
                    StagingValue::Present(v) => Some(v.clone()),
                    StagingValue::Deleted => None,
                });
            }
            return Ok(tick.read_snapshot.get(key).cloned());
        }
        Ok(match self.keys.get(key) {
            Some(KeyState::Present { value, .. }) => Some(value.clone()),
            Some(KeyState::Deleted { .. }) | None => None,
        })
    }

    pub fn delete_key(&mut self, key: &str) -> Result<bool> {
        validate_key(key)?;
        if let Some(tick) = &self.active_tick {
            if tick.locks.contains(key) {
                return Err(Error::msg(format!(
                    "key locked during tick {}: {key}",
                    tick.number
                )));
            }
        }
        if let Some(tick) = &mut self.active_tick {
            let existed = tick.staging.contains_key(key) || tick.read_snapshot.contains_key(key);
            if !existed {
                return Ok(false);
            }
            tick.staging.insert(key.to_string(), StagingValue::Deleted);
            return Ok(true);
        }
        let existed = matches!(self.keys.get(key), Some(KeyState::Present { .. }));
        if !existed {
            return Ok(false);
        }
        self.keys
            .insert(key.to_string(), KeyState::Deleted { dirty: true });
        Ok(true)
    }

    pub fn send(&mut self, recipient: &str, body: &str) -> Result<Uuid> {
        self.send_from(crate::envelope::DEFAULT_SENDER, recipient, body)
    }

    pub fn send_from(&mut self, from: &str, recipient: &str, body: &str) -> Result<Uuid> {
        validate_segment(from, "from")?;
        validate_segment(recipient, "recipient")?;
        let env = crate::envelope::Envelope::new(
            from,
            recipient,
            crate::envelope::Envelope::parse_payload(body),
        );
        let id = env.id;
        self.mailboxes
            .entry(recipient.to_string())
            .or_default()
            .insert(
                id,
                MsgState {
                    body: env.to_json()?,
                    dirty: true,
                },
            );
        Ok(id)
    }

    pub fn post_cron(&mut self, schedule: &str, body: &str) -> Result<Uuid> {
        CronSchedule::parse(schedule)?;
        let id = Uuid::new_v4();
        self.cron.entry(schedule.to_string()).or_default().insert(
            id,
            MsgState {
                body: body.to_string(),
                dirty: true,
            },
        );
        Ok(id)
    }

    pub fn poll_mailbox(&self, recipient: &str) -> Result<Vec<Message>> {
        validate_segment(recipient, "recipient")?;
        Ok(self.collect_messages(
            MessageKind::Mailbox,
            recipient,
            self.mailboxes.get(recipient),
        ))
    }

    pub fn poll_cron(&self) -> Result<Vec<Message>> {
        let mut out = Vec::new();
        for (schedule, msgs) in &self.cron {
            let parsed = CronSchedule::parse(schedule)?;
            if !parsed.matches_now() {
                continue;
            }
            out.extend(self.collect_messages(MessageKind::Cron, schedule, Some(msgs)));
        }
        out.sort_by(|a, b| a.path.cmp(&b.path));
        Ok(out)
    }

    pub fn list_dir(&self, home: &UnifierHome, subpath: &str) -> Result<Vec<Message>> {
        let dir = resolve_under_root(home.path(), subpath)?;
        let rel = dir
            .strip_prefix(home.path())
            .map_err(|_| Error::msg("path escapes store root"))?;
        let parts: Vec<_> = rel.iter().collect();
        match parts.as_slice() {
            [p1, p2] if p1.to_string_lossy() == MAILBOX => self.poll_mailbox(&p2.to_string_lossy()),
            [p1, p2] if p1.to_string_lossy() == CRON => {
                let schedule = p2.to_string_lossy();
                Ok(self.collect_messages(
                    MessageKind::Cron,
                    &schedule,
                    self.cron.get(schedule.as_ref()),
                ))
            }
            _ => Ok(Vec::new()),
        }
    }

    pub fn ack(&mut self, home: &UnifierHome, id_or_path: &str) -> Result<bool> {
        if id_or_path.contains('/') || Path::new(id_or_path).is_absolute() {
            let path = if Path::new(id_or_path).is_absolute() {
                PathBuf::from(id_or_path)
            } else {
                resolve_under_root(home.path(), id_or_path)?
            };
            return self.ack_path(home, &path);
        }
        let id = Uuid::parse_str(id_or_path)
            .map_err(|_| Error::msg(format!("invalid message id or path: {id_or_path}")))?;
        self.ack_id(home, &id)
    }

    pub fn flush(&mut self, home: &UnifierHome) -> Result<()> {
        if self.active_tick.is_some() {
            return Err(Error::msg(
                "cannot flush while a tick is active; run tick end first",
            ));
        }
        home.ensure()?;
        self.flush_keys(home)?;
        self.flush_mailboxes(home)?;
        self.flush_cron(home)?;
        self.flush_events(home)?;
        self.removed_messages.clear();
        Ok(())
    }

    fn load_keys(&mut self, home: &UnifierHome) -> Result<()> {
        let root = home.path().join(KEYS);
        if !root.is_dir() {
            return Ok(());
        }
        let mut prefix = Vec::new();
        self.walk_key_files(&root, &mut prefix)?;
        Ok(())
    }

    fn walk_key_files(&mut self, dir: &Path, prefix: &mut Vec<String>) -> Result<()> {
        for entry in fs::read_dir(dir)? {
            let entry = entry?;
            let name = entry.file_name().to_string_lossy().into_owned();
            if entry.file_type()?.is_dir() {
                prefix.push(name);
                self.walk_key_files(&entry.path(), prefix)?;
                prefix.pop();
            } else if entry.file_type()?.is_file() {
                prefix.push(name);
                let key = prefix.join("/");
                prefix.pop();
                let value = read_text(&entry.path())?;
                self.keys.insert(
                    key,
                    KeyState::Present {
                        value,
                        dirty: false,
                    },
                );
            }
        }
        Ok(())
    }

    fn flush_keys(&mut self, home: &UnifierHome) -> Result<()> {
        let mut to_remove = Vec::new();
        for (key, state) in &mut self.keys {
            match state {
                KeyState::Present { value, dirty: true } => {
                    write_text_atomic(&key_path(home, key), value)?;
                    *state = KeyState::Present {
                        value: value.clone(),
                        dirty: false,
                    };
                }
                KeyState::Deleted { dirty: true } => {
                    let path = key_path(home, key);
                    if path.is_file() {
                        fs::remove_file(&path)?;
                    }
                    to_remove.push(key.clone());
                }
                _ => {}
            }
        }
        for key in to_remove {
            self.keys.remove(&key);
        }
        Ok(())
    }

    fn flush_mailboxes(&mut self, home: &UnifierHome) -> Result<()> {
        flush_message_buckets(
            home,
            MessageKind::Mailbox,
            &mut self.mailboxes,
            &mut self.removed_messages,
            mailbox_dir,
        )
    }

    fn flush_cron(&mut self, home: &UnifierHome) -> Result<()> {
        flush_message_buckets(
            home,
            MessageKind::Cron,
            &mut self.cron,
            &mut self.removed_messages,
            cron_dir,
        )
    }

    fn ack_path(&mut self, home: &UnifierHome, path: &Path) -> Result<bool> {
        let rel = path.strip_prefix(home.path()).ok();
        let Some(rel) = rel else {
            return Ok(false);
        };
        let parts: Vec<_> = rel
            .iter()
            .map(|p| p.to_string_lossy().into_owned())
            .collect();
        if parts.len() != 3 {
            return Ok(false);
        }
        let Some(id) = parse_message_id(&parts[2]) else {
            return Ok(false);
        };
        let existed = match parts[0].as_str() {
            MAILBOX => {
                let had = self
                    .mailboxes
                    .get(&parts[1])
                    .is_some_and(|m| m.contains_key(&id))
                    || path.is_file();
                if had {
                    self.removed_messages
                        .insert((MessageKind::Mailbox, parts[1].clone(), id));
                    if let Some(msgs) = self.mailboxes.get_mut(&parts[1]) {
                        msgs.remove(&id);
                    }
                }
                had
            }
            CRON => {
                let had = self
                    .cron
                    .get(&parts[1])
                    .is_some_and(|m| m.contains_key(&id))
                    || path.is_file();
                if had {
                    self.removed_messages
                        .insert((MessageKind::Cron, parts[1].clone(), id));
                    if let Some(msgs) = self.cron.get_mut(&parts[1]) {
                        msgs.remove(&id);
                    }
                }
                had
            }
            _ => false,
        };
        Ok(existed)
    }

    fn ack_id(&mut self, home: &UnifierHome, id: &Uuid) -> Result<bool> {
        if let Some(recipient) = self
            .mailboxes
            .iter()
            .find_map(|(name, msgs)| msgs.contains_key(id).then(|| name.clone()))
        {
            self.removed_messages
                .insert((MessageKind::Mailbox, recipient.clone(), *id));
            if let Some(m) = self.mailboxes.get_mut(&recipient) {
                m.remove(id);
            }
            return Ok(true);
        }
        if let Some(schedule) = self
            .cron
            .iter()
            .find_map(|(name, msgs)| msgs.contains_key(id).then(|| name.clone()))
        {
            self.removed_messages
                .insert((MessageKind::Cron, schedule.clone(), *id));
            if let Some(m) = self.cron.get_mut(&schedule) {
                m.remove(id);
            }
            return Ok(true);
        }
        // Message may exist on disk but not loaded into a bucket (empty dir walk).
        let filename = format!("{}.txt", id.hyphenated());
        for sub in [MAILBOX, CRON] {
            let base = home.path().join(sub);
            if !base.is_dir() {
                continue;
            }
            for entry in fs::read_dir(&base)? {
                let entry = entry?;
                if entry.file_type()?.is_dir() {
                    let candidate = entry.path().join(&filename);
                    if candidate.is_file() {
                        let bucket = entry.file_name().to_string_lossy().into_owned();
                        let kind = if sub == MAILBOX {
                            MessageKind::Mailbox
                        } else {
                            MessageKind::Cron
                        };
                        self.removed_messages.insert((kind, bucket, *id));
                        return Ok(true);
                    }
                }
            }
        }
        Ok(false)
    }

    fn collect_messages(
        &self,
        kind: MessageKind,
        bucket: &str,
        msgs: Option<&BTreeMap<Uuid, MsgState>>,
    ) -> Vec<Message> {
        let Some(msgs) = msgs else {
            return Vec::new();
        };
        let base = match kind {
            MessageKind::Mailbox => PathBuf::from(MAILBOX).join(bucket),
            MessageKind::Cron => PathBuf::from(CRON).join(bucket),
        };
        let mut out = Vec::new();
        for (id, msg) in msgs {
            out.push(Message {
                id: *id,
                path: base.join(format!("{}.txt", id.hyphenated())),
                body: msg.body.clone(),
            });
        }
        out.sort_by_key(|m| m.id);
        out
    }

    fn has_dirty_messages(&self, buckets: &BTreeMap<String, BTreeMap<Uuid, MsgState>>) -> bool {
        buckets.values().any(|msgs| msgs.values().any(|m| m.dirty))
    }
}

fn load_message_buckets(
    home: &UnifierHome,
    top: &str,
) -> Result<BTreeMap<String, BTreeMap<Uuid, MsgState>>> {
    let mut target = BTreeMap::new();
    let base = home.path().join(top);
    if !base.is_dir() {
        return Ok(target);
    }
    for entry in fs::read_dir(&base)? {
        let entry = entry?;
        if !entry.file_type()?.is_dir() {
            continue;
        }
        let bucket = entry.file_name().to_string_lossy().into_owned();
        let mut msgs = BTreeMap::new();
        for msg_entry in fs::read_dir(entry.path())? {
            let msg_entry = msg_entry?;
            if !msg_entry.file_type()?.is_file() {
                continue;
            }
            let name = msg_entry.file_name().to_string_lossy().into_owned();
            let Some(id) = parse_message_id(&name) else {
                continue;
            };
            msgs.insert(
                id,
                MsgState {
                    body: read_text(&msg_entry.path())?,
                    dirty: false,
                },
            );
        }
        if !msgs.is_empty() {
            target.insert(bucket, msgs);
        }
    }
    Ok(target)
}

fn flush_message_buckets<F>(
    home: &UnifierHome,
    kind: MessageKind,
    buckets: &mut BTreeMap<String, BTreeMap<Uuid, MsgState>>,
    removed_messages: &mut BTreeSet<(MessageKind, String, Uuid)>,
    dir_for: F,
) -> Result<()>
where
    F: Fn(&UnifierHome, &str) -> PathBuf,
{
    for (bucket, msgs) in buckets.iter_mut() {
        for (id, msg) in msgs.iter_mut() {
            if msg.dirty {
                write_text(&message_path(&dir_for(home, bucket), id), &msg.body)?;
                msg.dirty = false;
            }
        }
    }
    let removed: Vec<_> = removed_messages
        .iter()
        .filter(|(k, _, _)| *k == kind)
        .cloned()
        .collect();
    for (_, bucket, id) in removed {
        let path = message_path(&dir_for(home, &bucket), &id);
        if path.is_file() {
            fs::remove_file(path)?;
        }
        if let Some(msgs) = buckets.get_mut(&bucket) {
            msgs.remove(&id);
        }
        removed_messages.remove(&(kind, bucket.clone(), id));
    }
    buckets.retain(|_, msgs| !msgs.is_empty());
    Ok(())
}

pub(crate) fn validate_key(key: &str) -> Result<()> {
    if key.is_empty() {
        return Err(Error::msg("key must not be empty"));
    }
    if key.contains("..") {
        return Err(Error::msg("key must not contain '..'"));
    }
    Ok(())
}

pub(crate) fn validate_segment(segment: &str, label: &str) -> Result<()> {
    if segment.is_empty() || segment.contains('/') || segment.contains("..") {
        return Err(Error::msg(format!("invalid {label}: {segment}")));
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::home::UnifierHome;
    use tempfile::tempdir;

    #[test]
    fn put_get_without_flush_leaves_disk_clean() {
        let tmp = tempdir().unwrap();
        let home = UnifierHome::resolve(Some(tmp.path().to_path_buf()), None).unwrap();
        let mut store = HotStore::load(&home).unwrap();

        store.put_key("app/theme", "dark").unwrap();
        assert_eq!(store.get_key("app/theme").unwrap(), Some("dark".into()));
        assert!(store.is_dirty());
        assert!(!home.path().join("keys/app/theme").exists());

        store.flush(&home).unwrap();
        assert!(!store.is_dirty());
        assert!(home.path().join("keys/app/theme").is_file());
    }

    #[test]
    fn send_without_flush_then_flush_persists() {
        let tmp = tempdir().unwrap();
        let home = UnifierHome::resolve(Some(tmp.path().to_path_buf()), None).unwrap();
        let mut store = HotStore::load(&home).unwrap();

        store.send("worker", "hello").unwrap();
        assert!(store.poll_mailbox("worker").unwrap().len() == 1);
        assert!(fs::read_dir(home.path().join("mailbox")).is_err());

        store.flush(&home).unwrap();
        assert!(home.path().join("mailbox/worker").is_dir());
    }

    #[test]
    fn load_existing_keys_from_disk() {
        let tmp = tempdir().unwrap();
        let home = UnifierHome::resolve(Some(tmp.path().to_path_buf()), None).unwrap();
        put_key_fs(&home, "k", "v").unwrap();

        let store = HotStore::load(&home).unwrap();
        assert_eq!(store.get_key("k").unwrap(), Some("v".into()));
    }

    fn put_key_fs(home: &UnifierHome, key: &str, value: &str) -> Result<()> {
        write_text_atomic(&key_path(home, key), value)
    }
}