use crate::{AetherError, Result, Slot, SlotKind};
use regex::Regex;
use std::sync::OnceLock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;
const SLOT_PATTERN: &str = r"\{\{AI:([a-zA-Z_][a-zA-Z0-9_]*)(?::([a-zA-Z]+))?\}\}";
static SLOT_REGEX: OnceLock<Regex> = OnceLock::new();
fn get_slot_regex() -> &'static Regex {
SLOT_REGEX.get_or_init(|| Regex::new(SLOT_PATTERN).expect("Invalid slot pattern regex"))
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Template {
pub content: String,
pub name: String,
pub slots: HashMap<String, Slot>,
pub metadata: TemplateMetadata,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct TemplateMetadata {
pub description: Option<String>,
pub language: Option<String>,
pub author: Option<String>,
pub version: Option<String>,
}
#[derive(Debug, Clone)]
pub struct SlotLocation {
pub name: String,
pub start: usize,
pub end: usize,
pub kind: Option<SlotKind>,
}
impl Template {
pub fn new(content: impl Into<String>) -> Self {
let content = content.into();
let slots = Self::parse_slots(&content);
Self {
content,
name: String::from("unnamed"),
slots,
metadata: TemplateMetadata::default(),
}
}
pub async fn from_file(path: impl AsRef<Path>) -> Result<Self> {
let path = path.as_ref();
let content = tokio::fs::read_to_string(path).await?;
let name = path
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or("unnamed")
.to_string();
Ok(Self {
name,
slots: Self::parse_slots(&content),
content,
metadata: TemplateMetadata::default(),
})
}
pub fn with_name(mut self, name: impl Into<String>) -> Self {
self.name = name.into();
self
}
pub fn with_metadata(mut self, metadata: TemplateMetadata) -> Self {
self.metadata = metadata;
self
}
pub fn with_slot(mut self, name: impl Into<String>, prompt: impl Into<String>) -> Self {
let name = name.into();
if let Some(slot) = self.slots.get_mut(&name) {
slot.prompt = prompt.into();
} else {
self.slots.insert(name.clone(), Slot::new(name, prompt));
}
self
}
pub fn configure_slot(mut self, slot: Slot) -> Self {
self.slots.insert(slot.name.clone(), slot);
self
}
fn parse_slots(content: &str) -> HashMap<String, Slot> {
let re = get_slot_regex();
let mut slots = HashMap::new();
for cap in re.captures_iter(content) {
let name = cap[1].to_string();
let kind = cap.get(2).map(|m| Self::parse_kind(m.as_str()));
let mut slot = Slot::new(&name, format!("Generate code for: {}", name));
if let Some(k) = kind {
slot = slot.with_kind(k);
}
slots.insert(name, slot);
}
slots
}
fn parse_kind(s: &str) -> SlotKind {
match s.to_lowercase().as_str() {
"raw" => SlotKind::Raw,
"function" | "fn" => SlotKind::Function,
"class" | "struct" => SlotKind::Class,
"html" => SlotKind::Html,
"css" => SlotKind::Css,
"js" | "javascript" => SlotKind::JavaScript,
"component" => SlotKind::Component,
other => SlotKind::Custom(other.to_string()),
}
}
fn find_locations(&self) -> Vec<SlotLocation> {
let re = get_slot_regex();
let mut locations = Vec::new();
for cap in re.captures_iter(&self.content) {
let full_match = cap.get(0).unwrap();
locations.push(SlotLocation {
name: cap[1].to_string(),
start: full_match.start(),
end: full_match.end(),
kind: cap.get(2).map(|m| Self::parse_kind(m.as_str())),
});
}
locations.sort_by(|a, b| b.start.cmp(&a.start));
locations
}
pub fn render(&self, injections: &HashMap<String, String>) -> Result<String> {
let mut result = self.content.clone();
let locations = self.find_locations();
for loc in locations {
let code = if let Some(code) = injections.get(&loc.name) {
code.clone()
} else if let Some(slot) = self.slots.get(&loc.name) {
if slot.required {
return Err(AetherError::SlotNotFound(loc.name));
}
slot.default.clone().unwrap_or_default()
} else {
return Err(AetherError::SlotNotFound(loc.name));
};
result.replace_range(loc.start..loc.end, &code);
}
Ok(result)
}
pub fn slot_names(&self) -> Vec<&str> {
self.slots.keys().map(|s| s.as_str()).collect()
}
pub fn validate(&self, injections: &HashMap<String, String>) -> Result<()> {
for (name, slot) in &self.slots {
if slot.required && !injections.contains_key(name) {
return Err(AetherError::SlotNotFound(name.clone()));
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_slots() {
let template = Template::new("Hello {{AI:greeting}} World {{AI:content:html}}");
assert_eq!(template.slots.len(), 2);
assert!(template.slots.contains_key("greeting"));
assert!(template.slots.contains_key("content"));
}
#[test]
fn test_render_template() {
let template = Template::new("<div>{{AI:content}}</div>");
let mut injections = HashMap::new();
injections.insert("content".to_string(), "<p>Hello</p>".to_string());
let result = template.render(&injections).unwrap();
assert_eq!(result, "<div><p>Hello</p></div>");
}
#[test]
fn test_slot_kind_parsing() {
let template = Template::new("{{AI:func:function}} {{AI:style:css}}");
assert_eq!(template.slots.get("func").unwrap().kind, SlotKind::Function);
assert_eq!(template.slots.get("style").unwrap().kind, SlotKind::Css);
}
}