use owo_colors::OwoColorize;
use owo_colors::colors::Green;

use crate::ast::{Module, Predicate};
use crate::prover::body::BodyProver;
use crate::prover::constraints::Constraints;
use crate::prover::tracing::{ProofType, Tracer};
use crate::prover::{Counter, GlobalCounter};

pub struct PredicateProver<'a, T: Tracer>
{
    module: &'a Module,
    predicate: Predicate,
    constraints: Constraints,
    counter_snapshot: Counter,
    tracer: T,

    global_counter: GlobalCounter,
    current_clause: usize,
    sub_proof: Option<BodyProver<'a>>,
}

impl<'a, T: Tracer + 'a> PredicateProver<'a, T>
{
    pub fn new(
        module: &'a Module,
        global_counter: GlobalCounter,
        predicate: Predicate,
        constraints: Constraints,
        tracer: &mut T,
    ) -> PredicateProver<'a, T>
    {
        //info!(target: "PredicateProver", "Proving {}", predicate);
        let mut predicate_prover = tracer.begin_proof(ProofType::Predicate);
        predicate_prover.print_step(format!(
            "{} '{}'",
            "Proving predicate".fg::<Green>(),
            predicate
        ));
        PredicateProver {
            module,
            tracer: predicate_prover,
            predicate,
            constraints,
            current_clause: 0,
            sub_proof: None,
            counter_snapshot: global_counter.snapshot(),
            global_counter,
        }
    }
}

impl<'a, T: Tracer + 'a> Iterator for PredicateProver<'a, T>
{
    type Item = Constraints;

    fn next(&mut self) -> Option<Self::Item>
    {
        match self.sub_proof.as_mut()
        {
            None =>
            {
                if self.current_clause == self.module.clauses.len()
                {
                    None
                }
                else
                {
                    let clause = &self.module.clauses[self.current_clause]
                        .make_unique(self.global_counter.clone());
                    let uni = self.predicate.matches(&clause.head);
                    let full_constraints = uni.and_then(|x| x.and(&self.constraints));

                    if let Some(c) = &full_constraints
                    {
                        self.tracer.print_step(format!(
                            "Unifying '{}' aginst '{}'",
                            self.predicate, clause
                        ));
                        self.tracer.print_step(format!(
                            "↳ {}: {}",
                            "Matches".fg::<Green>(),
                            c.simplified()
                        ));
                    }

                    match full_constraints
                    {
                        Some(constraints) if clause.body.is_none() =>
                        {
                            self.current_clause += 1;
                            Some(constraints)
                        }
                        Some(constraints) =>
                        {
                            self.current_clause += 1;
                            self.counter_snapshot = self.global_counter.snapshot();
                            self.sub_proof = Some(BodyProver::new(
                                self.module,
                                self.global_counter.clone(),
                                clause.body.clone().unwrap(),
                                constraints,
                                &mut self.tracer,
                            ));
                            self.next()
                        }
                        None =>
                        {
                            self.global_counter.restore(self.counter_snapshot);
                            self.current_clause += 1;
                            self.next()
                        }
                    }
                }
            }
            Some(prover) =>
            {
                let next = prover.next();
                match next
                {
                    Some(constraints) => Some(constraints),
                    None =>
                    {
                        self.global_counter.restore(self.counter_snapshot);
                        self.sub_proof = None;
                        self.next()
                    }
                }
            }
        }
    }
}