mod common;

use std::collections::BTreeSet;

use common::*;
use pickaxe_rs::formula;
use pickaxe_rs::prelude::*;

/// Build and ground the running example from Modgil & Prakken 2014, Example 3.7.
///
/// K_n = {p}, K_p = {s, u, x}
/// R_s = {s1: p,q → r;  s2: v → ¬s}
/// R_d = {d1: p ⇒ q;  d2: s ⇒ t;  d3: t ⇒ ¬d1;
///        d4: u ⇒ v;  d5: v,x ⇒ ¬t;  d6: s ⇒ ¬p}
/// n(d_i) = d_i for all i.
fn build_example_3_7() -> GroundedSystem {
    let (system, _warnings) = AspicSystemBuilder::new()
        .with_classical_negation()
        .axiom(formula!(p))
        .ordinary_premise(formula!(s))
        .ordinary_premise(formula!(u))
        .ordinary_premise(formula!(x))
        .strict_rule("s1", vec![formula!(p), formula!(q)], formula!(r))
        .strict_rule("s2", vec![formula!(v)], formula!(!s))
        .named_defeasible_rule("d1", formula!(d1), vec![formula!(p)], formula!(q))
        .named_defeasible_rule("d2", formula!(d2), vec![formula!(s)], formula!(t))
        .named_defeasible_rule("d3", formula!(d3), vec![formula!(t)], formula!(!d1))
        .named_defeasible_rule("d4", formula!(d4), vec![formula!(u)], formula!(v))
        .named_defeasible_rule(
            "d5",
            formula!(d5),
            vec![formula!(v), formula!(x)],
            formula!(!t),
        )
        .named_defeasible_rule("d6", formula!(d6), vec![formula!(s)], formula!(!p))
        .build();
    NaiveGrounder.ground(&system).unwrap()
}

fn construct(system: &GroundedSystem) -> ArgumentStore {
    ForwardChainingConstructor.construct(system)
}

fn rule_ids(rules: &[DefeasibleRule]) -> BTreeSet<String> {
    rules.iter().map(|r| r.id.0.clone()).collect()
}

fn prem_formulas(prems: &[Premise]) -> BTreeSet<Formula> {
    prems.iter().map(|p| p.formula.clone()).collect()
}

// ---

#[test]
fn test_example_3_7_all_arguments_present() {
    let store = construct(&build_example_3_7());

    // 4 premise arguments
    assert!(
        find_premise_arg(&store, &formula!(p)).is_some(),
        "premise p"
    );
    assert!(
        find_premise_arg(&store, &formula!(s)).is_some(),
        "premise s"
    );
    assert!(
        find_premise_arg(&store, &formula!(u)).is_some(),
        "premise u"
    );
    assert!(
        find_premise_arg(&store, &formula!(x)).is_some(),
        "premise x"
    );

    // 8 inference arguments
    assert!(
        find_arg_by_conclusion_and_rule(&store, &formula!(q), "d1").is_some(),
        "q via d1"
    );
    assert!(
        find_arg_by_conclusion_and_rule(&store, &formula!(t), "d2").is_some(),
        "t via d2"
    );
    assert!(
        find_arg_by_conclusion_and_rule(&store, &formula!(!d1), "d3").is_some(),
        "¬d1 via d3"
    );
    assert!(
        find_arg_by_conclusion_and_rule(&store, &formula!(v), "d4").is_some(),
        "v via d4"
    );
    assert!(
        find_arg_by_conclusion_and_rule(&store, &formula!(!t), "d5").is_some(),
        "¬t via d5"
    );
    assert!(
        find_arg_by_conclusion_and_rule(&store, &formula!(!s), "s2").is_some(),
        "¬s via s2"
    );
    assert!(
        find_arg_by_conclusion_and_rule(&store, &formula!(r), "s1").is_some(),
        "r via s1"
    );
    assert!(
        find_arg_by_conclusion_and_rule(&store, &formula!(!p), "d6").is_some(),
        "¬p via d6"
    );

    assert_eq!(
        store.len(),
        12,
        "Example 3.7 should produce exactly 12 arguments"
    );
}

