//! Define the main evaluation stack of the Nickel abstract machine and related operations.
//!
//! See [eval](../eval/index.html).
use super::cache::{Cache, CacheIndex};
use super::operation::OperationCont;
use crate::eval::{Closure, Environment};
use crate::position::TermPos;
use crate::term::{BindingType, RichTerm, StrChunk};

pub struct StrAccData {
    /// The accumulator.
    pub acc: String,
    /// The common environment of chunks.
    pub env: Environment,
    /// The indentation level of the chunk currently being evaluated.
    pub curr_indent: usize,
    /// The position of the original (unevaluated) expression of the chunk currently being
    /// evaluated.
    pub curr_pos: TermPos,
}

/// An element of the stack.
pub enum Marker<C: Cache> {
    /// An equality to test.
    ///
    /// When evaluating one equality `t1 == t2`, the abstract machine may generate several new
    /// equalities to test (for example, take `t1` and `t2` to be `[1, 2, 3]`). In this case, the
    /// first equality is evaluated and the remaining ones - the continuation of the whole
    /// computation - are put on the stack as `Eq` elements. If an equality evaluates to `false` at
    /// some point, all the consecutive `Eq` elements at the top of the stack are discarded.
    Eq(Closure, Closure),

    /// An argument of an application.
    Arg(Closure, TermPos),

    /// A tracked argument. Behave the same as a standard argument, but is given directly as a cache
    /// index, such that it can be shared with other part of the program.
    ///
    /// In particular, contract arguments are tracked, in order to report the actual, evaluated
    /// offending term in case of blame.
    TrackedArg(CacheIndex, TermPos),

    /// An update index, which is a pointer to a mutable memory cell to be updated.
    UpdateIndex(C::UpdateIndex),

    /// The continuation of a primitive operation.
    Cont(
        OperationCont,
        usize,   /*callStack size*/
        TermPos, /*position span of the operation*/
    ),

    /// A string chunk.
    ///
    /// Generated by `ChunksConcat` when evaluating a string with interpolated expressions. Such
    /// strings are represented by a list of chunks to evaluate and concatenate, a chunk being
    /// either an interpolated expression or a string literals. The shared environment is stored in
    /// the top element of the stack, which must be `StrAcc`.
    StrChunk(StrChunk<RichTerm>),

    /// A string accumulator. Used by `ChunksConcat` to store additional state, that is the string
    /// being constructed, the indentation of the chunk being evaluated, and the common initial
    /// environment of chunks.
    StrAcc(StrAccData),
}

impl<C: Cache> std::fmt::Debug for Marker<C> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Marker::Eq(_, _) => write!(f, "Eq"),
            Marker::Arg(_, _) => write!(f, "Arg"),
            Marker::TrackedArg(_, _) => write!(f, "TrackedArg"),
            Marker::UpdateIndex(_) => write!(f, "UpdateIndex"),
            Marker::Cont(op, sz, _) => write!(f, "Cont {op:?} (callstack size {sz})"),
            Marker::StrChunk(_) => write!(f, "StrChunk"),
            Marker::StrAcc { .. } => write!(f, "StrAcc"),
        }
    }
}

impl<C: Cache> Marker<C> {
    pub fn is_arg(&self) -> bool {
        matches!(*self, Marker::Arg(..) | Marker::TrackedArg(..))
    }

    pub fn is_idx(&self) -> bool {
        matches!(*self, Marker::UpdateIndex(..))
    }

    pub fn is_cont(&self) -> bool {
        matches!(*self, Marker::Cont(..))
    }

    pub fn is_eq(&self) -> bool {
        matches!(*self, Marker::Eq(..))
    }

    pub fn is_str_chunk(&self) -> bool {
        matches!(*self, Marker::StrChunk(..))
    }

    pub fn is_str_acc(&self) -> bool {
        matches!(*self, Marker::StrAcc { .. })
    }
}

