//! SAT-based semantics solver using propositional encodings.
//!
//! Encodes Dung-style argumentation semantics as SAT problems (following
//! Besnard & Doutre 2004; Cerutti et al. 2014) and uses incremental SAT
//! solving for model enumeration. Generic over the SAT backend via
//! rustsat's `Solve + SolveIncremental` traits.
//!
//! # Encodings
//!
//! Each argument `a` maps to a propositional variable `x_a`. The encodings
//! layer constraints:
//!
//! - **Conflict-free**: for each defeat (a, b), `¬x_a ∨ ¬x_b`.
//! - **Defence (admissibility)**: for each attacker `a` of `b`, if `b` is
//!   accepted then some counter-attacker of `a` must be accepted.
//! - **Stable range**: each argument is either accepted or attacked by an
//!   accepted argument.
//! - **Complete closure**: auxiliary variables encode the characteristic
//!   function fixpoint condition.
//!
//! Preferred and semi-stable extensions use iterative maximisation on top
//! of admissibility/completeness encodings.

use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::marker::PhantomData;

use rustsat::solvers::{Solve, SolveIncremental, SolverResult};
use rustsat::types::{Lit, TernaryVal, Var};

use crate::framework::abstract_af::AbstractFramework;
use crate::framework::argument::ArgumentId;
use crate::framework::semantics::{Extension, Semantics};
use crate::pipeline::traits::SemanticsSolver;

/// SAT-based semantics solver, generic over the SAT backend.
///
/// Uses propositional encodings of argumentation semantics (Besnard & Doutre,
/// Cerutti et al.) with incremental solving for model enumeration.
///
/// The default backend is `BatSat` (pure Rust, no FFI). Users can swap in any
/// solver implementing rustsat's `Solve + SolveIncremental` traits, e.g.:
/// ```ignore
/// // For production performance (requires rustsat-cadical dependency):
/// SatSolver::<rustsat_cadical::CaDiCaL>::default()
/// ```
pub struct SatSolver<S = rustsat_batsat::BasicSolver> {
    _marker: PhantomData<fn() -> S>,
}

impl<S> Default for SatSolver<S> {
    fn default() -> Self {
        SatSolver {
            _marker: PhantomData,
        }
    }
}

impl<S> SemanticsSolver for SatSolver<S>
where
    S: Solve + SolveIncremental + Default + 'static,
{
    fn extensions(&self, af: &AbstractFramework, sem: Semantics) -> Vec<Extension> {
        if af.arguments.is_empty() {
            return vec![BTreeSet::new()];
        }

        match sem {
            Semantics::Stable => enumerate_stable::<S>(af),
            Semantics::Complete => enumerate_complete::<S>(af),
            Semantics::Preferred => enumerate_preferred::<S>(af),
            Semantics::Grounded => vec![super::naive_solver::grounded_extension(af)],
            Semantics::Admissible => super::naive_solver::admissible_extensions(af),
            Semantics::SemiStable => enumerate_semi_stable::<S>(af),
            Semantics::Ideal => vec![compute_ideal::<S>(af)],
            Semantics::Eager => vec![compute_eager::<S>(af)],
        }
    }

    fn find_one(&self, af: &AbstractFramework, sem: Semantics) -> Option<Extension> {
        if af.arguments.is_empty() {
            return Some(BTreeSet::new());
        }

        match sem {
            Semantics::Admissible => find_one_admissible::<S>(af),
            Semantics::Stable => find_one_stable::<S>(af),
            Semantics::Complete => find_one_complete::<S>(af),
            Semantics::Preferred => find_one_preferred::<S>(af),
            Semantics::Grounded => Some(super::naive_solver::grounded_extension(af)),
            Semantics::SemiStable => enumerate_semi_stable::<S>(af).into_iter().next(),
            Semantics::Ideal => Some(compute_ideal::<S>(af)),
            Semantics::Eager => Some(compute_eager::<S>(af)),
        }
    }

    fn count(&self, af: &AbstractFramework, sem: Semantics) -> usize {
        if af.arguments.is_empty() {
            return 1;
        }

        match sem {
            Semantics::Stable => count_stable::<S>(af),
            Semantics::Complete => count_complete::<S>(af),
            Semantics::Preferred => count_preferred::<S>(af),
            Semantics::Admissible => super::naive_solver::admissible_extensions(af).len(),
            Semantics::Grounded | Semantics::Ideal | Semantics::Eager => 1,
            Semantics::SemiStable => enumerate_semi_stable::<S>(af).len(),
        }
    }
}