#[test]
fn test_empty_knowledge_base() {
    let system = GroundedSystem {
        knowledge_base: KnowledgeBase::default(),
        strict_rules: vec![],
        defeasible_rules: vec![],
        contrariness: ContrarinessFunction::new(),
    };
    let store = construct(&system);
    assert!(store.is_empty());
    assert_eq!(store.len(), 0);
}

#[test]
fn test_premise_only_no_rules() {
    let (sys, _) = AspicSystemBuilder::new()
        .axiom(formula!(a))
        .axiom(formula!(b))
        .axiom(formula!(c))
        .build();
    let grounded = NaiveGrounder.ground(&sys).unwrap();
    let store = construct(&grounded);

    assert_eq!(store.len(), 3);
    for (_, arg) in store.all_arguments() {
        assert!(
            matches!(arg, Argument::Premise { .. }),
            "all arguments should be Premise variants"
        );
    }
}

#[test]
fn test_zero_antecedent_rule() {
    let system = GroundedSystem {
        knowledge_base: KnowledgeBase::default(),
        strict_rules: vec![],
        defeasible_rules: vec![DefeasibleRule {
            id: RuleId("d1".into()),
            antecedents: vec![],
            consequent: formula!(q),
            name: None,
        }],
        contrariness: ContrarinessFunction::new(),
    };
    let store = construct(&system);

    assert_eq!(store.len(), 1);
    let (_, arg) = store.all_arguments().next().unwrap();
    assert!(matches!(arg, Argument::Inference { sub_arguments, .. } if sub_arguments.is_empty()));
    assert_eq!(*arg.conclusion(), formula!(q));
}

#[test]
fn test_no_duplicate_arguments() {
    let store = construct(&build_example_3_7());

    let inferences: Vec<_> = store
        .all_arguments()
        .filter_map(|(_, arg)| {
            if let Argument::Inference {
                rule,
                sub_arguments,
                ..
            } = arg
            {
                Some((rule.id().clone(), sub_arguments.clone()))
            } else {
                None
            }
        })
        .collect();

    for i in 0..inferences.len() {
        for j in (i + 1)..inferences.len() {
            assert!(
                inferences[i] != inferences[j],
                "duplicate inference: rule={}, subs={:?}",
                inferences[i].0,
                inferences[i].1
            );
        }
    }
}

#[test]
fn test_multiple_proof_paths_preserved() {
    let (sys, _) = AspicSystemBuilder::new()
        .axiom(formula!(a))
        .named_defeasible_rule("d1", formula!(d1), vec![formula!(a)], formula!(c))
        .named_defeasible_rule("d2", formula!(d2), vec![formula!(a)], formula!(c))
        .build();
    let grounded = NaiveGrounder.ground(&sys).unwrap();
    let store = construct(&grounded);

    let c_args = store.arguments_for(&formula!(c));
    assert_eq!(
        c_args.len(),
        2,
        "two distinct arguments for c via d1 and d2"
    );
}

#[test]
fn test_argument_ids_sequential() {
    let store = construct(&build_example_3_7());

    for (i, (id, _)) in store.all_arguments().enumerate() {
        assert_eq!(
            id,
            ArgumentId(i),
            "ID at position {i} should be ArgumentId({i})"
        );
    }
}

// ---

#[test]
fn test_accessor_a1_premise() {
    let store = construct(&build_example_3_7());
    let a1_id = find_premise_arg(&store, &formula!(p)).expect("A1");
    let a1 = store.get(a1_id).unwrap();

    assert_eq!(*a1.conclusion(), formula!(p));
    assert_eq!(
        prem_formulas(&a1.premises(&store)),
        BTreeSet::from([formula!(p)])
    );
    assert_eq!(a1.sub_argument_ids(&store), vec![a1_id]);
    assert!(a1.defeasible_rules(&store).is_empty());
    assert!(a1.top_rule().is_none());
    assert!(a1.is_strict(&store));
    assert!(!a1.is_defeasible(&store));
    assert!(a1.is_firm(&store));
    assert!(!a1.is_plausible(&store));
}

