1use crate::{AetherError, Result, Slot, SlotKind};
7use regex::Regex;
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10use std::path::Path;
11
12const SLOT_PATTERN: &str = r"\{\{AI:([a-zA-Z_][a-zA-Z0-9_]*)(?::([a-zA-Z]+))?\}\}";
15
16#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct Template {
19 pub content: String,
21
22 pub name: String,
24
25 pub slots: HashMap<String, Slot>,
27
28 pub metadata: TemplateMetadata,
30}
31
32#[derive(Debug, Clone, Default, Serialize, Deserialize)]
34pub struct TemplateMetadata {
35 pub description: Option<String>,
37
38 pub language: Option<String>,
40
41 pub author: Option<String>,
43
44 pub version: Option<String>,
46}
47
48#[derive(Debug, Clone)]
50pub struct SlotLocation {
51 pub name: String,
53
54 pub start: usize,
56
57 pub end: usize,
59
60 pub kind: Option<SlotKind>,
62}
63
64impl Template {
65 pub fn new(content: impl Into<String>) -> Self {
80 let content = content.into();
81 let slots = Self::parse_slots(&content);
82
83 Self {
84 content,
85 name: String::from("unnamed"),
86 slots,
87 metadata: TemplateMetadata::default(),
88 }
89 }
90
91 pub async fn from_file(path: impl AsRef<Path>) -> Result<Self> {
93 let path = path.as_ref();
94 let content = tokio::fs::read_to_string(path).await?;
95 let name = path
96 .file_stem()
97 .and_then(|s| s.to_str())
98 .unwrap_or("unnamed")
99 .to_string();
100
101 Ok(Self {
102 name,
103 slots: Self::parse_slots(&content),
104 content,
105 metadata: TemplateMetadata::default(),
106 })
107 }
108
109 pub fn with_name(mut self, name: impl Into<String>) -> Self {
111 self.name = name.into();
112 self
113 }
114
115 pub fn with_metadata(mut self, metadata: TemplateMetadata) -> Self {
117 self.metadata = metadata;
118 self
119 }
120
121 pub fn with_slot(mut self, name: impl Into<String>, prompt: impl Into<String>) -> Self {
128 let name = name.into();
129 if let Some(slot) = self.slots.get_mut(&name) {
130 slot.prompt = prompt.into();
131 } else {
132 self.slots.insert(name.clone(), Slot::new(name, prompt));
133 }
134 self
135 }
136
137 pub fn configure_slot(mut self, slot: Slot) -> Self {
139 self.slots.insert(slot.name.clone(), slot);
140 self
141 }
142
143 fn parse_slots(content: &str) -> HashMap<String, Slot> {
145 let re = Regex::new(SLOT_PATTERN).expect("Invalid slot pattern regex");
146 let mut slots = HashMap::new();
147
148 for cap in re.captures_iter(content) {
149 let name = cap[1].to_string();
150 let kind = cap.get(2).map(|m| Self::parse_kind(m.as_str()));
151
152 let mut slot = Slot::new(&name, format!("Generate code for: {}", name));
153 if let Some(k) = kind {
154 slot = slot.with_kind(k);
155 }
156 slots.insert(name, slot);
157 }
158
159 slots
160 }
161
162 fn parse_kind(s: &str) -> SlotKind {
164 match s.to_lowercase().as_str() {
165 "raw" => SlotKind::Raw,
166 "function" | "fn" => SlotKind::Function,
167 "class" | "struct" => SlotKind::Class,
168 "html" => SlotKind::Html,
169 "css" => SlotKind::Css,
170 "js" | "javascript" => SlotKind::JavaScript,
171 "component" => SlotKind::Component,
172 other => SlotKind::Custom(other.to_string()),
173 }
174 }
175
176 pub fn find_locations(&self) -> Vec<SlotLocation> {
178 let re = Regex::new(SLOT_PATTERN).expect("Invalid slot pattern regex");
179 let mut locations = Vec::new();
180
181 for cap in re.captures_iter(&self.content) {
182 let full_match = cap.get(0).unwrap();
183 locations.push(SlotLocation {
184 name: cap[1].to_string(),
185 start: full_match.start(),
186 end: full_match.end(),
187 kind: cap.get(2).map(|m| Self::parse_kind(m.as_str())),
188 });
189 }
190
191 locations.sort_by(|a, b| b.start.cmp(&a.start));
193 locations
194 }
195
196 pub fn render(&self, injections: &HashMap<String, String>) -> Result<String> {
202 let mut result = self.content.clone();
203 let locations = self.find_locations();
204
205 for loc in locations {
206 let code = if let Some(code) = injections.get(&loc.name) {
207 code.clone()
208 } else if let Some(slot) = self.slots.get(&loc.name) {
209 if slot.required {
210 return Err(AetherError::SlotNotFound(loc.name));
211 }
212 slot.default.clone().unwrap_or_default()
213 } else {
214 return Err(AetherError::SlotNotFound(loc.name));
215 };
216
217 result.replace_range(loc.start..loc.end, &code);
218 }
219
220 Ok(result)
221 }
222
223 pub fn slot_names(&self) -> Vec<&str> {
225 self.slots.keys().map(|s| s.as_str()).collect()
226 }
227
228 pub fn validate(&self, injections: &HashMap<String, String>) -> Result<()> {
230 for (name, slot) in &self.slots {
231 if slot.required && !injections.contains_key(name) {
232 return Err(AetherError::SlotNotFound(name.clone()));
233 }
234 }
235 Ok(())
236 }
237}
238
239#[cfg(test)]
240mod tests {
241 use super::*;
242
243 #[test]
244 fn test_parse_slots() {
245 let template = Template::new("Hello {{AI:greeting}} World {{AI:content:html}}");
246 assert_eq!(template.slots.len(), 2);
247 assert!(template.slots.contains_key("greeting"));
248 assert!(template.slots.contains_key("content"));
249 }
250
251 #[test]
252 fn test_render_template() {
253 let template = Template::new("<div>{{AI:content}}</div>");
254 let mut injections = HashMap::new();
255 injections.insert("content".to_string(), "<p>Hello</p>".to_string());
256
257 let result = template.render(&injections).unwrap();
258 assert_eq!(result, "<div><p>Hello</p></div>");
259 }
260
261 #[test]
262 fn test_slot_kind_parsing() {
263 let template = Template::new("{{AI:func:function}} {{AI:style:css}}");
264 assert_eq!(template.slots.get("func").unwrap().kind, SlotKind::Function);
265 assert_eq!(template.slots.get("style").unwrap().kind, SlotKind::Css);
266 }
267}