/// Bidirectional mapping between [`ArgumentId`]s and SAT variables.
///
/// Assigns each argument a unique 0-indexed propositional variable and
/// tracks auxiliary variables introduced by complete-closure and
/// preferred-maximisation encodings.
struct VarMap {
    id_to_var: HashMap<ArgumentId, Var>,
    var_to_id: BTreeMap<u32, ArgumentId>,
    next_aux: u32,
}

impl VarMap {
    #[allow(clippy::cast_possible_truncation)]
    fn new(args: &BTreeSet<ArgumentId>) -> Self {
        let mut id_to_var = HashMap::new();
        let mut var_to_id = BTreeMap::new();
        for (i, &arg) in args.iter().enumerate() {
            let var = Var::new(i as u32);
            id_to_var.insert(arg, var);
            var_to_id.insert(i as u32, arg);
        }
        VarMap {
            id_to_var,
            var_to_id,
            next_aux: args.len() as u32,
        }
    }

    fn lit(&self, arg: ArgumentId) -> Lit {
        self.id_to_var[&arg].pos_lit()
    }

    fn neg_lit(&self, arg: ArgumentId) -> Lit {
        self.id_to_var[&arg].neg_lit()
    }

    fn new_aux_var(&mut self) -> Var {
        let v = Var::new(self.next_aux);
        self.next_aux += 1;
        v
    }

    fn extract_extension(&self, solver: &impl Solve) -> Extension {
        self.var_to_id
            .iter()
            .filter(|(&idx, _)| {
                let var = Var::new(idx);
                matches!(solver.lit_val(var.pos_lit()), Ok(TernaryVal::True))
            })
            .map(|(_, &id)| id)
            .collect()
    }
}

/// Add conflict-free constraints: for each defeat (a, b), !`x_a` || !`x_b`.
fn add_conflict_free(solver: &mut impl Solve, af: &AbstractFramework, vm: &VarMap) {
    for &(from, to) in &af.defeats {
        solver.add_binary(vm.neg_lit(from), vm.neg_lit(to)).unwrap();
    }
}

/// Add admissibility (defence) constraints: for each argument b and each
/// attacker a of b, if b is in then some counter-attacker of a must be in.
///
/// Encoded as: !`x_b` || x_{d1} || x_{d2} || ... for each attacker a of b,
/// where `d_i` are the counter-attackers of a.
fn add_defence(solver: &mut impl Solve, af: &AbstractFramework, vm: &VarMap) {
    for &arg in &af.arguments {
        for &attacker in af.attackers_of(arg) {
            let defenders = af.attackers_of(attacker);
            let mut lits: Vec<Lit> = Vec::with_capacity(1 + defenders.len());
            lits.push(vm.neg_lit(arg));
            for &defender in defenders {
                lits.push(vm.lit(defender));
            }
            solver.add_clause_ref(lits.as_slice()).unwrap();
        }
    }
}

/// Add stable range constraints: for each argument b, b is in OR some
/// attacker of b is in.
fn add_stable_range(solver: &mut impl Solve, af: &AbstractFramework, vm: &VarMap) {
    for &arg in &af.arguments {
        let attackers = af.attackers_of(arg);
        let mut lits = Vec::with_capacity(1 + attackers.len());
        lits.push(vm.lit(arg));
        for &attacker in attackers {
            lits.push(vm.lit(attacker));
        }
        solver.add_clause_ref(lits.as_slice()).unwrap();
    }
}