#[test]
fn test_accessor_a2_defeasible_inference() {
    let store = construct(&build_example_3_7());
    let a1_id = find_premise_arg(&store, &formula!(p)).expect("A1");
    let a2_id = find_arg_by_conclusion_and_rule(&store, &formula!(q), "d1").expect("A2");
    let a2 = store.get(a2_id).unwrap();

    assert_eq!(*a2.conclusion(), formula!(q));
    assert_eq!(
        prem_formulas(&a2.premises(&store)),
        BTreeSet::from([formula!(p)])
    );

    let sub_ids: BTreeSet<_> = a2.sub_argument_ids(&store).into_iter().collect();
    assert!(sub_ids.contains(&a1_id), "Sub(A2) should contain A1");
    assert!(sub_ids.contains(&a2_id), "Sub(A2) should contain A2");
    assert_eq!(sub_ids.len(), 2);

    assert_eq!(
        rule_ids(&a2.defeasible_rules(&store)),
        BTreeSet::from(["d1".into()])
    );

    let top = a2.top_rule().expect("A2 has a top rule");
    assert!(top.is_defeasible(), "A2 top rule d1 is defeasible");
    assert_eq!(top.id().0, "d1");

    assert!(!a2.is_strict(&store));
    assert!(a2.is_defeasible(&store));
    assert!(a2.is_firm(&store));
    assert!(!a2.is_plausible(&store));
}

#[test]
fn test_accessor_a3_strict_over_defeasible() {
    let store = construct(&build_example_3_7());
    let a1_id = find_premise_arg(&store, &formula!(p)).expect("A1");
    let a2_id = find_arg_by_conclusion_and_rule(&store, &formula!(q), "d1").expect("A2");
    let a3_id = find_arg_by_conclusion_and_rule(&store, &formula!(r), "s1").expect("A3");
    let a3 = store.get(a3_id).unwrap();

    assert_eq!(*a3.conclusion(), formula!(r));
    assert_eq!(
        prem_formulas(&a3.premises(&store)),
        BTreeSet::from([formula!(p)])
    );

    let sub_ids: BTreeSet<_> = a3.sub_argument_ids(&store).into_iter().collect();
    assert!(sub_ids.contains(&a1_id));
    assert!(sub_ids.contains(&a2_id));
    assert!(sub_ids.contains(&a3_id));
    assert_eq!(sub_ids.len(), 3);

    assert_eq!(
        rule_ids(&a3.defeasible_rules(&store)),
        BTreeSet::from(["d1".into()])
    );

    let top = a3.top_rule().expect("A3 has a top rule");
    assert!(top.is_strict(), "A3 top rule s1 is strict");
    assert_eq!(top.id().0, "s1");

    assert!(!a3.is_strict(&store), "A3 not strict (DefRules ≠ ∅)");
    assert!(a3.is_defeasible(&store));
    assert!(a3.is_firm(&store));
    assert!(!a3.is_plausible(&store));
}

