aethellib 0.8.2

//! standard combinator functions that each return a [`Rule`].

use std::ops::Range;

use rand::{Rng, RngExt};

use super::{AethelError, ComposedValue, GenerationContext, InlineRule, Rule};

/// selects a random value from the pool identified by `section` and `field`.
pub fn pick(
    name: impl Into<String>,
    section: impl Into<String>,
    field: impl Into<String>,
) -> impl Rule + 'static {
    let section = section.into();
    let field = field.into();
    InlineRule::new(
        name,
        move |ctx: &GenerationContext<'_>, rng: &mut dyn Rng| {
            let values = ctx
                .corpus
                .pooled_values_for_field_section(&field, &section)
                .ok_or_else(|| AethelError::PoolNotFound {
                    section: section.clone(),
                    field: field.clone(),
                })?;

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

            let index = rng.random_range(0..values.len());
            let selected = &values[index];

            Ok(ComposedValue {
                value: selected.value.clone(),
                provenance: selected.provenance.clone(),
            })
        },
    )
}

/// executes `rule_a` and `rule_b` sequentially and merges their results into one [`ComposedValue`].
pub fn concat(
    name: impl Into<String>,
    rule_a: impl Rule + 'static,
    rule_b: impl Rule + 'static,
) -> impl Rule + 'static {
    InlineRule::new(
        name,
        move |ctx: &GenerationContext<'_>, rng: &mut dyn Rng| {
            let val_a = rule_a.execute(ctx, rng)?;
            let val_b = rule_b.execute(ctx, rng)?;
            Ok(val_a.merge(val_b))
        },
    )
}

/// tries `primary`; on any error executes `secondary` instead.
pub fn fallback(
    name: impl Into<String>,
    primary: impl Rule + 'static,
    secondary: impl Rule + 'static,
) -> impl Rule + 'static {
    InlineRule::new(
        name,
        move |ctx: &GenerationContext<'_>, rng: &mut dyn Rng| match primary.execute(ctx, rng) {
            Ok(val) => Ok(val),
            Err(_) => secondary.execute(ctx, rng),
        },
    )
}

/// performs a single weighted RNG roll to select and execute one rule from `choices`.
///
/// returns [`AethelError::Custom`] if the total weight of all choices is zero.
pub fn weighted_choice(
    name: impl Into<String>,
    choices: Vec<(u32, Box<dyn Rule>)>,
) -> impl Rule + 'static {
    InlineRule::new(
        name,
        move |ctx: &GenerationContext<'_>, rng: &mut dyn Rng| {
            let total_weight: u32 = choices.iter().map(|(w, _)| w).sum();

            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 choice".to_string(),
            ))
        },
    )
}

/// evaluates `probability` (0.0–1.0) against a single RNG roll.
/// on success executes `rule`; on failure returns an empty [`ComposedValue`].
pub fn chance(
    name: impl Into<String>,
    probability: f64,
    rule: impl Rule + 'static,
) -> impl Rule + 'static {
    InlineRule::new(
        name,
        move |ctx: &GenerationContext<'_>, rng: &mut dyn Rng| {
            let probability = probability.clamp(0.0, 1.0);
            let roll = rng.random::<f64>();

            if roll < probability {
                rule.execute(ctx, rng)
            } else {
                Ok(ComposedValue {
                    value: String::new(),
                    provenance: Vec::new(),
                })
            }
        },
    )
}
/// returns a rule that always produces a fixed string with no provenance.
pub fn lit(text: &'static str) -> impl Rule + 'static {
    InlineRule::new(
        format!("LITERAL({})", text),
        move |_ctx: &GenerationContext<'_>, _rng: &mut dyn Rng| {
            Ok(ComposedValue {
                value: text.to_string(),
                provenance: vec![],
            })
        },
    )
}

/// fetches a previously generated value from the context history.
/// returns an error if the key does not exist in the history.
pub fn recall(name: impl Into<String>, target_key: impl Into<String>) -> impl Rule + 'static {
    let target_key = target_key.into();
    InlineRule::new(
        name,
        move |ctx: &GenerationContext<'_>, _rng: &mut dyn Rng| {
            ctx.get_previous(&target_key)
                .cloned()
                .ok_or_else(|| AethelError::MissingDependency(target_key.clone()))
        },
    )
}