/// Enumerate stable extensions via SAT.
///
/// Encoding: conflict-free + every non-member is attacked by a member.
fn enumerate_stable<S: Solve + SolveIncremental + Default>(
    af: &AbstractFramework,
) -> Vec<Extension> {
    let vm = VarMap::new(&af.arguments);
    let mut solver = S::default();
    add_conflict_free(&mut solver, af, &vm);
    add_stable_range(&mut solver, af, &vm);
    enumerate_models(&mut solver, &vm)
}

/// Add complete-extension closure constraints.
///
/// Auxiliary variable `defended_a`: true iff some member of the extension
/// attacks a. Then for each b: (all attackers defended) => b must be in.
fn add_complete_closure(solver: &mut impl Solve, af: &AbstractFramework, vm: &mut VarMap) {
    let mut defended_vars: HashMap<ArgumentId, Var> = HashMap::new();

    for &arg in &af.arguments {
        let attackers = af.attackers_of(arg);
        if attackers.is_empty() {
            continue;
        }
        for &attacker in attackers {
            if defended_vars.contains_key(&attacker) {
                continue;
            }
            let def_var = vm.new_aux_var();
            defended_vars.insert(attacker, def_var);

            let counter_attackers = af.attackers_of(attacker);

            for &d in counter_attackers {
                solver.add_binary(vm.neg_lit(d), def_var.pos_lit()).unwrap();
            }

            let mut lits = Vec::with_capacity(1 + counter_attackers.len());
            lits.push(def_var.neg_lit());
            for &d in counter_attackers {
                lits.push(vm.lit(d));
            }
            solver.add_clause_ref(lits.as_slice()).unwrap();

            if counter_attackers.is_empty() {
                solver.add_unit(def_var.neg_lit()).unwrap();
            }
        }
    }

    for &arg in &af.arguments {
        let attackers = af.attackers_of(arg);
        if attackers.is_empty() {
            solver.add_unit(vm.lit(arg)).unwrap();
        } else {
            let mut lits = Vec::with_capacity(1 + attackers.len());
            lits.push(vm.lit(arg));
            for &attacker in attackers {
                if let Some(&def_var) = defended_vars.get(&attacker) {
                    lits.push(def_var.neg_lit());
                }
            }
            solver.add_clause_ref(lits.as_slice()).unwrap();
        }
    }
}

/// Enumerate complete extensions via SAT.
///
/// Encoding: conflict-free + defence + characteristic-function closure.
fn enumerate_complete<S: Solve + SolveIncremental + Default>(
    af: &AbstractFramework,
) -> Vec<Extension> {
    let mut vm = VarMap::new(&af.arguments);
    let mut solver = S::default();
    add_conflict_free(&mut solver, af, &vm);
    add_defence(&mut solver, af, &vm);
    add_complete_closure(&mut solver, af, &mut vm);
    enumerate_models(&mut solver, &vm)
}

/// Enumerate preferred extensions via SAT.
///
/// Preferred = maximal admissible. Strategy: find admissible sets, maximize
/// each via incremental solving with assumption-guarded extension constraints,
/// then block and repeat.
fn enumerate_preferred<S: Solve + SolveIncremental + Default>(
    af: &AbstractFramework,
) -> Vec<Extension> {
    let mut vm = VarMap::new(&af.arguments);
    let mut solver = S::default();

    add_conflict_free(&mut solver, af, &vm);
    add_defence(&mut solver, af, &vm);

    let mut preferred = Vec::new();
    let args: Vec<ArgumentId> = af.arguments.iter().copied().collect();

    while let Ok(SolverResult::Sat) = solver.solve() {
        let mut current = vm.extract_extension(&solver);

        loop {
            let not_in: Vec<ArgumentId> = args
                .iter()
                .filter(|a| !current.contains(a))
                .copied()
                .collect();

            if not_in.is_empty() {
                break;
            }

            let act = vm.new_aux_var();
            let mut guard_clause = Vec::with_capacity(1 + not_in.len());
            guard_clause.push(act.neg_lit());
            for &a in &not_in {
                guard_clause.push(vm.lit(a));
            }
            solver.add_clause_ref(guard_clause.as_slice()).unwrap();

            let mut assumps: Vec<Lit> = Vec::with_capacity(current.len() + 1);
            for &a in &current {
                assumps.push(vm.lit(a));
            }
            assumps.push(act.pos_lit());

            match solver.solve_assumps(&assumps) {
                Ok(SolverResult::Sat) => {
                    current = vm.extract_extension(&solver);
                }
                _ => break,
            }
        }

        preferred.push(current.clone());

        if current.is_empty() {
            break;
        }
        let blocking: Vec<Lit> = current.iter().map(|&a| vm.neg_lit(a)).collect();
        solver.add_clause_ref(blocking.as_slice()).unwrap();
    }

    preferred.sort();
    preferred.dedup();
    filter_maximal(&preferred)
}