/// The evaluation stack.
#[derive(Default)]
pub struct Stack<C: Cache>(Vec<Marker<C>>);

impl<C: Cache> IntoIterator for Stack<C> {
    type Item = Marker<C>;
    type IntoIter = ::std::vec::IntoIter<Marker<C>>;

    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

impl<C: Cache> Stack<C> {
    pub fn new() -> Stack<C> {
        Stack(Vec::new())
    }

    /// Count the number of consecutive elements satisfying `pred` from the top of the stack.
    fn count<P>(&self, pred: P) -> usize
    where
        P: Fn(&Marker<C>) -> bool,
    {
        let mut count = 0;
        for marker in self.0.iter().rev() {
            if pred(marker) {
                count += 1;
            } else {
                break;
            }
        }
        count
    }

    /// Pops all items in the stack and resets the state of the [Cache] elements it encounters.
    pub fn reset(&mut self, cache: &mut C) {
        while let Some(marker) = self.0.pop() {
            if let Marker::UpdateIndex(mut uidx) = marker {
                cache.reset_index_state(&mut uidx);
            }
        }
    }

    /// Count the number of arguments at the top of the stack.
    pub fn count_args(&self) -> usize {
        Stack::count(self, Marker::is_arg)
    }

    pub fn push_arg(&mut self, arg: Closure, pos: TermPos) {
        self.0.push(Marker::Arg(arg, pos))
    }

    pub fn push_tracked_arg(&mut self, idx: CacheIndex, pos: TermPos) {
        self.0.push(Marker::TrackedArg(idx, pos))
    }

    pub fn push_update_index(&mut self, uidx: C::UpdateIndex) {
        self.0.push(Marker::UpdateIndex(uidx))
    }

    pub fn push_op_cont(&mut self, cont: OperationCont, len: usize, pos: TermPos) {
        self.0.push(Marker::Cont(cont, len, pos))
    }

    /// Push a sequence of equalities on the stack.
    pub fn push_eqs<I>(&mut self, it: I)
    where
        I: Iterator<Item = (Closure, Closure)>,
    {
        self.0.extend(it.map(|(t1, t2)| Marker::Eq(t1, t2)));
    }

    /// Push a sequence of string chunks on the stack.
    pub fn push_str_chunks<I>(&mut self, it: I)
    where
        I: Iterator<Item = StrChunk<RichTerm>>,
    {
        self.0.extend(it.map(Marker::StrChunk));
    }

    /// Push a string accumulator on the stack.
    pub fn push_str_acc(&mut self, str_acc_data: StrAccData) {
        self.0.push(Marker::StrAcc(str_acc_data));
    }

    /// Try to pop an argument from the top of the stack. If `None` is returned, the top element
    /// was not an argument and the stack is left unchanged.
    ///
    /// If the argument is tracked, it is automatically converted into an owned closure.
    pub fn pop_arg(&mut self, cache: &C) -> Option<(Closure, TermPos)> {
        match self.0.pop() {
            Some(Marker::Arg(arg, pos)) => Some((arg, pos)),
            Some(Marker::TrackedArg(arg_idx, pos)) => Some((cache.get(arg_idx), pos)),
            Some(m) => {
                self.0.push(m);
                None
            }
            _ => None,
        }
    }

    /// Try to pop an argument from the top of the stack and return it as an index. If `None` is
    /// returned, the top element was not an argument and the stack is left unchanged.
    ///
    /// If the argument is not tracked, it is directly returned.
    pub fn pop_arg_as_idx(&mut self, cache: &mut C) -> Option<(CacheIndex, TermPos)> {
        match self.0.pop() {
            Some(Marker::Arg(arg, pos)) => Some((cache.add(arg, BindingType::Normal), pos)),
            Some(Marker::TrackedArg(arg_idx, pos)) => Some((arg_idx, pos)),
            Some(m) => {
                self.0.push(m);
                None
            }
            _ => None,
        }
    }

