use super::types::{DryRunError, VariablePreview};
use crate::cook::execution::mapreduce::{MapPhase, ReducePhase, SetupPhase};
use serde_json::Value;
use std::collections::{HashMap, HashSet};
use tracing::debug;
pub struct VariableProcessor;
impl VariableProcessor {
pub fn new() -> Self {
Self
}
pub fn create_preview(
&self,
map_phase: &MapPhase,
work_items: &[Value],
setup_phase: Option<&SetupPhase>,
reduce_phase: Option<&ReducePhase>,
) -> Result<VariablePreview, DryRunError> {
debug!("Creating variable interpolation preview");
let setup_variables = self.extract_setup_variables(setup_phase);
let item_variables = self.extract_item_variables(work_items, 5)?; let reduce_variables = self.extract_reduce_variables(reduce_phase, work_items.len());
let undefined_references = self.find_undefined_references(
map_phase,
&setup_variables,
&item_variables,
&reduce_variables,
);
Ok(VariablePreview {
setup_variables,
item_variables,
reduce_variables,
undefined_references,
})
}
fn extract_setup_variables(&self, setup_phase: Option<&SetupPhase>) -> HashMap<String, String> {
let captured_outputs: HashMap<String, String> = setup_phase
.map(|setup| {
setup
.capture_outputs
.keys()
.map(|var_name| {
(
format!("setup.{}", var_name),
"<captured from setup command>".to_string(),
)
})
.collect()
})
.unwrap_or_default();
let default_vars = [
("setup.phase".to_string(), "setup".to_string()),
("setup.status".to_string(), "completed".to_string()),
];
default_vars.into_iter().chain(captured_outputs).collect()
}
fn extract_item_variables(
&self,
work_items: &[Value],
sample_size: usize,
) -> Result<Vec<HashMap<String, String>>, DryRunError> {
work_items
.iter()
.take(sample_size)
.enumerate()
.map(|(idx, item)| self.extract_single_item_variables(idx, item))
.collect()
}
fn extract_single_item_variables(
&self,
idx: usize,
item: &Value,
) -> Result<HashMap<String, String>, DryRunError> {
let index_var = ("item.index".to_string(), idx.to_string());
let item_vars: Vec<(String, String)> = match item {
Value::Object(map) => map
.iter()
.map(|(key, value)| (format!("item.{}", key), self.value_to_string(value)))
.collect(),
Value::String(s) => vec![
("item".to_string(), s.clone()),
("item.value".to_string(), s.clone()),
],
Value::Number(n) => vec![
("item".to_string(), n.to_string()),
("item.value".to_string(), n.to_string()),
],
_ => vec![("item".to_string(), serde_json::to_string(item)?)],
};
Ok(std::iter::once(index_var).chain(item_vars).collect())
}
fn extract_reduce_variables(
&self,
reduce_phase: Option<&ReducePhase>,
work_item_count: usize,
) -> HashMap<String, String> {
reduce_phase
.map(|_| {
[
(
"map.successful".to_string(),
"<count of successful items>".to_string(),
),
(
"map.failed".to_string(),
"<count of failed items>".to_string(),
),
("map.total".to_string(), work_item_count.to_string()),
(
"map.results".to_string(),
"<aggregated results>".to_string(),
),
("reduce.phase".to_string(), "reduce".to_string()),
]
.into_iter()
.collect()
})
.unwrap_or_default()
}
fn find_undefined_references(
&self,
map_phase: &MapPhase,
setup_variables: &HashMap<String, String>,
item_variables: &[HashMap<String, String>],
reduce_variables: &HashMap<String, String>,
) -> Vec<String> {
let shell_vars = [
"shell.output",
"shell.stdout",
"shell.stderr",
"shell.exit_code",
]
.into_iter()
.map(String::from);
let all_defined: HashSet<String> = setup_variables
.keys()
.cloned()
.chain(reduce_variables.keys().cloned())
.chain(
item_variables
.first()
.into_iter()
.flat_map(|item| item.keys().cloned()),
)
.chain(shell_vars)
.collect();
map_phase
.agent_template
.iter()
.flat_map(|cmd| self.extract_variable_references(cmd))
.filter(|var_ref| !all_defined.contains(var_ref) && !self.is_dynamic_variable(var_ref))
.collect::<HashSet<_>>()
.into_iter()
.collect()
}
fn extract_variable_references(
&self,
command: &crate::cook::workflow::WorkflowStep,
) -> Vec<String> {
let regex = regex::Regex::new(r"\$\{([^}]+)\}").expect("Invalid regex");
[command.claude.as_deref(), command.shell.as_deref()]
.into_iter()
.flatten()
.flat_map(|text| {
regex
.captures_iter(text)
.map(|cap| cap[1].to_string())
.collect::<Vec<_>>()
})
.collect()
}
fn is_dynamic_variable(&self, var_name: &str) -> bool {
let dynamic_prefixes = [
"env.", "runtime.", "timestamp.", "random.", ];
dynamic_prefixes
.iter()
.any(|prefix| var_name.starts_with(prefix))
}
fn value_to_string(&self, value: &Value) -> String {
match value {
Value::String(s) => s.clone(),
Value::Number(n) => n.to_string(),
Value::Bool(b) => b.to_string(),
Value::Null => "null".to_string(),
Value::Array(arr) => format!("[array with {} items]", arr.len()),
Value::Object(obj) => format!("[object with {} fields]", obj.len()),
}
}
}
impl Default for VariableProcessor {
fn default() -> Self {
Self::new()
}
}