/// Enumerate all models by iteratively adding blocking clauses.
fn enumerate_models(solver: &mut impl SolveIncremental, vm: &VarMap) -> Vec<Extension> {
    let mut extensions = Vec::new();

    while let Ok(SolverResult::Sat) = solver.solve() {
        let ext = vm.extract_extension(solver);

        let blocking: Vec<Lit> = vm
            .var_to_id
            .keys()
            .map(|&idx| {
                let var = Var::new(idx);
                let val = solver.lit_val(var.pos_lit()).unwrap();
                if val == TernaryVal::True {
                    var.neg_lit()
                } else {
                    var.pos_lit()
                }
            })
            .collect();

        solver.add_clause_ref(blocking.as_slice()).unwrap();
        extensions.push(ext);
    }

    extensions
}

/// Count all models by iteratively adding blocking clauses without
/// allocating per-model extension sets.
fn count_models(solver: &mut impl SolveIncremental, vm: &VarMap) -> usize {
    let mut count = 0;

    while let Ok(SolverResult::Sat) = solver.solve() {
        count += 1;

        let blocking: Vec<Lit> = vm
            .var_to_id
            .keys()
            .map(|&idx| {
                let var = Var::new(idx);
                let val = solver.lit_val(var.pos_lit()).unwrap();
                if val == TernaryVal::True {
                    var.neg_lit()
                } else {
                    var.pos_lit()
                }
            })
            .collect();

        solver.add_clause_ref(blocking.as_slice()).unwrap();
    }

    count
}

/// Count stable extensions without materialising per-extension sets.
fn count_stable<S: Solve + SolveIncremental + Default>(af: &AbstractFramework) -> usize {
    let vm = VarMap::new(&af.arguments);
    let mut solver = S::default();
    add_conflict_free(&mut solver, af, &vm);
    add_stable_range(&mut solver, af, &vm);
    count_models(&mut solver, &vm)
}

/// Count complete extensions without materialising per-extension sets.
fn count_complete<S: Solve + SolveIncremental + Default>(af: &AbstractFramework) -> usize {
    let mut vm = VarMap::new(&af.arguments);
    let mut solver = S::default();
    add_conflict_free(&mut solver, af, &vm);
    add_defence(&mut solver, af, &vm);
    add_complete_closure(&mut solver, af, &mut vm);
    count_models(&mut solver, &vm)
}

/// Count preferred extensions by iterative maximisation with deduplication.
fn count_preferred<S: Solve + SolveIncremental + Default>(af: &AbstractFramework) -> usize {
    let mut vm = VarMap::new(&af.arguments);
    let mut solver = S::default();
    add_conflict_free(&mut solver, af, &vm);
    add_defence(&mut solver, af, &vm);

    let mut count = 0;
    let args: Vec<ArgumentId> = af.arguments.iter().copied().collect();
    let mut seen: Vec<Extension> = Vec::new();

    while let Ok(SolverResult::Sat) = solver.solve() {
        let mut current = vm.extract_extension(&solver);

        loop {
            let not_in: Vec<ArgumentId> = args
                .iter()
                .filter(|a| !current.contains(a))
                .copied()
                .collect();

            if not_in.is_empty() {
                break;
            }

            let act = vm.new_aux_var();
            let mut guard_clause = Vec::with_capacity(1 + not_in.len());
            guard_clause.push(act.neg_lit());
            for &a in &not_in {
                guard_clause.push(vm.lit(a));
            }
            solver.add_clause_ref(guard_clause.as_slice()).unwrap();

            let mut assumps: Vec<Lit> = Vec::with_capacity(current.len() + 1);
            for &a in &current {
                assumps.push(vm.lit(a));
            }
            assumps.push(act.pos_lit());

            match solver.solve_assumps(&assumps) {
                Ok(SolverResult::Sat) => {
                    current = vm.extract_extension(&solver);
                }
                _ => break,
            }
        }

        if !seen.iter().any(|s| current.is_subset(s)) {
            seen.push(current.clone());
            count += 1;
        }

        if current.is_empty() {
            break;
        }
        let blocking: Vec<Lit> = current.iter().map(|&a| vm.neg_lit(a)).collect();
        solver.add_clause_ref(blocking.as_slice()).unwrap();
    }

    count
}