    /// Try to pop an index from the top of the stack. If `None` is returned, the top element was
    /// not an index and the stack is left unchanged.
    pub fn pop_update_index(&mut self) -> Option<C::UpdateIndex> {
        match self.0.pop() {
            Some(Marker::UpdateIndex(uidx)) => Some(uidx),
            Some(m) => {
                self.0.push(m);
                None
            }
            _ => None,
        }
    }

    /// Try to pop an operator continuation from the top of the stack. If `None` is returned, the
    /// top element was not an operator continuation and the stack is left unchanged.
    pub fn pop_op_cont(&mut self) -> Option<(OperationCont, usize, TermPos)> {
        match self.0.pop() {
            Some(Marker::Cont(cont, len, pos)) => Some((cont, len, pos)),
            Some(m) => {
                self.0.push(m);
                None
            }
            _ => None,
        }
    }

    pub fn peek_op_cont(&self) -> Option<OperationCont> {
        let mut it = self
            .0
            .iter()
            .rev()
            .skip_while(|&marker| matches!(marker, Marker::UpdateIndex(..)));

        match it.next() {
            Some(Marker::Cont(cont, _, _)) => Some(cont.clone()),
            _ => None,
        }
    }

    /// Try to pop an equality from the top of the stack. If `None` is returned, the top element
    /// was not an equality and the stack is left unchanged.
    pub fn pop_eq(&mut self) -> Option<(Closure, Closure)> {
        if self.0.last().map(Marker::is_eq).unwrap_or(false) {
            match self.0.pop() {
                Some(Marker::Eq(c1, c2)) => Some((c1, c2)),
                _ => unreachable!(),
            }
        } else {
            None
        }
    }

    /// Try to pop the a string accumulator from the stack. If `None` is returned, the top element
    /// was not a string accumulator and the stack is left unchanged.
    pub fn pop_str_acc(&mut self) -> Option<StrAccData> {
        if self.0.last().map(Marker::is_str_acc).unwrap_or(false) {
            match self.0.pop() {
                Some(Marker::StrAcc(str_acc_data)) => Some(str_acc_data),
                _ => unreachable!(),
            }
        } else {
            None
        }
    }

    /// Try to pop a string chunk from the top of the stack. If `None` is returned, the top element
    /// was not a string chunk and the stack is left unchanged.
    pub fn pop_str_chunk(&mut self) -> Option<StrChunk<RichTerm>> {
        if self.0.last().map(Marker::is_str_chunk).unwrap_or(false) {
            match self.0.pop() {
                Some(Marker::StrChunk(c)) => Some(c),
                _ => unreachable!(),
            }
        } else {
            None
        }
    }

    /// Check if the top element is a [CacheIndex].
    pub fn is_top_idx(&self) -> bool {
        self.0.last().map(Marker::is_idx).unwrap_or(false)
    }

    /// Check if the top element is an operation continuation.
    pub fn is_top_cont(&self) -> bool {
        self.0.last().map(Marker::is_cont).unwrap_or(false)
    }

    /// Discard all the consecutive equality from the top of the stack. This drops the continuation
    /// of the equality being currently evaluated.
    pub fn clear_eqs(&mut self) {
        while self.pop_eq().is_some() {}
    }