#[test]
fn test_property_combinations() {
    let store = construct(&build_example_3_7());

    // (a) A1: strict + firm (axiom premise, no rules)
    let a1 = store
        .get(find_premise_arg(&store, &formula!(p)).unwrap())
        .unwrap();
    assert!(a1.is_strict(&store) && a1.is_firm(&store));
    assert!(!a1.is_defeasible(&store) && !a1.is_plausible(&store));

    // (b) B1: strict + plausible (ordinary premise, no rules)
    let b1 = store
        .get(find_premise_arg(&store, &formula!(s)).unwrap())
        .unwrap();
    assert!(b1.is_strict(&store) && b1.is_plausible(&store));
    assert!(!b1.is_defeasible(&store) && !b1.is_firm(&store));

    // (c) A2: defeasible + firm (axiom premise p, defeasible rule d1)
    let a2 = store
        .get(find_arg_by_conclusion_and_rule(&store, &formula!(q), "d1").unwrap())
        .unwrap();
    assert!(a2.is_defeasible(&store) && a2.is_firm(&store));
    assert!(!a2.is_strict(&store) && !a2.is_plausible(&store));

    // (d) B2: defeasible + plausible (ordinary premise s, defeasible rule d2)
    let b2 = store
        .get(find_arg_by_conclusion_and_rule(&store, &formula!(t), "d2").unwrap())
        .unwrap();
    assert!(b2.is_defeasible(&store) && b2.is_plausible(&store));
    assert!(!b2.is_strict(&store) && !b2.is_firm(&store));
}

#[test]
fn test_sub_argument_ids_includes_self() {
    let store = construct(&build_example_3_7());

    let cases = [
        (find_premise_arg(&store, &formula!(p)).unwrap(), "A1"),
        (
            find_arg_by_conclusion_and_rule(&store, &formula!(q), "d1").unwrap(),
            "A2",
        ),
        (
            find_arg_by_conclusion_and_rule(&store, &formula!(r), "s1").unwrap(),
            "A3",
        ),
        (
            find_arg_by_conclusion_and_rule(&store, &formula!(!d1), "d3").unwrap(),
            "B3",
        ),
    ];

    for (id, label) in &cases {
        let arg = store.get(*id).unwrap();
        assert!(
            arg.sub_argument_ids(&store).contains(id),
            "{label}: sub_argument_ids must include self"
        );
    }
}

#[test]
fn test_premises_multi_level() {
    let store = construct(&build_example_3_7());

    // B3 (¬d1 via d3) chains B1(s) → B2(t via d2) → B3(¬d1 via d3)
    let b3 = store
        .get(find_arg_by_conclusion_and_rule(&store, &formula!(!d1), "d3").unwrap())
        .unwrap();
    let b3_prems = prem_formulas(&b3.premises(&store));
    assert_eq!(
        b3_prems,
        BTreeSet::from([formula!(s)]),
        "B3 premises should be {{s}}"
    );

    // D4 (¬t via d5) uses C2(v via d4, prem u) and D3_prem(x)
    let d4 = store
        .get(find_arg_by_conclusion_and_rule(&store, &formula!(!t), "d5").unwrap())
        .unwrap();
    let d4_prems = prem_formulas(&d4.premises(&store));
    assert_eq!(
        d4_prems,
        BTreeSet::from([formula!(u), formula!(x)]),
        "D4 premises should be {{u, x}}"
    );
}

#[test]
fn test_direct_sub_ids() {
    let store = construct(&build_example_3_7());

    // A3 (r via s1) has direct sub-arguments A1 and A2
    let a1_id = find_premise_arg(&store, &formula!(p)).unwrap();
    let a2_id = find_arg_by_conclusion_and_rule(&store, &formula!(q), "d1").unwrap();
    let a3 = store
        .get(find_arg_by_conclusion_and_rule(&store, &formula!(r), "s1").unwrap())
        .unwrap();
    let direct: BTreeSet<_> = a3.direct_sub_ids().iter().copied().collect();
    assert!(direct.contains(&a1_id), "A3 direct subs should include A1");
    assert!(direct.contains(&a2_id), "A3 direct subs should include A2");
    assert_eq!(direct.len(), 2);

    // A1 (premise) has no direct sub-arguments
    let a1 = store.get(a1_id).unwrap();
    assert!(
        a1.direct_sub_ids().is_empty(),
        "premise has no direct sub-arguments"
    );
}

// ---

