aethellib 0.10.1

//! standard combinator functions that each return executable rules.

use std::{ops::RangeInclusive, sync::Arc};

use rand::{Rng, RngExt};

use super::{AethelError, ComposedValue, GenerationContext, PlanError, PoolRef, Rule, RuleKey};

pub(crate) type RuleBox = Box<dyn Rule>;
type RuleFn = dyn for<'a> Fn(&GenerationContext<'a>, &mut dyn Rng) -> Result<ComposedValue, AethelError>
    + Send
    + Sync;
type ValidateFn = dyn Fn(&RuleKey) -> Vec<PlanError> + Send + Sync;

/// closure-backed [`Rule`] implementation for concise built-in and custom rules.
struct InlineRule {
    dependencies: Vec<RuleKey>,
    pool_refs: Vec<PoolRef>,
    logic: Arc<RuleFn>,
    validate: Arc<ValidateFn>,
}

impl InlineRule {
    /// creates a rule from executable logic and no declared validation metadata.
    fn new<F>(logic: F) -> Self
    where
        F: for<'a> Fn(&GenerationContext<'a>, &mut dyn Rng) -> Result<ComposedValue, AethelError>
            + Send
            + Sync
            + 'static,
    {
        Self {
            dependencies: Vec::new(),
            pool_refs: Vec::new(),
            logic: Arc::new(logic),
            validate: Arc::new(|_| Vec::new()),
        }
    }

    /// declares rule dependencies used by plan validation and ordering.
    fn with_dependencies(mut self, dependencies: impl IntoIterator<Item = RuleKey>) -> Self {
        self.dependencies = dependencies.into_iter().collect();
        self
    }

    /// declares corpus pool references used by plan validation.
    fn with_pool_refs(mut self, pool_refs: impl IntoIterator<Item = PoolRef>) -> Self {
        self.pool_refs = pool_refs.into_iter().collect();
        self
    }

    /// declares additional validation logic for this rule.
    fn with_validator<F>(mut self, validate: F) -> Self
    where
        F: Fn(&RuleKey) -> Vec<PlanError> + Send + Sync + 'static,
    {
        self.validate = Arc::new(validate);
        self
    }
}

impl Rule for InlineRule {
    fn dependencies(&self) -> Vec<RuleKey> {
        self.dependencies.clone()
    }

    fn pool_refs(&self) -> Vec<PoolRef> {
        self.pool_refs.clone()
    }

    fn validate(&self, rule_key: &RuleKey) -> Vec<PlanError> {
        (self.validate)(rule_key)
    }

    fn execute<'a>(
        &self,
        ctx: &GenerationContext<'a>,
        rng: &mut dyn Rng,
    ) -> Result<ComposedValue, AethelError> {
        (self.logic)(ctx, rng)
    }
}

impl Rule for Box<dyn Rule> {
    fn dependencies(&self) -> Vec<RuleKey> {
        self.as_ref().dependencies()
    }

    fn pool_refs(&self) -> Vec<PoolRef> {
        self.as_ref().pool_refs()
    }

    fn validate(&self, rule_key: &RuleKey) -> Vec<super::PlanError> {
        self.as_ref().validate(rule_key)
    }

    fn execute<'a>(
        &self,
        ctx: &GenerationContext<'a>,
        rng: &mut dyn Rng,
    ) -> Result<ComposedValue, AethelError> {
        self.as_ref().execute(ctx, rng)
    }
}

/// creates a boxed rule.
pub(crate) fn boxed(rule: impl Rule + 'static) -> RuleBox {
    Box::new(rule)
}

/// creates a typed pool pick rule that selects `amount` values with replacement.
pub fn pick_many(pool: PoolRef, amount: usize, delimiter: impl Into<String>) -> Box<dyn Rule> {
    pick_range(pool, amount..=amount, delimiter)
}

/// creates a typed pool pick rule that selects one value.
pub fn pick(pool: PoolRef) -> Box<dyn Rule> {
    pick_many(pool, 1, "")
}