    /// Turning the top element of the stack into a tracked arg if it was not already. Returns the
    /// corresponding index, or `None` if the top element wasn't an argument.
    pub fn track_arg(&mut self, cache: &mut C) -> Option<CacheIndex> {
        match self.0.last_mut() {
            Some(Marker::TrackedArg(idx, _)) => Some(idx.clone()),
            Some(Marker::Arg(..)) => {
                let (closure, pos) = self.pop_arg(cache).unwrap();
                let idx = cache.add(closure, BindingType::Normal);
                self.push_tracked_arg(idx.clone(), pos);
                Some(idx)
            }
            _ => None,
        }
    }
}

impl<C: Cache> std::fmt::Debug for Stack<C> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "--- STACK ---")?;
        for marker in self.0.iter().rev() {
            writeln!(f, "| {marker:?}")?;
        }
        writeln!(f, "---  END  ---")
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::eval::cache::CacheImpl;
    use crate::term::{Term, UnaryOp};
    use assert_matches::assert_matches;

    impl Stack<CacheImpl> {
        /// Count the number of indices at the top of the stack.
        pub fn count_thunks(&self) -> usize {
            Stack::count(self, Marker::is_idx)
        }

        /// Count the number of operation continuation at the top of the stack.
        pub fn count_conts(&self) -> usize {
            Stack::count(self, Marker::is_cont)
        }
    }

    fn some_closure() -> Closure {
        Closure::atomic_closure(Term::Bool(true).into())
    }

    fn some_cont() -> OperationCont {
        OperationCont::Op1(UnaryOp::Typeof, TermPos::None)
    }

    fn some_arg_marker() -> Marker<CacheImpl> {
        Marker::Arg(some_closure(), TermPos::None)
    }

    fn some_thunk_marker(eval_cache: &mut CacheImpl) -> Marker<CacheImpl> {
        let mut idx = eval_cache.add(some_closure(), BindingType::Normal);
        let uidx = eval_cache.make_update_index(&mut idx).unwrap();
        Marker::UpdateIndex(uidx)
    }

    fn some_cont_marker() -> Marker<CacheImpl> {
        Marker::Cont(some_cont(), 42, TermPos::None)
    }

    #[test]
    fn marker_differentiates() {
        let mut eval_cache = CacheImpl::new();
        assert!(some_arg_marker().is_arg());
        assert!(some_thunk_marker(&mut eval_cache).is_idx());
        assert!(some_cont_marker().is_cont());
    }

    #[test]
    fn pushing_and_popping_args() {
        let mut s = Stack::new();
        assert_eq!(0, s.count_args());

        s.push_arg(some_closure(), TermPos::None);
        s.push_arg(some_closure(), TermPos::None);
        assert_eq!(2, s.count_args());
        assert_eq!(
            some_closure(),
            s.pop_arg(&CacheImpl::new()).expect("Already checked").0
        );
        assert_eq!(1, s.count_args());
    }

    #[test]
    fn pushing_and_popping_thunks() {
        let mut s = Stack::new();
        assert_eq!(0, s.count_thunks());

        let mut eval_cache = CacheImpl::new();

        let mut idx = eval_cache.add(some_closure(), BindingType::Normal);
        s.push_update_index(eval_cache.make_update_index(&mut idx).unwrap());
        idx = eval_cache.add(some_closure(), BindingType::Normal);
        s.push_update_index(eval_cache.make_update_index(&mut idx).unwrap());

        assert_eq!(2, s.count_thunks());
        s.pop_update_index().expect("Already checked");
        assert_eq!(1, s.count_thunks());
    }

    #[test]
    fn thunk_blackhole() {
        let mut eval_cache = CacheImpl::new();
        let mut idx = eval_cache.add(some_closure(), BindingType::Normal);
        let idx_upd = eval_cache.make_update_index(&mut idx);
        assert_matches!(idx_upd, Ok(..));
        assert_matches!(eval_cache.make_update_index(&mut idx), Err(..));
        eval_cache.update(some_closure(), idx_upd.unwrap());
        assert_matches!(eval_cache.make_update_index(&mut idx), Ok(..));
    }

    #[test]
    fn pushing_and_popping_conts() {
        let mut s = Stack::new();
        assert_eq!(0, s.count_conts());

        s.push_op_cont(some_cont(), 3, TermPos::None);
        s.push_op_cont(some_cont(), 4, TermPos::None);
        assert_eq!(2, s.count_conts());
        assert_eq!(
            (some_cont(), 4, TermPos::None),
            s.pop_op_cont().expect("Already checked")
        );
        assert_eq!(1, s.count_conts());
    }
}
