//! Core argument representation for ASPIC+.
//!
//! An ASPIC+ argument is a tree-structured proof built from premises and
//! inference rules. This module defines the two data types at the heart of
//! every pipeline stage:
//!
//! - [`ArgumentId`] — a lightweight, copy-cheap handle used to refer to
//!   arguments throughout the framework without carrying the full structure.
//! - [`Argument`] — the argument itself, either a bare [`Premise`] leaf or
//!   an [`Inference`](Argument::Inference) step that applies a rule to
//!   sub-arguments.
//!
//! Most tree-walking operations (e.g. collecting all premises or defeasible
//! rules) take an [`ArgumentStore`](super::argument_store::ArgumentStore)
//! reference so they can resolve sub-argument IDs without each node owning
//! its children.

use std::collections::HashSet;

use crate::language::formula::Formula;
use crate::system::knowledge::{Premise, PremiseKind};
use crate::system::rules::{DefeasibleRule, Rule, RuleId};

/// Lightweight handle into an [`ArgumentStore`](super::argument_store::ArgumentStore).
///
/// Wraps a dense `usize` index, so copies, hashes, and comparisons are all
/// O(1). The inner value corresponds to the insertion order in the store.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, Ord, PartialOrd)]
pub struct ArgumentId(pub usize);

/// A single ASPIC+ argument.
///
/// In ASPIC+ an argument is either:
///
/// - **Premise** — a leaf node whose conclusion is taken directly from the
///   knowledge base (axiom or ordinary premise).
/// - **Inference** — a node that applies one rule to zero or more
///   sub-arguments, deriving a new conclusion.
///
/// Most queries (`premises`, `defeasible_rules`, etc.) walk the sub-argument
/// tree via an [`ArgumentStore`](super::argument_store::ArgumentStore)
/// reference, keeping `Argument` itself small and cheaply clonable.
#[derive(Clone, Debug)]
pub enum Argument {
    /// A leaf argument drawn directly from the knowledge base.
    Premise { id: ArgumentId, premise: Premise },
    /// An inference step applying `rule` to `sub_arguments`, yielding `conclusion`.
    Inference {
        id: ArgumentId,
        sub_arguments: Vec<ArgumentId>,
        rule: Rule,
        conclusion: Formula,
    },
}

impl Argument {
    /// Returns this argument's unique identifier.
    pub fn id(&self) -> ArgumentId {
        match self {
            Argument::Premise { id, .. } => *id,
            Argument::Inference { id, .. } => *id,
        }
    }

    /// The formula that this argument derives (Conc).
    pub fn conclusion(&self) -> &Formula {
        match self {
            Argument::Premise { premise, .. } => &premise.formula,
            Argument::Inference { conclusion, .. } => conclusion,
        }
    }

    /// The top-level rule of the argument (`TopRule`), or `None` for premise
    /// arguments.
    pub fn top_rule(&self) -> Option<&Rule> {
        match self {
            Argument::Premise { .. } => None,
            Argument::Inference { rule, .. } => Some(rule),
        }
    }

    /// All sub-argument IDs (Sub), including this argument itself.
    ///
    /// Walks the tree recursively through `store`, collecting every
    /// transitive sub-argument exactly once.
    pub fn sub_argument_ids(
        &self,
        store: &super::argument_store::ArgumentStore,
    ) -> Vec<ArgumentId> {
        let mut seen = HashSet::new();
        let mut result = Vec::new();
        seen.insert(self.id());
        result.push(self.id());
        if let Argument::Inference { sub_arguments, .. } = self {
            for &sub_id in sub_arguments {
                if let Some(sub_arg) = store.get(sub_id) {
                    for id in sub_arg.sub_argument_ids(store) {
                        if seen.insert(id) {
                            result.push(id);
                        }
                    }
                }
            }
        }
        result
    }