/// executes a rule and applies a transformation function to the resulting string.
/// provenance metadata is preserved exactly as it was.
pub fn map<F>(
    name: impl Into<String>,
    rule: impl Rule + 'static,
    transform: F,
) -> impl Rule + 'static
where
    F: Fn(String) -> String + 'static,
{
    InlineRule::new(
        name,
        move |ctx: &GenerationContext<'_>, rng: &mut dyn Rng| {
            let mut result = rule.execute(ctx, rng)?;
            result.value = transform(result.value);
            Ok(result)
        },
    )
}

/// selects a variable number of unique values from a pool without replacement.
/// merges the results using `separator`.
pub fn pick_multiple(
    name: impl Into<String>,
    section: impl Into<String>,
    field: impl Into<String>,
    count_range: Range<usize>,
    separator: Option<&'static str>,
) -> impl Rule + 'static {
    let section = section.into();
    let field = field.into();
    InlineRule::new(
        name,
        move |ctx: &GenerationContext<'_>, rng: &mut dyn Rng| {
            let values = ctx
                .corpus
                .pooled_values_for_field_section(&field, &section)
                .ok_or_else(|| AethelError::PoolNotFound {
                    section: section.clone(),
                    field: field.clone(),
                })?;

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

            let count = rng.random_range(count_range.clone());
            let mut selected_indices = Vec::new();
            let mut final_result = ComposedValue {
                value: String::new(),
                provenance: Vec::new(),
            };

            for _ in 0..count.min(values.len()) {
                let mut index;
                loop {
                    index = rng.random_range(0..values.len());
                    if !selected_indices.contains(&index) {
                        selected_indices.push(index);
                        break;
                    }
                }

                let selected = &values[index];
                if !final_result.value.is_empty()
                    && let Some(sep) = separator
                {
                    final_result.value.push_str(sep);
                }
                final_result.value.push_str(&selected.value);
                final_result.provenance.extend(selected.provenance.clone());
            }

            Ok(final_result)
        },
    )
}

/// evaluates `condition` and executes either `if_true` or `if_false` accordingly.
pub fn conditional(
    name: impl Into<String>,
    condition: bool,
    if_true: impl Rule + 'static,
    if_false: impl Rule + 'static,
) -> impl Rule + 'static {
    InlineRule::new(
        name,
        move |ctx: &GenerationContext<'_>, rng: &mut dyn Rng| {
            if condition {
                if_true.execute(ctx, rng)
            } else {
                if_false.execute(ctx, rng)
            }
        },
    )
}

/// executes a sequence of rules and joins them together with an optional separator.
pub fn sequence(
    name: impl Into<String>,
    rules: Vec<Box<dyn Rule>>,
    separator: Option<&'static str>,
) -> impl Rule + 'static {
    InlineRule::new(
        name,
        move |ctx: &GenerationContext<'_>, rng: &mut dyn Rng| {
            let mut iter = rules.iter();

            // grab the first rule to initialize our base ComposedValue
            let first_rule = match iter.next() {
                Some(r) => r,
                None => {
                    return Ok(ComposedValue {
                        value: String::new(),
                        provenance: Vec::new(),
                    });
                }
            };

            let mut final_result = first_rule.execute(ctx, rng)?;

            // iterate through the rest, appending the separator and the next rule's result
            for rule in iter {
                if let Some(sep) = separator {
                    final_result.value.push_str(sep);
                }

                let next_result = rule.execute(ctx, rng)?;
                final_result = final_result.merge(next_result);
            }

            Ok(final_result)
        },
    )
}

#[cfg(test)]
mod tests {
    use rand::{SeedableRng, rngs::StdRng};

    use crate::{
        corpus::Corpus,
        engine::{AethelError, GenerationContext, Rule},
    };

    use super::{chance, fallback, lit, pick, recall, sequence, weighted_choice};

    fn test_corpus() -> Corpus {
        let raw = r#"
[header]
title = "combinator test"
target = "weapon"

[name]
first = ["ash", "birch"]
"#;

        Corpus::builder("weapon")
            .add_str("combinator-test", raw)
            .build()
            .expect("corpus should build")
    }

