aethellib 0.9.6

//! plan construction types for unified rule execution.
//!
//! this module provides typed rule keys, plan registration, and the validated
//! intermediate plan that can be compiled for repeated generation.

use std::collections::HashMap;
use std::sync::Arc;

use rand::Rng;

use crate::{
    corpus::{Corpus, PooledValue},
    engine::{AethelError, GenerationContext, validation::validate_pool_ref},
};

use super::{
    combinators::{RuleExpr, eval_expr},
    error::{PlanError, PlanErrorReport},
    validation,
};

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
/// validated identifier for one rule output.
pub struct RuleKey(Arc<str>);

impl RuleKey {
    /// creates a new rule key and validates its format.
    pub fn new(name: impl AsRef<str>) -> Result<Self, PlanError> {
        let name = name.as_ref().trim();

        if name.is_empty() {
            return Err(PlanError::InvalidIdentifier {
                kind: "rule key".to_string(),
                value: name.to_string(),
                reason: "must not be empty".to_string(),
            });
        }

        if !name
            .chars()
            .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '_' | '-' | '.'))
        {
            return Err(PlanError::InvalidIdentifier {
                kind: "rule key".to_string(),
                value: name.to_string(),
                reason: "allowed characters are ascii letters, digits, '_', '-', and '.'"
                    .to_string(),
            });
        }

        Ok(Self(Arc::from(name)))
    }

    /// returns this key as `&str`.
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl From<&RuleKey> for RuleKey {
    fn from(value: &RuleKey) -> Self {
        value.clone()
    }
}

impl AsRef<RuleKey> for RuleKey {
    fn as_ref(&self) -> &RuleKey {
        self
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
/// validated selector for one exact corpus value pool.
pub struct PoolRef {
    section: Arc<str>,
    field: Arc<str>,
}

impl PoolRef {
    /// creates a new pool selector after validating section and field names.
    pub fn new(section: impl AsRef<str>, field: impl AsRef<str>) -> Result<Self, PlanError> {
        let section = section.as_ref().trim();
        let field = field.as_ref().trim();

        validate_pool_ref(section, "section")?;
        validate_pool_ref(field, "field")?;

        Ok(Self {
            section: Arc::from(section),
            field: Arc::from(field),
        })
    }

    /// returns the section name.
    pub fn section(&self) -> &str {
        &self.section
    }

    /// returns the field name.
    pub fn field(&self) -> &str {
        &self.field
    }
}

#[derive(Clone)]
/// one named rule node in a plan.
pub struct PlanNode {
    pub key: RuleKey,
    pub expr: RuleExpr,
}

/// builder for typed generation plans.
pub struct PlanBuilder<'a> {
    pub(crate) corpus: &'a Corpus,
    pub(crate) nodes: Vec<PlanNode>,
}

impl<'a> PlanBuilder<'a> {
    /// creates a new plan builder for a corpus.
    pub fn new(corpus: &'a Corpus) -> Self {
        Self {
            corpus,
            nodes: Vec::new(),
        }
    }

    /// appends a named rule expression to the plan.
    pub fn rule(mut self, key: impl Into<RuleKey>, expr: RuleExpr) -> Self {
        self.nodes.push(PlanNode {
            key: key.into(),
            expr,
        });
        self
    }

    /// validates the plan and returns a validated intermediate representation.
    pub fn validate(self) -> Result<ValidatedPlan<'a>, PlanErrorReport> {
        validation::validate_plan(self)
    }
}

/// successfully validated plan ready for compilation.
pub struct ValidatedPlan<'a> {
    pub(crate) corpus: &'a Corpus,
    pub(crate) nodes: Vec<PlanNode>,
}

/// reusable compiled plan with indexed pools and deterministic execution order.
pub struct CompiledPlan<'a> {
    pub(crate) corpus: &'a Corpus,
    pub(crate) nodes: Vec<PlanNode>,
    pub(crate) order: Vec<usize>,
    pub(crate) pool_index: HashMap<PoolRef, Vec<PooledValue>>,
}