    /// All premises used in this argument (Prem).
    ///
    /// Returns every leaf premise reachable through the sub-argument tree,
    /// deduplicated by formula.
    pub fn premises(&self, store: &super::argument_store::ArgumentStore) -> Vec<Premise> {
        match self {
            Argument::Premise { premise, .. } => vec![premise.clone()],
            Argument::Inference { sub_arguments, .. } => {
                let mut seen = HashSet::new();
                let mut prems = Vec::new();
                for &sub_id in sub_arguments {
                    if let Some(sub_arg) = store.get(sub_id) {
                        for p in sub_arg.premises(store) {
                            if seen.insert(p.formula.clone()) {
                                prems.push(p);
                            }
                        }
                    }
                }
                prems
            }
        }
    }

    /// All defeasible rules used anywhere in this argument (`DefRules`).
    pub fn defeasible_rules(
        &self,
        store: &super::argument_store::ArgumentStore,
    ) -> Vec<DefeasibleRule> {
        let mut seen = HashSet::new();
        let mut rules = Vec::new();
        self.collect_defeasible_rules(store, &mut seen, &mut rules);
        rules
    }

    fn collect_defeasible_rules(
        &self,
        store: &super::argument_store::ArgumentStore,
        seen: &mut HashSet<RuleId>,
        rules: &mut Vec<DefeasibleRule>,
    ) {
        if let Argument::Inference {
            rule,
            sub_arguments,
            ..
        } = self
        {
            if let Rule::Defeasible(dr) = rule {
                if seen.insert(dr.id.clone()) {
                    rules.push(dr.clone());
                }
            }
            for &sub_id in sub_arguments {
                if let Some(sub_arg) = store.get(sub_id) {
                    sub_arg.collect_defeasible_rules(store, seen, rules);
                }
            }
        }
    }

    /// The last defeasible rules — those applied at the topmost inference
    /// steps before reaching a strict rule or a premise.
    ///
    /// For a defeasible top rule this is just that rule; for a strict top
    /// rule this recurses into the immediate sub-arguments.
    pub fn last_defeasible_rules(
        &self,
        store: &super::argument_store::ArgumentStore,
    ) -> Vec<DefeasibleRule> {
        match self {
            Argument::Premise { .. } => vec![],
            Argument::Inference {
                rule,
                sub_arguments,
                ..
            } => {
                if let Rule::Defeasible(dr) = rule {
                    vec![dr.clone()]
                } else {
                    let mut seen = HashSet::new();
                    let mut result = Vec::new();
                    for &sub_id in sub_arguments {
                        if let Some(sub_arg) = store.get(sub_id) {
                            for dr in sub_arg.last_defeasible_rules(store) {
                                if seen.insert(dr.id.clone()) {
                                    result.push(dr);
                                }
                            }
                        }
                    }
                    result
                }
            }
        }
    }

    /// Whether the argument is strict (all inference rules are strict;
    /// equivalently, no defeasible rules are used anywhere).
    pub fn is_strict(&self, store: &super::argument_store::ArgumentStore) -> bool {
        self.defeasible_rules(store).is_empty()
    }

    /// Whether the argument is firm (all premises are axioms).
    pub fn is_firm(&self, store: &super::argument_store::ArgumentStore) -> bool {
        self.premises(store)
            .iter()
            .all(|p| p.kind == PremiseKind::Axiom)
    }

    /// Whether the argument is defeasible (not strict; uses at least
    /// one defeasible rule).
    pub fn is_defeasible(&self, store: &super::argument_store::ArgumentStore) -> bool {
        !self.defeasible_rules(store).is_empty()
    }

    /// Whether the argument is plausible (not firm; uses at least one
    /// non-axiom premise).
    pub fn is_plausible(&self, store: &super::argument_store::ArgumentStore) -> bool {
        !self.is_firm(store)
    }

    /// Direct sub-argument IDs (the immediate children, not recursive).
    pub fn direct_sub_ids(&self) -> &[ArgumentId] {
        match self {
            Argument::Premise { .. } => &[],
            Argument::Inference { sub_arguments, .. } => sub_arguments,
        }
    }
}