#[test]
fn test_last_def_rules_definition_3_20_example() {
    let (sys, _) = AspicSystemBuilder::new()
        .ordinary_premise(formula!(p))
        .ordinary_premise(formula!(q))
        .strict_rule("s1", vec![formula!(r), formula!(s)], formula!(t))
        .named_defeasible_rule("d1", formula!(d1), vec![formula!(p)], formula!(r))
        .named_defeasible_rule("d2", formula!(d2), vec![formula!(q)], formula!(s))
        .build();
    let grounded = NaiveGrounder.ground(&sys).unwrap();
    let store = construct(&grounded);

    let t_id =
        find_arg_by_conclusion_and_rule(&store, &formula!(t), "s1").expect("argument for t via s1");
    let t_arg = store.get(t_id).unwrap();
    let last_def = rule_ids(&t_arg.last_defeasible_rules(&store));

    assert_eq!(
        last_def,
        BTreeSet::from(["d1".into(), "d2".into()]),
        "LastDefRules of t should be {{d1, d2}}"
    );
}

#[test]
fn test_last_def_rules_premise_arg() {
    let store = construct(&build_example_3_7());
    let a1 = store
        .get(find_premise_arg(&store, &formula!(p)).unwrap())
        .unwrap();
    assert!(
        a1.last_defeasible_rules(&store).is_empty(),
        "premise: LastDefRules = ∅"
    );
}

#[test]
fn test_last_def_rules_defeasible_top() {
    let store = construct(&build_example_3_7());
    let a2 = store
        .get(find_arg_by_conclusion_and_rule(&store, &formula!(q), "d1").unwrap())
        .unwrap();
    let last_def = rule_ids(&a2.last_defeasible_rules(&store));
    assert_eq!(
        last_def,
        BTreeSet::from(["d1".into()]),
        "A2 LastDefRules = {{d1}}"
    );
}

#[test]
fn test_last_def_rules_strict_top_propagation() {
    let store = construct(&build_example_3_7());
    let a3 = store
        .get(find_arg_by_conclusion_and_rule(&store, &formula!(r), "s1").unwrap())
        .unwrap();
    let last_def = rule_ids(&a3.last_defeasible_rules(&store));
    assert_eq!(
        last_def,
        BTreeSet::from(["d1".into()]),
        "A3 LastDefRules = {{d1}} (propagated from A2)"
    );
}

// ---

#[test]
fn test_store_get_valid_and_invalid_ids() {
    let store = construct(&build_example_3_7());

    for i in 0..store.len() {
        let arg = store.get(ArgumentId(i));
        assert!(arg.is_some(), "get(ArgumentId({i})) should return Some");
        assert_eq!(arg.unwrap().id(), ArgumentId(i));
    }

    assert!(store.get(ArgumentId(store.len())).is_none());
    assert!(store.get(ArgumentId(usize::MAX)).is_none());
}

#[test]
fn test_store_arguments_for() {
    let store = construct(&build_example_3_7());

    let q_args = store.arguments_for(&formula!(q));
    assert!(!q_args.is_empty(), "should have argument(s) for q");
    for &id in q_args {
        assert_eq!(*store.get(id).unwrap().conclusion(), formula!(q));
    }

    let absent = store.arguments_for(&formula!(zzz_nonexistent));
    assert!(absent.is_empty(), "no arguments for absent formula");
}

#[test]
fn test_store_conclusion_index_consistency() {
    let store = construct(&build_example_3_7());

    for (id, arg) in store.all_arguments() {
        let indexed = store.arguments_for(arg.conclusion());
        assert!(
            indexed.contains(&id),
            "argument {:?} should appear in arguments_for({:?})",
            id,
            arg.conclusion()
        );
    }
}

#[test]
fn test_store_all_arguments_order() {
    let store = construct(&build_example_3_7());
    let mut count = 0;
    for (id, _) in store.all_arguments() {
        assert_eq!(id.0, count, "all_arguments should yield sequential IDs");
        count += 1;
    }
    assert_eq!(count, store.len());
}