impl<'a> ValidatedPlan<'a> {
    /// compiles a validated plan into an immutable runtime representation.
    pub fn compile(self) -> Result<CompiledPlan<'a>, AethelError> {
        let key_to_index: HashMap<RuleKey, usize> = self
            .nodes
            .iter()
            .enumerate()
            .map(|(idx, node)| (node.key.clone(), idx))
            .collect();

        let mut adjacency: Vec<Vec<usize>> = vec![Vec::new(); self.nodes.len()];
        let mut indegree: Vec<usize> = vec![0; self.nodes.len()];

        for (idx, node) in self.nodes.iter().enumerate() {
            let mut deps = Vec::new();
            validation::collect_dependencies(&node.expr, &mut deps);
            for dep in deps {
                if let Some(dep_idx) = key_to_index.get(&dep) {
                    adjacency[*dep_idx].push(idx);
                    indegree[idx] += 1;
                }
            }
        }

        let mut queue: Vec<usize> = indegree
            .iter()
            .enumerate()
            .filter(|(_, degree)| **degree == 0)
            .map(|(idx, _)| idx)
            .collect();

        let mut order = Vec::with_capacity(self.nodes.len());
        while let Some(node) = queue.pop() {
            order.push(node);
            for next in &adjacency[node] {
                indegree[*next] = indegree[*next].saturating_sub(1);
                if indegree[*next] == 0 {
                    queue.push(*next);
                }
            }
        }

        if order.len() != self.nodes.len() {
            return Err(AethelError::Custom(
                "validation bug: cyclic plan reached compile".to_string(),
            ));
        }

        let mut used_pool_refs = Vec::new();
        for node in &self.nodes {
            validation::collect_pool_refs(&node.expr, &mut used_pool_refs);
        }

        let mut pool_index = HashMap::new();
        for pool_ref in used_pool_refs {
            if pool_index.contains_key(&pool_ref) {
                continue;
            }

            let values = self
                .corpus
                .pooled_values_for_field_section(pool_ref.field(), pool_ref.section())
                .ok_or_else(|| AethelError::PoolNotFound {
                    section: pool_ref.section().to_string(),
                    field: pool_ref.field().to_string(),
                })?;

            pool_index.insert(pool_ref, values.to_vec());
        }

        Ok(CompiledPlan {
            corpus: self.corpus,
            nodes: self.nodes,
            order,
            pool_index,
        })
    }
}

impl<'a> CompiledPlan<'a> {
    /// executes the compiled plan and returns the populated generation context.
    pub fn generate(&self, rng: &mut dyn Rng) -> Result<GenerationContext<'a>, AethelError> {
        let mut ctx = GenerationContext::new(self.corpus);

        for idx in &self.order {
            let node = &self.nodes[*idx];
            let result = eval_expr(&node.expr, &ctx, &self.pool_index, rng)?;
            ctx.insert_typed(node.key.clone(), result);
        }