/// Enumerate semi-stable extensions via SAT.
///
/// Semi-stable = complete extensions with maximal range (S ∪ S⁺).
/// (Caminada et al. 2012, Definition 2)
fn enumerate_semi_stable<S: Solve + SolveIncremental + Default>(
    af: &AbstractFramework,
) -> Vec<Extension> {
    let complete = enumerate_complete::<S>(af);
    if complete.is_empty() {
        return vec![];
    }

    let with_range: Vec<(Extension, usize)> = complete
        .into_iter()
        .map(|ext| {
            let attacked: BTreeSet<ArgumentId> = ext
                .iter()
                .flat_map(|&a| af.attacked_by(a).iter().copied())
                .collect();
            let range_size = ext.len() + attacked.difference(&ext).count();
            (ext, range_size)
        })
        .collect();

    let max_range = with_range.iter().map(|(_, r)| *r).max().unwrap_or(0);

    with_range
        .into_iter()
        .filter(|(_, r)| *r == max_range)
        .map(|(ext, _)| ext)
        .collect()
}

/// Find the unique maximal admissible subset of `candidates` within the
/// given framework. Used by ideal and eager extension computation.
fn maximal_admissible_in<S: Solve + SolveIncremental + Default>(
    af: &AbstractFramework,
    candidates: &BTreeSet<ArgumentId>,
) -> Extension {
    if candidates.is_empty() {
        return BTreeSet::new();
    }

    let mut vm = VarMap::new(&af.arguments);
    let mut solver = S::default();

    add_conflict_free(&mut solver, af, &vm);
    add_defence(&mut solver, af, &vm);

    for &arg in &af.arguments {
        if !candidates.contains(&arg) {
            solver.add_unit(vm.neg_lit(arg)).unwrap();
        }
    }

    let mut current = match solver.solve() {
        Ok(SolverResult::Sat) => vm.extract_extension(&solver),
        _ => return BTreeSet::new(),
    };

    let candidate_args: Vec<ArgumentId> = candidates.iter().copied().collect();

    loop {
        let not_in: Vec<ArgumentId> = candidate_args
            .iter()
            .filter(|a| !current.contains(a))
            .copied()
            .collect();

        if not_in.is_empty() {
            break;
        }

        let act = vm.new_aux_var();
        let mut guard_clause = Vec::with_capacity(1 + not_in.len());
        guard_clause.push(act.neg_lit());
        for &a in &not_in {
            guard_clause.push(vm.lit(a));
        }
        solver.add_clause_ref(guard_clause.as_slice()).unwrap();

        let mut assumps: Vec<Lit> = Vec::with_capacity(current.len() + 1);
        for &a in &current {
            assumps.push(vm.lit(a));
        }
        assumps.push(act.pos_lit());

        match solver.solve_assumps(&assumps) {
            Ok(SolverResult::Sat) => {
                current = vm.extract_extension(&solver);
            }
            _ => break,
        }
    }

    current
}

/// Compute the ideal extension via SAT.
///
/// Ideal = maximal admissible subset of the intersection of all preferred
/// extensions. (Dung et al. 2007, Theorem 3.3)
fn compute_ideal<S: Solve + SolveIncremental + Default>(af: &AbstractFramework) -> Extension {
    let preferred = enumerate_preferred::<S>(af);
    if preferred.is_empty() {
        return BTreeSet::new();
    }

    let mut intersection = preferred[0].clone();
    for ext in &preferred[1..] {
        intersection = intersection.intersection(ext).copied().collect();
    }

    maximal_admissible_in::<S>(af, &intersection)
}