#[test]
fn test_store_next_id() {
    let store = construct(&build_example_3_7());
    assert_eq!(store.next_id(), ArgumentId(store.len()));
}

// --- Non-circular argument construction (Li et al. 2017)

#[test]
fn test_two_rule_cycle_terminates() {
    // d1: a => b, d2: b => a, KB: {a}
    // Without the non-circular check this would loop forever.
    // Expected: premise a, inference b (via d1). The candidate a (via d2
    // applied to b) is rejected because a already appears in Sub.
    let (sys, _) = AspicSystemBuilder::new()
        .axiom(formula!(a))
        .named_defeasible_rule("d1", formula!(d1), vec![formula!(a)], formula!(b))
        .named_defeasible_rule("d2", formula!(d2), vec![formula!(b)], formula!(a))
        .build();
    let grounded = NaiveGrounder.ground(&sys).unwrap();
    let store = construct(&grounded);

    assert_eq!(store.len(), 2, "premise a + inference b only");
    assert!(find_premise_arg(&store, &formula!(a)).is_some());
    assert!(find_arg_by_conclusion_and_rule(&store, &formula!(b), "d1").is_some());
    // The circular candidate (a via d2) must not exist
    assert!(
        find_arg_by_conclusion_and_rule(&store, &formula!(a), "d2").is_none(),
        "circular argument a via d2 should be rejected"
    );
}

#[test]
fn test_three_rule_cycle_terminates() {
    // d1: a => b, d2: b => c, d3: c => a, KB: {a}
    // Expected: premise a, inference b, inference c.
    // The closing step c => a is rejected (a in Sub of the chain).
    let (sys, _) = AspicSystemBuilder::new()
        .axiom(formula!(a))
        .named_defeasible_rule("d1", formula!(d1), vec![formula!(a)], formula!(b))
        .named_defeasible_rule("d2", formula!(d2), vec![formula!(b)], formula!(c))
        .named_defeasible_rule("d3", formula!(d3), vec![formula!(c)], formula!(a))
        .build();
    let grounded = NaiveGrounder.ground(&sys).unwrap();
    let store = construct(&grounded);

    assert_eq!(store.len(), 3, "premise a + inference b + inference c");
    assert!(find_premise_arg(&store, &formula!(a)).is_some());
    assert!(find_arg_by_conclusion_and_rule(&store, &formula!(b), "d1").is_some());
    assert!(find_arg_by_conclusion_and_rule(&store, &formula!(c), "d2").is_some());
    assert!(
        find_arg_by_conclusion_and_rule(&store, &formula!(a), "d3").is_none(),
        "circular argument a via d3 should be rejected"
    );
}

#[test]
fn test_non_circular_paths_preserved() {
    // d1: a => c, d2: b => c, KB: {a, b}
    // Both d1(a) and d2(b) conclude c via independent sub-trees.
    // Neither is circular — both should be constructed.
    let (sys, _) = AspicSystemBuilder::new()
        .axiom(formula!(a))
        .axiom(formula!(b))
        .named_defeasible_rule("d1", formula!(d1), vec![formula!(a)], formula!(c))
        .named_defeasible_rule("d2", formula!(d2), vec![formula!(b)], formula!(c))
        .build();
    let grounded = NaiveGrounder.ground(&sys).unwrap();
    let store = construct(&grounded);

    assert_eq!(
        store.len(),
        4,
        "2 premises + 2 independent inferences for c"
    );
    assert!(find_premise_arg(&store, &formula!(a)).is_some());
    assert!(find_premise_arg(&store, &formula!(b)).is_some());
    assert!(find_arg_by_conclusion_and_rule(&store, &formula!(c), "d1").is_some());
    assert!(find_arg_by_conclusion_and_rule(&store, &formula!(c), "d2").is_some());
}