/// creates a typed pool pick rule that selects a random amount of values with replacement.
pub fn pick_range(
    pool: PoolRef,
    amount: RangeInclusive<usize>,
    delimiter: impl Into<String>,
) -> Box<dyn Rule> {
    let delimiter = delimiter.into();
    boxed(
        InlineRule::new({
            let pool = pool.clone();
            move |ctx: &GenerationContext<'_>, rng: &mut dyn Rng| {
                let values = values_for_pool(ctx, &pool)?;
                let amount = rng.random_range(amount.clone());
                let mut selected_values = Vec::new();

                for _ in 0..amount {
                    let idx = rng.random_range(0..values.len());
                    selected_values.push(values[idx].clone());
                }

                Ok(compose_values(selected_values, &delimiter))
            }
        })
        .with_pool_refs([pool]),
    )
}

/// creates a typed pool pick rule that selects `amount` unique values without replacement.
pub fn pick_unique(pool: PoolRef, amount: usize, delimiter: impl Into<String>) -> Box<dyn Rule> {
    pick_unique_range(pool, amount..=amount, delimiter)
}

/// creates a typed pool pick rule that selects a random amount of unique values without replacement.
pub fn pick_unique_range(
    pool: PoolRef,
    amount: RangeInclusive<usize>,
    delimiter: impl Into<String>,
) -> Box<dyn Rule> {
    let delimiter = delimiter.into();
    boxed(
        InlineRule::new({
            let pool = pool.clone();
            move |ctx: &GenerationContext<'_>, rng: &mut dyn Rng| {
                let values = values_for_pool(ctx, &pool)?;
                let amount = rng.random_range(amount.clone()).min(values.len());
                let mut available = (0..values.len()).collect::<Vec<_>>();
                let mut selected_values = Vec::new();

                for _ in 0..amount {
                    let available_idx = rng.random_range(0..available.len());
                    let value_idx = available.swap_remove(available_idx);
                    selected_values.push(values[value_idx].clone());
                }

                Ok(compose_values(selected_values, &delimiter))
            }
        })
        .with_pool_refs([pool]),
    )
}

/// creates a dependency recall rule.
pub fn recall(key: impl Into<RuleKey>) -> Box<dyn Rule> {
    let key = key.into();
    boxed(
        InlineRule::new({
            let key = key.clone();
            move |ctx: &GenerationContext<'_>, _rng: &mut dyn Rng| ctx.require(&key).cloned()
        })
        .with_dependencies([key]),
    )
}

/// creates a literal rule.
pub fn lit(text: impl Into<String>) -> Box<dyn Rule> {
    let text = text.into();
    boxed(InlineRule::new(
        move |_ctx: &GenerationContext<'_>, _rng: &mut dyn Rng| {
            Ok(ComposedValue {
                value: text.clone(),
                provenance: Vec::new(),
            })
        },
    ))
}

/// creates a rule that concatenates child rule outputs in order.
pub fn join(parts: impl IntoIterator<Item = impl Rule + 'static>) -> Box<dyn Rule> {
    let parts: Vec<RuleBox> = parts.into_iter().map(boxed).collect();
    let dependencies = collect_dependencies(&parts);
    let pool_refs = collect_pool_refs(&parts);
    let validation_errors = collect_validation_errors(&parts);

    boxed(
        InlineRule::new(move |ctx: &GenerationContext<'_>, rng: &mut dyn Rng| {
            let mut result = ComposedValue {
                value: String::new(),
                provenance: Vec::new(),
            };

            for part in &parts {
                result = result.merge(part.execute(ctx, rng)?);
            }

            Ok(result)
        })
        .with_dependencies(dependencies)
        .with_pool_refs(pool_refs)
        .with_validator(move |rule_key| retarget_validation_errors(&validation_errors, rule_key)),
    )
}

/// creates a rule that evaluates `inner` with probability `probability`.
pub fn chance(probability: f64, inner: impl Rule + 'static) -> Box<dyn Rule> {
    let inner = boxed(inner);
    let dependencies = inner.dependencies();
    let pool_refs = inner.pool_refs();
    let validation_errors = inner.validate(rule_key_placeholder());

    boxed(
        InlineRule::new(move |ctx: &GenerationContext<'_>, rng: &mut dyn Rng| {
            if rng.random::<f64>() < probability {
                inner.execute(ctx, rng)
            } else {
                Ok(ComposedValue {
                    value: String::new(),
                    provenance: Vec::new(),
                })
            }
        })
        .with_dependencies(dependencies)
        .with_pool_refs(pool_refs)
        .with_validator(move |rule_key| {
            let mut errors = retarget_validation_errors(&validation_errors, rule_key);
            if !(0.0..=1.0).contains(&probability) {
                errors.push(PlanError::InvalidChanceProbability {
                    rule: rule_key.as_str().to_string(),
                    value: probability,
                });
            }
            errors
        }),
    )
}