/// Compute the eager extension via SAT.
///
/// Eager = maximal admissible subset of the intersection of all semi-stable
/// extensions. (Caminada 2007)
fn compute_eager<S: Solve + SolveIncremental + Default>(af: &AbstractFramework) -> Extension {
    let semi_stable = enumerate_semi_stable::<S>(af);
    if semi_stable.is_empty() {
        return BTreeSet::new();
    }

    let mut intersection = semi_stable[0].clone();
    for ext in &semi_stable[1..] {
        intersection = intersection.intersection(ext).copied().collect();
    }

    maximal_admissible_in::<S>(af, &intersection)
}

// ---------------------------------------------------------------------------
// find_one helpers — single-solve variants for each semantics
// ---------------------------------------------------------------------------

fn find_one_admissible<S: Solve + Default>(af: &AbstractFramework) -> Option<Extension> {
    let vm = VarMap::new(&af.arguments);
    let mut solver = S::default();
    add_conflict_free(&mut solver, af, &vm);
    add_defence(&mut solver, af, &vm);
    match solver.solve() {
        Ok(SolverResult::Sat) => Some(vm.extract_extension(&solver)),
        _ => None,
    }
}

fn find_one_stable<S: Solve + Default>(af: &AbstractFramework) -> Option<Extension> {
    let vm = VarMap::new(&af.arguments);
    let mut solver = S::default();
    add_conflict_free(&mut solver, af, &vm);
    add_stable_range(&mut solver, af, &vm);
    match solver.solve() {
        Ok(SolverResult::Sat) => Some(vm.extract_extension(&solver)),
        _ => None,
    }
}

fn find_one_complete<S: Solve + Default>(af: &AbstractFramework) -> Option<Extension> {
    let mut vm = VarMap::new(&af.arguments);
    let mut solver = S::default();
    add_conflict_free(&mut solver, af, &vm);
    add_defence(&mut solver, af, &vm);
    add_complete_closure(&mut solver, af, &mut vm);
    match solver.solve() {
        Ok(SolverResult::Sat) => Some(vm.extract_extension(&solver)),
        _ => None,
    }
}

/// Find one preferred extension by finding an admissible set and maximising
/// it once, without enumerating further.
fn find_one_preferred<S: Solve + SolveIncremental + Default>(
    af: &AbstractFramework,
) -> Option<Extension> {
    let mut vm = VarMap::new(&af.arguments);
    let mut solver = S::default();
    add_conflict_free(&mut solver, af, &vm);
    add_defence(&mut solver, af, &vm);

    let args: Vec<ArgumentId> = af.arguments.iter().copied().collect();

    let mut current = match solver.solve() {
        Ok(SolverResult::Sat) => vm.extract_extension(&solver),
        _ => return None,
    };

    loop {
        let not_in: Vec<ArgumentId> = args
            .iter()
            .filter(|a| !current.contains(a))
            .copied()
            .collect();

        if not_in.is_empty() {
            break;
        }

        let act = vm.new_aux_var();
        let mut guard_clause = Vec::with_capacity(1 + not_in.len());
        guard_clause.push(act.neg_lit());
        for &a in &not_in {
            guard_clause.push(vm.lit(a));
        }
        solver.add_clause_ref(guard_clause.as_slice()).unwrap();

        let mut assumps: Vec<Lit> = Vec::with_capacity(current.len() + 1);
        for &a in &current {
            assumps.push(vm.lit(a));
        }
        assumps.push(act.pos_lit());

        match solver.solve_assumps(&assumps) {
            Ok(SolverResult::Sat) => {
                current = vm.extract_extension(&solver);
            }
            _ => break,
        }
    }

    Some(current)
}

/// Filter `sets` to retain only those not properly contained in another.
fn filter_maximal(sets: &[Extension]) -> Vec<Extension> {
    crate::algorithms::maximal::maximal_by_inclusion(sets)
}
