//! Forward-chaining argument construction via fixpoint saturation.
//!
//! Seeds the argument store with premise arguments, then repeatedly fires
//! every rule against the current store until no new arguments can be
//! derived. A non-circularity guard (Li et al. 2017 / Grooters & Prakken)
//! prevents arguments whose conclusion already appears in their own
//! transitive sub-argument tree.

use std::collections::HashSet;

use crate::framework::argument::{Argument, ArgumentId};
use crate::framework::argument_store::ArgumentStore;
use crate::language::formula::Formula;
use crate::language::grounding::GroundedSystem;
use crate::pipeline::traits::ArgumentConstructor;
use crate::system::rules::Rule;

/// Forward-chaining saturation argument constructor.
///
/// Delegates to `saturate` for the actual fixpoint computation.
/// See that function's documentation for algorithmic details.
pub struct ForwardChainingConstructor;

impl ArgumentConstructor for ForwardChainingConstructor {
    fn construct(&self, system: &GroundedSystem) -> ArgumentStore {
        saturate(system)
    }
}

/// Build every non-circular argument derivable from `system`.
///
/// The algorithm proceeds in two phases:
///
/// 1. **Seed** — create one [`Argument::Premise`] for each entry in the
///    knowledge base. Each premise argument records the set of conclusions
///    reachable through its (trivial) sub-argument tree, which is just the
///    premise formula itself.
///
/// 2. **Saturate** — repeatedly scan every rule and try to fire it against
///    the current argument store. For each rule whose antecedents can be
///    matched by existing argument conclusions, every combination of
///    supporting sub-arguments is considered. A candidate inference is
///    accepted only when:
///    - Its conclusion does **not** already appear among the transitive
///      sub-conclusions of the chosen sub-arguments (the *non-circularity*
///      guard from Li et al. 2017 / Grooters & Prakken).
///    - No structurally identical argument (same rule + same sub-argument
///      list) has been created before.
///
///    The loop terminates when a full pass over every rule produces no new
///    arguments -- i.e. a fixpoint has been reached. Convergence is
///    guaranteed: the store grows monotonically and the set of possible
///    arguments is finite (bounded by rules × sub-argument combinations),
///    so the fixpoint is reached in finitely many passes.
///
/// `sub_conclusions` is a parallel `Vec` indexed by `ArgumentId`: entry *i*
/// holds the set of all conclusion formulas reachable through the transitive
/// sub-arguments of argument *i* (including its own conclusion). This makes
/// the circularity check an O(1) `HashSet::contains` rather than a recursive
/// tree walk.
fn saturate(system: &GroundedSystem) -> ArgumentStore {
    let mut store = ArgumentStore::new();
    let mut sub_conclusions: Vec<HashSet<Formula>> = Vec::new();

    // Phase 1: seed premise arguments.
    for premise in &system.knowledge_base.premises {
        let id = store.next_id();
        let formula = premise.formula.clone();
        store.add(Argument::Premise {
            id,
            premise: premise.clone(),
        });
        let mut set = HashSet::new();
        set.insert(formula);
        sub_conclusions.push(set);
    }

    let all_rules: Vec<Rule> = system
        .strict_rules
        .iter()
        .map(|r| Rule::Strict(r.clone()))
        .chain(
            system
                .defeasible_rules
                .iter()
                .map(|r| Rule::Defeasible(r.clone())),
        )
        .collect();

    // Phase 2: fixpoint saturation.
    let mut changed = true;
    while changed {
        changed = false;
        for rule in &all_rules {
            let antecedents = rule.antecedents();
            let consequent = rule.consequent();

            if antecedents.is_empty() {
                if !has_argument_for(&store, consequent, rule) {
                    let id = store.next_id();
                    store.add(Argument::Inference {
                        id,
                        sub_arguments: vec![],
                        rule: rule.clone(),
                        conclusion: consequent.clone(),
                    });
                    let mut set = HashSet::new();
                    set.insert(consequent.clone());
                    sub_conclusions.push(set);
                    changed = true;
                }
                continue;
            }

            let combinations = find_antecedent_combinations(&store, antecedents);

            for combo in combinations {
                let mut merged: HashSet<Formula> = HashSet::new();
                for &sub_id in &combo {
                    if let Some(set) = sub_conclusions.get(sub_id.0) {
                        merged.extend(set.iter().cloned());
                    }
                }
                if merged.contains(consequent) {
                    continue;
                }

                if !has_argument_with_subs(&store, consequent, rule, &combo) {
                    let id = store.next_id();
                    store.add(Argument::Inference {
                        id,
                        sub_arguments: combo,
                        rule: rule.clone(),
                        conclusion: consequent.clone(),
                    });
                    merged.insert(consequent.clone());
                    sub_conclusions.push(merged);
                    changed = true;
                }
            }
        }
    }

    store
}

fn has_argument_for(store: &ArgumentStore, conclusion: &Formula, rule: &Rule) -> bool {
    for &arg_id in store.arguments_for(conclusion) {
        if let Some(Argument::Inference {
            rule: r,
            sub_arguments,
            ..
        }) = store.get(arg_id)
        {
            if sub_arguments.is_empty() && r.id() == rule.id() {
                return true;
            }
        }
    }
    false
}

fn has_argument_with_subs(
    store: &ArgumentStore,
    conclusion: &Formula,
    rule: &Rule,
    subs: &[ArgumentId],
) -> bool {
    for &arg_id in store.arguments_for(conclusion) {
        if let Some(Argument::Inference {
            rule: r,
            sub_arguments,
            ..
        }) = store.get(arg_id)
        {
            if r.id() == rule.id() && sub_arguments == subs {
                return true;
            }
        }
    }
    false
}

fn find_antecedent_combinations(
    store: &ArgumentStore,
    antecedents: &[Formula],
) -> Vec<Vec<ArgumentId>> {
    if antecedents.is_empty() {
        return vec![vec![]];
    }

    let first = &antecedents[0];
    let rest = &antecedents[1..];

    let first_args: Vec<ArgumentId> = store.arguments_for(first).to_vec();
    if first_args.is_empty() {
        return vec![];
    }

    let rest_combos = find_antecedent_combinations(store, rest);
    if rest_combos.is_empty() {
        return vec![];
    }

    let mut result = Vec::new();
    for &arg_id in &first_args {
        for rest_combo in &rest_combos {
            let mut combo = vec![arg_id];
            combo.extend_from_slice(rest_combo);
            result.push(combo);
        }
    }
    result
}