/// creates a weighted choice rule that executes exactly one child.
pub fn weighted(choices: impl IntoIterator<Item = (u32, impl Rule + 'static)>) -> Box<dyn Rule> {
    let choices: Vec<(u32, RuleBox)> = choices
        .into_iter()
        .map(|(weight, rule)| (weight, boxed(rule)))
        .collect();
    let dependencies = choices
        .iter()
        .flat_map(|(_, rule)| rule.dependencies())
        .collect::<Vec<_>>();
    let pool_refs = choices
        .iter()
        .flat_map(|(_, rule)| rule.pool_refs())
        .collect::<Vec<_>>();
    let validation_errors = choices
        .iter()
        .flat_map(|(_, rule)| rule.validate(rule_key_placeholder()))
        .collect::<Vec<_>>();
    let total_weight: u32 = choices.iter().map(|(weight, _)| *weight).sum();

    boxed(
        InlineRule::new(move |ctx: &GenerationContext<'_>, rng: &mut dyn Rng| {
            if total_weight == 0 {
                return Err(AethelError::Custom(
                    "weighted choice has a total weight of 0".to_string(),
                ));
            }

            let mut roll = rng.random_range(0..total_weight);
            for (weight, rule) in &choices {
                if roll < *weight {
                    return rule.execute(ctx, rng);
                }
                roll -= weight;
            }

            Err(AethelError::Custom(
                "mathematical error in weighted expression".to_string(),
            ))
        })
        .with_dependencies(dependencies)
        .with_pool_refs(pool_refs)
        .with_validator(move |rule_key| {
            let mut errors = retarget_validation_errors(&validation_errors, rule_key);
            if total_weight == 0 {
                errors.push(PlanError::WeightedTotalZero {
                    rule: rule_key.as_str().to_string(),
                });
            }
            errors
        }),
    )
}

/// creates an equal-weight choice rule that executes exactly one child.
pub fn choice(choices: impl IntoIterator<Item = impl Rule + 'static>) -> Box<dyn Rule> {
    weighted(choices.into_iter().map(|rule| (1, rule)))
}

/// creates a rule that executes `secondary` when `primary` returns an error.
pub fn fallback(primary: impl Rule + 'static, secondary: impl Rule + 'static) -> Box<dyn Rule> {
    let primary = boxed(primary);
    let secondary = boxed(secondary);
    let dependencies = primary
        .dependencies()
        .into_iter()
        .chain(secondary.dependencies())
        .collect::<Vec<_>>();
    let pool_refs = primary
        .pool_refs()
        .into_iter()
        .chain(secondary.pool_refs())
        .collect::<Vec<_>>();
    let validation_errors = primary
        .validate(rule_key_placeholder())
        .into_iter()
        .chain(secondary.validate(rule_key_placeholder()))
        .collect::<Vec<_>>();

    boxed(
        InlineRule::new(move |ctx: &GenerationContext<'_>, rng: &mut dyn Rng| {
            primary
                .execute(ctx, rng)
                .or_else(|_| secondary.execute(ctx, rng))
        })
        .with_dependencies(dependencies)
        .with_pool_refs(pool_refs)
        .with_validator(move |rule_key| retarget_validation_errors(&validation_errors, rule_key)),
    )
}

/// creates a rule that transforms a child output string while preserving provenance.
pub fn map<F>(inner: impl Rule + 'static, transform: F) -> Box<dyn Rule>
where
    F: Fn(String) -> String + Send + Sync + 'static,
{
    let inner = boxed(inner);
    let dependencies = inner.dependencies();
    let pool_refs = inner.pool_refs();
    let validation_errors = inner.validate(rule_key_placeholder());

    boxed(
        InlineRule::new(move |ctx: &GenerationContext<'_>, rng: &mut dyn Rng| {
            let mut composed = inner.execute(ctx, rng)?;
            composed.value = transform(composed.value);
            Ok(composed)
        })
        .with_dependencies(dependencies)
        .with_pool_refs(pool_refs)
        .with_validator(move |rule_key| retarget_validation_errors(&validation_errors, rule_key)),
    )
}

