use base64::engine::general_purpose;
use base64::Engine;
use rust_mcp_sdk::id_generator::IdGenerator;
use std::time::{SystemTime, UNIX_EPOCH};
pub struct TimeBase64Generator {
prefix: &'static str,
}
impl TimeBase64Generator {
pub fn new(prefix: Option<&'static str>) -> Self {
Self {
prefix: prefix.unwrap_or(""),
}
}
fn current_millis() -> u64 {
let now = SystemTime::now();
let duration = now
.duration_since(UNIX_EPOCH)
.expect("invalid system time!");
duration.as_millis() as u64
}
}
impl<T> IdGenerator<T> for TimeBase64Generator
where
T: From<String>,
{
fn generate(&self) -> T {
let timestamp = Self::current_millis();
let bytes = timestamp.to_le_bytes();
let encoded = general_purpose::URL_SAFE_NO_PAD.encode(bytes);
if self.prefix.is_empty() {
T::from(encoded)
} else {
T::from(format!("{}{}", self.prefix, encoded))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn generates_non_empty_id() {
let generator = TimeBase64Generator::new(None);
let id: String = generator.generate();
assert!(!id.is_empty(), "ID should not be empty");
}
#[test]
fn generates_id_with_prefix() {
let prefix = "ts_";
let generator = TimeBase64Generator::new(Some(prefix));
let id: String = generator.generate();
assert!(id.starts_with(prefix), "ID should start with prefix");
}
#[test]
fn ids_change_over_time() {
let generator = TimeBase64Generator::new(None);
let id1: String = generator.generate();
std::thread::sleep(std::time::Duration::from_millis(2));
let id2: String = generator.generate();
assert_ne!(id1, id2, "IDs generated at different times should differ");
}
#[test]
fn base64_decodes_to_timestamp() {
let generator = TimeBase64Generator::new(None);
let id: String = generator.generate();
let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD
.decode(&id)
.expect("Should decode base64");
let timestamp = u64::from_le_bytes(decoded.try_into().unwrap());
assert!(timestamp > 0, "Timestamp should be positive");
}
}