kiromi-ai-memory 0.2.2

// SPDX-License-Identifier: Apache-2.0 OR MIT
//! Partition scheme parsing, partition-path encoding, call-site `Partitions` builder.

use std::collections::BTreeMap;
use std::fmt;
use std::str::FromStr;

use percent_encoding::{AsciiSet, CONTROLS, utf8_percent_encode};
use serde::{Deserialize, Serialize};

/// Characters that must be percent-encoded inside a partition value.
/// We keep `=`, `/`, and control chars off-limits; everything else passes through.
const PARTITION_VALUE_BAD: &AsciiSet = &CONTROLS.add(b'=').add(b'/').add(b' ').add(b'\\');

/// A user-declared partition layout, e.g. `user={user}/year={year}`.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PartitionScheme {
    keys: Vec<String>,
}

impl PartitionScheme {
    /// Parse a scheme template. Each level is `KEY={KEY}` separated by `/`.
    pub fn parse(template: &str) -> Result<Self, InvalidScheme> {
        if template.is_empty() {
            return Err(InvalidScheme::Empty);
        }
        let mut keys = Vec::new();
        for level in template.split('/') {
            let (k, placeholder) =
                level
                    .split_once('=')
                    .ok_or_else(|| InvalidScheme::MissingEquals {
                        level: level.to_string(),
                    })?;
            if !placeholder.starts_with('{') || !placeholder.ends_with('}') {
                return Err(InvalidScheme::BadPlaceholder {
                    level: level.to_string(),
                });
            }
            let inner = &placeholder[1..placeholder.len() - 1];
            if k != inner {
                return Err(InvalidScheme::PlaceholderMismatch {
                    key: k.to_string(),
                    placeholder: inner.to_string(),
                });
            }
            if !is_valid_key(k) {
                return Err(InvalidScheme::BadKey { key: k.to_string() });
            }
            keys.push(k.to_string());
        }
        Ok(PartitionScheme { keys })
    }

    /// Number of levels.
    #[must_use]
    pub fn depth(&self) -> usize {
        self.keys.len()
    }

    /// Borrow the keys in order.
    #[must_use]
    pub fn keys(&self) -> &[String] {
        &self.keys
    }
}

impl fmt::Display for PartitionScheme {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (i, k) in self.keys.iter().enumerate() {
            if i > 0 {
                f.write_str("/")?;
            }
            write!(f, "{k}={{{k}}}")?;
        }
        Ok(())
    }
}

fn is_valid_key(k: &str) -> bool {
    !k.is_empty()
        && k.chars()
            .next()
            .is_some_and(|c| c.is_ascii_alphabetic() || c == '_')
        && k.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
}

/// Errors produced when parsing a `PartitionScheme`.
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum InvalidScheme {
    /// Template was empty.
    #[error("partition scheme is empty")]
    Empty,
    /// A level is missing its `=`.
    #[error("partition scheme level missing '=': {level:?}")]
    MissingEquals {
        /// The offending level.
        level: String,
    },
    /// Placeholder is not surrounded by `{ }`.
    #[error("partition scheme placeholder must be {{...}}: {level:?}")]
    BadPlaceholder {
        /// The offending level.
        level: String,
    },
    /// Key on the LHS doesn't match the placeholder name.
    #[error("partition scheme placeholder {placeholder:?} does not match key {key:?}")]
    PlaceholderMismatch {
        /// LHS key.
        key: String,
        /// Placeholder name.
        placeholder: String,
    },
    /// Key contains a disallowed char.
    #[error("partition scheme key {key:?} contains disallowed character")]
    BadKey {
        /// The offending key.
        key: String,
    },
}

/// Concrete call-site partition values, e.g. `[("user","alex"),("year","2026")]`.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Partitions {
    values: BTreeMap<String, String>,
}

impl Partitions {
    /// Empty.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Add a key/value level.
    #[must_use]
    pub fn with(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.values.insert(key.into(), value.into());
        self
    }

    /// Convert a parsed `PartitionPath` back into a `Partitions` builder.
    ///
    /// The path string is the canonical `key=value/key=value` form; values are
    /// percent-decoded so the round-trip through `resolve()` is lossless for
    /// any value the slice-1 charset allows.
    #[must_use]
    pub fn from_path(p: &PartitionPath) -> Self {
        let mut out = Partitions::new();
        for (k, v) in p.kv_pairs() {
            out = out.with(k.to_string(), v.into_owned());
        }
        out
    }