/// creates a custom rule with declared dependencies passed into the closure.
pub fn custom<K, F>(dependencies: impl IntoIterator<Item = K>, logic: F) -> Box<dyn Rule>
where
    K: Into<RuleKey>,
    F: for<'a> Fn(
            &GenerationContext<'a>,
            &mut dyn Rng,
            &[RuleKey],
        ) -> Result<ComposedValue, AethelError>
        + Send
        + Sync
        + 'static,
{
    let dependencies = dependencies.into_iter().map(Into::into).collect::<Vec<_>>();
    boxed(
        InlineRule::new({
            let dependencies = dependencies.clone();
            move |ctx, rng| logic(ctx, rng, &dependencies)
        })
        .with_dependencies(dependencies),
    )
}

/// creates a rule that evaluates `inner` only when `condition` returns a non-empty value.
pub fn when(condition: impl Rule + 'static, inner: impl Rule + 'static) -> Box<dyn Rule> {
    let condition = boxed(condition);
    let inner = boxed(inner);
    let dependencies = condition
        .dependencies()
        .into_iter()
        .chain(inner.dependencies())
        .collect::<Vec<_>>();
    let pool_refs = condition
        .pool_refs()
        .into_iter()
        .chain(inner.pool_refs())
        .collect::<Vec<_>>();
    let validation_errors = condition
        .validate(rule_key_placeholder())
        .into_iter()
        .chain(inner.validate(rule_key_placeholder()))
        .collect::<Vec<_>>();

    boxed(
        InlineRule::new(move |ctx: &GenerationContext<'_>, rng: &mut dyn Rng| {
            let condition_value = condition.execute(ctx, rng)?;
            if condition_value.value.is_empty() {
                Ok(ComposedValue {
                    value: String::new(),
                    provenance: Vec::new(),
                })
            } else {
                inner.execute(ctx, rng)
            }
        })
        .with_dependencies(dependencies)
        .with_pool_refs(pool_refs)
        .with_validator(move |rule_key| retarget_validation_errors(&validation_errors, rule_key)),
    )
}

fn collect_dependencies(rules: &[RuleBox]) -> Vec<RuleKey> {
    rules
        .iter()
        .flat_map(|rule| rule.dependencies())
        .collect::<Vec<_>>()
}

fn collect_pool_refs(rules: &[RuleBox]) -> Vec<PoolRef> {
    rules
        .iter()
        .flat_map(|rule| rule.pool_refs())
        .collect::<Vec<_>>()
}

fn values_for_pool<'a>(
    ctx: &'a GenerationContext<'_>,
    pool: &PoolRef,
) -> Result<&'a [crate::corpus::PooledValue], AethelError> {
    let values = ctx
        .corpus
        .pooled_values_for_field_section(pool.field(), pool.section())
        .ok_or_else(|| AethelError::PoolNotFound {
            section: pool.section().to_string(),
            field: pool.field().to_string(),
        })?;

    if values.is_empty() {
        return Err(AethelError::Custom("pool is empty".to_string()));
    }

    Ok(values)
}

fn compose_values(
    selected_values: impl IntoIterator<Item = crate::corpus::PooledValue>,
    delimiter: &str,
) -> ComposedValue {
    let selected_values = selected_values.into_iter().collect::<Vec<_>>();

    ComposedValue {
        value: selected_values
            .iter()
            .map(|value| value.value.clone())
            .collect::<Vec<_>>()
            .join(delimiter),
        provenance: selected_values
            .iter()
            .flat_map(|value| value.provenance.clone())
            .collect(),
    }
}

fn collect_validation_errors(rules: &[RuleBox]) -> Vec<PlanError> {
    rules
        .iter()
        .flat_map(|rule| rule.validate(rule_key_placeholder()))
        .collect::<Vec<_>>()
}

fn retarget_validation_errors(errors: &[PlanError], rule_key: &RuleKey) -> Vec<PlanError> {
    errors
        .iter()
        .cloned()
        .map(|error| match error {
            PlanError::InvalidChanceProbability { value, .. } => {
                PlanError::InvalidChanceProbability {
                    rule: rule_key.as_str().to_string(),
                    value,
                }
            }
            PlanError::WeightedTotalZero { .. } => PlanError::WeightedTotalZero {
                rule: rule_key.as_str().to_string(),
            },
            other => other,
        })
        .collect()
}

fn rule_key_placeholder() -> &'static RuleKey {
    static KEY: std::sync::OnceLock<RuleKey> = std::sync::OnceLock::new();
    KEY.get_or_init(|| RuleKey::new("__nested").expect("placeholder rule key should be valid"))
}