        Ok(ctx)
    }
}

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

    use crate::{
        corpus::Corpus,
        engine::{
            ComposedValue,
            combinators::{custom, join, lit, recall, when},
        },
    };

    use super::{PlanBuilder, PoolRef, RuleKey};

    #[test]
    fn rule_key_validation_rejects_invalid_names() {
        assert!(RuleKey::new("").is_err());
        assert!(RuleKey::new("has space").is_err());
        assert!(RuleKey::new("weapon-name_1").is_ok());
    }

    #[test]
    fn validate_reports_missing_dependency() {
        let raw = r#"
[header]
title = "plan-test"
target = "weapon"

[name]
prefix = ["ash"]
"#;

        let corpus = Corpus::builder("weapon")
            .add_str("plan-test", raw)
            .build()
            .expect("corpus should build");

        let out = RuleKey::new("out").expect("rule key should be valid");
        let missing = RuleKey::new("missing").expect("rule key should be valid");

        let result = PlanBuilder::new(&corpus)
            .rule(&out, recall(&missing))
            .rule(
                RuleKey::new("literal").expect("rule key should be valid"),
                lit("x"),
            )
            .validate();

        let err = match result {
            Ok(_) => panic!("validation should fail"),
            Err(err) => err,
        };

        assert!(err.to_string().contains("missing dependency"));
    }

    #[test]
    fn pool_ref_new_rejects_invalid_names() {
        let err = PoolRef::new("", "prefix").expect_err("empty section should fail");
        assert!(err.to_string().contains("section"));

        let err = PoolRef::new("name", "prefix value").expect_err("space should fail");
        assert!(err.to_string().contains("allowed characters"));
    }

    #[test]
    fn pool_ref_new_accepts_valid_names() {
        let pool = PoolRef::new("name", "prefix").expect("valid pool ref should build");
        assert_eq!(pool.section(), "name");
        assert_eq!(pool.field(), "prefix");
    }

    #[test]
    fn custom_expression_supports_user_defined_combinator_logic() {
        let raw = r#"
[header]
title = "plan-test"
target = "weapon"

[name]
prefix = ["ash"]
"#;

        let corpus = Corpus::builder("weapon")
            .add_str("plan-test", raw)
            .build()
            .expect("corpus should build");

        let base = RuleKey::new("base").expect("rule key should be valid");
        let out = RuleKey::new("out").expect("rule key should be valid");
        let base_for_closure = base;

        let compiled = PlanBuilder::new(&corpus)
            .rule(&base_for_closure, lit("ash"))
            .rule(
                &out,
                custom([base_for_closure.clone()], move |ctx, _rng| {
                    let prior = ctx.require(&base_for_closure)?;
                    Ok(ComposedValue {
                        value: format!("{} spear", prior.value),
                        provenance: prior.provenance.clone(),
                    })
                }),
            )
            .validate()
            .expect("plan should validate")
            .compile()
            .expect("plan should compile");

        let mut rng = StdRng::seed_from_u64(13);
        let ctx = compiled
            .generate(&mut rng)
            .expect("generation should succeed");
        let generated = ctx.require(&out).expect("out should exist");

        assert_eq!(generated.value, "ash spear");
    }

    #[test]
    fn when_combinator_includes_inner_when_condition_is_non_empty() {
        let raw = r#"
[header]
title = "plan-test"
target = "weapon"

[name]
prefix = ["ash"]
"#;

        let corpus = Corpus::builder("weapon")
            .add_str("plan-test", raw)
            .build()
            .expect("corpus should build");

        let condition = RuleKey::new("condition").expect("rule key should be valid");
        let out = RuleKey::new("out").expect("rule key should be valid");

        let compiled = PlanBuilder::new(&corpus)
            .rule(&condition, lit("enabled"))
            .rule(
                &out,
                join([
                    lit("start"),
                    when(recall(&condition), lit("-extra")),
                    lit("-end"),
                ]),
            )
            .validate()
            .expect("plan should validate")
            .compile()
            .expect("plan should compile");

        let mut rng = StdRng::seed_from_u64(9);
        let ctx = compiled
            .generate(&mut rng)
            .expect("generation should succeed");

        assert_eq!(
            ctx.require(&out).expect("out should exist").value,
            "start-extra-end"
        );
    }

    #[test]
    fn when_combinator_skips_inner_when_condition_is_empty() {
        let raw = r#"
[header]
title = "plan-test"
target = "weapon"

[name]
prefix = ["ash"]
"#;

        let corpus = Corpus::builder("weapon")
            .add_str("plan-test", raw)
            .build()
            .expect("corpus should build");

        let condition = RuleKey::new("condition").expect("rule key should be valid");
        let out = RuleKey::new("out").expect("rule key should be valid");

        let compiled = PlanBuilder::new(&corpus)
            .rule(&condition, lit(""))
            .rule(
                &out,
                join([
                    lit("start"),
                    when(recall(&condition), lit("-extra")),
                    lit("-end"),
                ]),
            )
            .validate()
            .expect("plan should validate")
            .compile()
            .expect("plan should compile");

        let mut rng = StdRng::seed_from_u64(10);
        let ctx = compiled
            .generate(&mut rng)
            .expect("generation should succeed");

        assert_eq!(
            ctx.require(&out).expect("out should exist").value,
            "start-end"
        );
    }
}