    #[test]
    fn pick_returns_a_value_and_provenance_from_pool() {
        let corpus = test_corpus();
        let ctx = GenerationContext::new(&corpus);
        let mut rng = StdRng::seed_from_u64(1);

        let rule = pick("name_pick", "name", "first");
        let result = rule.execute(&ctx, &mut rng).expect("pick should succeed");

        assert!(result.value == "ash" || result.value == "birch");
        assert_eq!(result.provenance.len(), 1);
        assert_eq!(result.provenance[0].section, "name");
        assert_eq!(result.provenance[0].field, "first");
    }

    #[test]
    fn pick_returns_pool_not_found_for_missing_pool() {
        let corpus = test_corpus();
        let ctx = GenerationContext::new(&corpus);
        let mut rng = StdRng::seed_from_u64(2);

        let rule = pick("missing", "name", "last");
        let err = rule.execute(&ctx, &mut rng).expect_err("pick should fail");

        match err {
            AethelError::PoolNotFound { section, field } => {
                assert_eq!(section, "name");
                assert_eq!(field, "last");
            }
            _ => panic!("unexpected error variant"),
        }
    }

    #[test]
    fn fallback_executes_secondary_when_primary_fails() {
        let corpus = test_corpus();
        let ctx = GenerationContext::new(&corpus);
        let mut rng = StdRng::seed_from_u64(3);

        let primary = pick("primary", "name", "does_not_exist");
        let secondary = lit("backup");

        let rule = fallback("with_fallback", primary, secondary);
        let result = rule
            .execute(&ctx, &mut rng)
            .expect("fallback should succeed");

        assert_eq!(result.value, "backup");
    }

    #[test]
    fn weighted_choice_errors_on_zero_total_weight() {
        let corpus = test_corpus();
        let ctx = GenerationContext::new(&corpus);
        let mut rng = StdRng::seed_from_u64(4);

        let rule = weighted_choice("weighted", vec![(0, Box::new(lit("never")))]);
        let err = rule
            .execute(&ctx, &mut rng)
            .expect_err("weighted choice should fail");

        match err {
            AethelError::Custom(msg) => {
                assert!(msg.contains("total weight of 0"));
            }
            _ => panic!("unexpected error variant"),
        }
    }

    #[test]
    fn chance_clamps_probability_bounds() {
        let corpus = test_corpus();
        let ctx = GenerationContext::new(&corpus);

        let mut rng_never = StdRng::seed_from_u64(5);
        let never = chance("never", -1.0, lit("x"))
            .execute(&ctx, &mut rng_never)
            .expect("chance should execute");
        assert_eq!(never.value, "");
        assert!(never.provenance.is_empty());

        let mut rng_always = StdRng::seed_from_u64(6);
        let always = chance("always", 2.0, lit("x"))
            .execute(&ctx, &mut rng_always)
            .expect("chance should execute");
        assert_eq!(always.value, "x");
    }

    #[test]
    fn recall_and_sequence_handle_missing_and_empty_cases() {
        let corpus = test_corpus();
        let ctx = GenerationContext::new(&corpus);
        let mut rng = StdRng::seed_from_u64(7);

        let recall_missing = recall("recall", "unknown")
            .execute(&ctx, &mut rng)
            .expect_err("recall should fail for missing key");
        match recall_missing {
            AethelError::MissingDependency(key) => assert_eq!(key, "unknown"),
            _ => panic!("unexpected error variant"),
        }

        let empty = sequence("empty", vec![], Some(" "))
            .execute(&ctx, &mut rng)
            .expect("empty sequence should succeed");
        assert_eq!(empty.value, "");
        assert!(empty.provenance.is_empty());
    }

    #[test]
    fn sequence_joins_values_with_separator_in_order() {
        let corpus = test_corpus();
        let ctx = GenerationContext::new(&corpus);
        let mut rng = StdRng::seed_from_u64(8);

        let rules: Vec<Box<dyn Rule>> = vec![
            Box::new(lit("one")),
            Box::new(lit("two")),
            Box::new(lit("three")),
        ];

        let result = sequence("seq", rules, Some("-"))
            .execute(&ctx, &mut rng)
            .expect("sequence should execute");

        assert_eq!(result.value, "one-two-three");
    }
}