    /// Resolve against a scheme to get a fully-qualified `PartitionPath`. Fails
    /// if any scheme key is missing or any extra keys are present.
    pub fn resolve(self, scheme: &PartitionScheme) -> Result<PartitionPath, PartitionResolveError> {
        for k in scheme.keys() {
            if !self.values.contains_key(k) {
                return Err(PartitionResolveError::MissingKey { key: k.clone() });
            }
        }
        for k in self.values.keys() {
            if !scheme.keys().iter().any(|s| s == k) {
                return Err(PartitionResolveError::UnknownKey { key: k.clone() });
            }
        }
        let mut segments = Vec::with_capacity(scheme.depth());
        for k in scheme.keys() {
            let v = self
                .values
                .get(k)
                .ok_or_else(|| PartitionResolveError::MissingKey { key: k.clone() })?;
            if v.is_empty() {
                return Err(PartitionResolveError::EmptyValue { key: k.clone() });
            }
            let encoded = utf8_percent_encode(v, PARTITION_VALUE_BAD).to_string();
            segments.push(format!("{k}={encoded}"));
        }
        Ok(PartitionPath(segments.join("/")))
    }
}

impl<S1, S2> From<&[(S1, S2)]> for Partitions
where
    S1: AsRef<str>,
    S2: AsRef<str>,
{
    fn from(pairs: &[(S1, S2)]) -> Self {
        let mut p = Partitions::new();
        for (k, v) in pairs {
            p.values
                .insert(k.as_ref().to_string(), v.as_ref().to_string());
        }
        p
    }
}

impl<S1, S2, const N: usize> From<[(S1, S2); N]> for Partitions
where
    S1: AsRef<str>,
    S2: AsRef<str>,
{
    fn from(arr: [(S1, S2); N]) -> Self {
        let mut p = Partitions::new();
        for (k, v) in arr {
            p.values
                .insert(k.as_ref().to_string(), v.as_ref().to_string());
        }
        p
    }
}

/// Errors produced when resolving `Partitions` against a `PartitionScheme`.
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum PartitionResolveError {
    /// Scheme key wasn't supplied.
    #[error("missing partition key {key:?}")]
    MissingKey {
        /// The missing key.
        key: String,
    },
    /// Caller supplied a key the scheme doesn't know about.
    #[error("unknown partition key {key:?}")]
    UnknownKey {
        /// The extra key.
        key: String,
    },
    /// A value was empty.
    #[error("partition value for {key:?} is empty")]
    EmptyValue {
        /// The key whose value was empty.
        key: String,
    },
}

/// Fully-resolved partition path, e.g. `user=alex/year=2026/month=05/topic=meetings`.
#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(transparent)]
pub struct PartitionPath(String);

impl PartitionPath {
    /// Borrow as `&str`.
    #[must_use]
    pub fn as_str(&self) -> &str {
        &self.0
    }

    /// Number of levels in the path.
    #[must_use]
    pub fn depth(&self) -> usize {
        self.0.split('/').count()
    }

    /// Iterate `(key, percent-decoded value)` pairs for each level. Returns a
    /// `Cow<'_, str>` so values that don't need decoding don't allocate.
    pub fn kv_pairs(&self) -> impl Iterator<Item = (&str, std::borrow::Cow<'_, str>)> + '_ {
        self.0.split('/').filter(|s| !s.is_empty()).map(|seg| {
            let (k, v) = seg.split_once('=').unwrap_or((seg, ""));
            let decoded = percent_encoding::percent_decode_str(v).decode_utf8_lossy();
            (k, decoded)
        })
    }

    /// Iterate parent paths, deepest first (excluding `self`). Empty for top-level.
    pub fn ancestors(&self) -> impl Iterator<Item = PartitionPath> + '_ {
        let mut s = self.0.as_str();
        std::iter::from_fn(move || {
            let idx = s.rfind('/')?;
            s = &s[..idx];
            Some(PartitionPath(s.to_string()))
        })
    }

    /// Build a `PartitionPath` from a previously-rendered string.
    ///
    /// Performs minimal validation: rejects the empty path and requires every
    /// non-empty segment to contain `=`. Used by FFI consumers (e.g.
    /// `kiromi-ai-swift`) that round-trip a partition through an opaque
    /// string previously returned by a `MemoryRef`.
    pub fn from_string(s: &str) -> Result<Self, PartitionResolveError> {
        if s.is_empty() {
            return Err(PartitionResolveError::MissingKey {
                key: "<empty>".into(),
            });
        }
        for seg in s.split('/') {
            if seg.is_empty() || !seg.contains('=') {
                return Err(PartitionResolveError::MissingKey {
                    key: seg.to_string(),
                });
            }
        }
        Ok(Self(s.to_string()))
    }
}

impl fmt::Display for PartitionPath {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(&self.0)
    }
}

impl FromStr for PartitionPath {
    type Err = PartitionResolveError;
    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        // Trust callers / SQL — only validate non-empty.
        if s.is_empty() {
            return Err(PartitionResolveError::EmptyValue {
                key: String::from("<root>"),
            });
        }
        Ok(PartitionPath(s.to_string()))
    }
}

/// Sentinel `partition_path` used to represent the tenant-root pseudo-node in
/// SQL state (`summary.subject_kind = 'tenant'` rows, hierarchical descent
/// dispatch, etc.).
///
/// The angle brackets keep the sentinel unable to collide with a real
/// partition path: `Partitions::resolve` rejects them in user-supplied keys
/// and values.
pub const TENANT_ROOT_SENTINEL: &str = "<root>";

/// Build the sentinel [`PartitionPath`] for the tenant-root node.
///
/// `unwrap` is safe here because `FromStr` rejects only the empty string and
/// `<root>` is non-empty.
#[must_use]
#[allow(clippy::expect_used)] // sentinel is a hard-coded non-empty string
pub fn tenant_root_path() -> PartitionPath {
    TENANT_ROOT_SENTINEL
        .parse()
        .expect("<root> sentinel always parses")
}

#[cfg(test)]
mod tests {
    use super::*;

    fn scheme() -> PartitionScheme {
        PartitionScheme::parse("user={user}/year={year}/month={month}/topic={topic}").unwrap()
    }

    #[test]
    fn parses_template() {
        let s = scheme();
        assert_eq!(s.depth(), 4);
        assert_eq!(s.keys(), &["user", "year", "month", "topic"]);
    }

    #[test]
    fn rejects_bad_template() {
        assert!(matches!(
            PartitionScheme::parse(""),
            Err(InvalidScheme::Empty)
        ));
        assert!(matches!(
            PartitionScheme::parse("user"),
            Err(InvalidScheme::MissingEquals { .. })
        ));
        assert!(matches!(
            PartitionScheme::parse("user={uid}"),
            Err(InvalidScheme::PlaceholderMismatch { .. })
        ));
        assert!(matches!(
            PartitionScheme::parse("9bad={9bad}"),
            Err(InvalidScheme::BadKey { .. })
        ));
    }

    #[test]
    fn resolves_partitions() {
        let s = scheme();
        let p = Partitions::from([
            ("user", "alex"),
            ("year", "2026"),
            ("month", "05"),
            ("topic", "meetings"),
        ])
        .resolve(&s)
        .unwrap();
        assert_eq!(p.as_str(), "user=alex/year=2026/month=05/topic=meetings");
        assert_eq!(p.depth(), 4);
    }

    #[test]
    fn percent_encodes_slashes_in_values() {
        let s = PartitionScheme::parse("topic={topic}").unwrap();
        let p = Partitions::from([("topic", "a/b c")]).resolve(&s).unwrap();
        assert_eq!(p.as_str(), "topic=a%2Fb%20c");
    }

    #[test]
    fn resolve_rejects_missing_or_unknown() {
        let s = scheme();
        assert!(matches!(
            Partitions::from([("user", "alex")])
                .resolve(&s)
                .unwrap_err(),
            PartitionResolveError::MissingKey { .. }
        ));
        assert!(matches!(
            Partitions::from([
                ("user", "alex"),
                ("year", "2026"),
                ("month", "05"),
                ("topic", "meetings"),
                ("extra", "x"),
            ])
            .resolve(&s)
            .unwrap_err(),
            PartitionResolveError::UnknownKey { .. }
        ));
    }

    #[test]
    fn ancestors_iterates_deepest_first() {
        let p = PartitionPath::from_str("user=alex/year=2026/month=05/topic=meetings").unwrap();
        let got: Vec<String> = p.ancestors().map(|a| a.0).collect();
        assert_eq!(
            got,
            vec![
                "user=alex/year=2026/month=05".to_string(),
                "user=alex/year=2026".to_string(),
                "user=alex".to_string(),
            ]
        );
    }

    #[test]
    fn from_string_accepts_valid_path() {
        let p = PartitionPath::from_string("user=alex/year=2026").unwrap();
        assert_eq!(p.as_str(), "user=alex/year=2026");
    }

    #[test]
    fn from_string_rejects_empty() {
        assert!(PartitionPath::from_string("").is_err());
    }

    #[test]
    fn from_string_rejects_missing_eq() {
        assert!(PartitionPath::from_string("alex/2026").is_err());
    }

    #[test]
    fn from_string_rejects_empty_segment() {
        assert!(PartitionPath::from_string("user=alex//year=2026").is_err());
    }
}