use std::cmp::Ordering;
use std::collections::BTreeMap;

use crate::atom_table;
use crate::heap_iter::{stackful_post_order_iter, NonListElider};
use crate::machine::machine_indices::VarKey;
use crate::machine::mock_wam::CompositeOpDir;
use crate::machine::{
    F64Offset, F64Ptr, Fixnum, Number, BREAK_FROM_DISPATCH_LOOP_LOC, LIB_QUERY_SUCCESS,
};
use crate::parser::ast::{Var, VarPtr};
use crate::parser::parser::{Parser, Tokens};
use crate::read::{write_term_to_heap, TermWriteResult};

use dashu::{Integer, Rational};
use indexmap::IndexMap;

use super::{streams::Stream, Atom, AtomCell, HeapCellValue, HeapCellValueTag, Machine};

#[cfg(test)]
mod lib_machine_tests;

/// Represents a leaf answer from a query.
#[derive(Debug, Clone, PartialEq)]
pub enum LeafAnswer {
    /// A `true` leaf answer.
    True,
    /// A `false` leaf answer.
    ///
    /// This means that there are no more answers for the query.
    False,
    /// An exception leaf answer.
    Exception(Term),
    /// A leaf answer with bindings and residual goals.
    LeafAnswer {
        /// The bindings of variables in the query.
        ///
        /// Can be empty.
        bindings: BTreeMap<String, Term>,
        /// Residual goals.
        ///
        /// Can be empty.
        residual_goals: Vec<Term>,
    },
}

/// Represents a Prolog term.
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq)]
pub enum Term {
    /// An arbitrary precision integer.
    Integer(Integer),
    /// An arbitrary precision rational.
    Rational(Rational),
    /// A float.
    Float(f64),
    /// A Prolog atom.
    Atom(String),
    /// A Prolog string.
    ///
    /// In particular, this represents Prolog lists of characters.
    String(String),
    /// A Prolog list.
    List(Vec<Term>),
    /// A Prolog compound term.
    Compound(String, Vec<Term>),
    /// A Prolog variable.
    Var(String),
}

impl Term {
    /// Creates an integer term.
    pub fn integer(value: impl Into<Integer>) -> Self {
        Term::Integer(value.into())
    }

    /// Creates a rational term.
    pub fn rational(value: impl Into<Rational>) -> Self {
        Term::Rational(value.into())
    }

    /// Creates a float term.
    pub fn float(value: impl Into<f64>) -> Self {
        Term::Float(value.into())
    }

    /// Creates an atom term.
    pub fn atom(value: impl Into<String>) -> Self {
        Term::Atom(value.into())
    }

    /// Creates a string term.
    ///
    /// In specific, this represents a list of chars in Prolog.
    pub fn string(value: impl Into<String>) -> Self {
        Term::String(value.into())
    }

    /// Creates a list term.
    pub fn list(value: impl IntoIterator<Item = Term>) -> Self {
        Term::List(value.into_iter().collect())
    }

    /// Creates a compound term.
    pub fn compound(functor: impl Into<String>, args: impl IntoIterator<Item = Term>) -> Self {
        Term::Compound(functor.into(), args.into_iter().collect())
    }

    /// Creates a variable.
    pub fn variable(value: impl Into<String>) -> Self {
        Term::Var(value.into())
    }

    /// Creates a conjunction, giving the atom `true` if empty.
    pub fn conjunction(value: impl IntoIterator<Item = Term>) -> Self {
        Term::try_conjunction(value).unwrap_or(Term::atom("true"))
    }

    /// Creates a conjunction, giving `None` if empty.
    pub fn try_conjunction(value: impl IntoIterator<Item = Term>) -> Option<Self> {
        let mut iter = value.into_iter();
        iter.next().map(|first| {
            Term::try_conjunction(iter)
                .map(|rest| Term::compound(",", [first.clone(), rest]))
                .unwrap_or(first)
        })
    }

    /// Creates a disjunction, giving the atom `false` if empty.
    pub fn disjunction(value: impl IntoIterator<Item = Term>) -> Self {
        Term::try_disjunction(value).unwrap_or(Term::atom("false"))
    }

    /// Creates a disjunction, giving `None` if empty.
    pub fn try_disjunction(value: impl IntoIterator<Item = Term>) -> Option<Self> {
        let mut iter = value.into_iter();
        iter.next().map(|first| {
            Term::try_disjunction(iter)
                .map(|rest| Term::compound(";", [first.clone(), rest]))
                .unwrap_or(first)
        })
    }
}

/// This is an auxiliary function to turn a count into names of anonymous variables like _A, _B,
/// _AB, etc...
fn count_to_letter_code(mut count: usize) -> String {
    let mut letters = Vec::new();

    loop {
        let letter_idx = (count % 26) as u32;
        letters.push(char::from_u32('A' as u32 + letter_idx).unwrap());
        count /= 26;

        if count == 0 {
            break;
        }
    }

    letters.into_iter().chain("_".chars()).rev().collect()
}

impl Term {
    pub(crate) fn from_heapcell(
        machine: &mut Machine,
        heap_cell: HeapCellValue,
        var_names: &mut IndexMap<HeapCellValue, VarPtr>,
    ) -> Self {
        // Adapted from MachineState::read_term_from_heap
        let mut term_stack = vec![];
        let iter = stackful_post_order_iter::<NonListElider>(
            &mut machine.machine_st.heap,
            &mut machine.machine_st.stack,
            heap_cell,
        );

        let mut anon_count: usize = 0;
        let var_ptr_cmp = |a, b| match a {
            Var::Named(name_a) => match b {
                Var::Named(name_b) => name_a.cmp(&name_b),
                _ => Ordering::Less,
            },
            _ => match b {
                Var::Named(_) => Ordering::Greater,
                _ => Ordering::Equal,
            },
        };

        for addr in iter {
            let addr = unmark_cell_bits!(addr);

            read_heap_cell!(addr,
                (HeapCellValueTag::Lis) => {
                    let tail = term_stack.pop().unwrap();
                    let head = term_stack.pop().unwrap();

                    let list = match tail {
                        Term::Atom(atom) if atom == "[]" => match head {
                            Term::Atom(ref a) if a.chars().collect::<Vec<_>>().len() == 1 => {
                                // Handle lists of char as strings
                                Term::String(a.to_string())
                            }
                            _ => Term::List(vec![head]),
                        },
                        Term::List(elems) if elems.is_empty() => match head {
                            Term::Atom(ref a) if a.chars().collect::<Vec<_>>().len() == 1 => {
                                // Handle lists of char as strings
                                Term::String(a.to_string())
                            },
                            _ => Term::List(vec![head]),
                        },
                        Term::List(mut elems) => {
                            elems.insert(0, head);
                            Term::List(elems)
                        },
                        Term::String(mut elems) => match head {
                            Term::Atom(ref a) if a.chars().collect::<Vec<_>>().len() == 1 => {
                                // Handle lists of char as strings
                                elems.insert(0, a.chars().next().unwrap());
                                Term::String(elems)
                            },
                            _ => {
                                let mut elems: Vec<Term> = elems
                                    .chars()
                                    .map(|x| Term::Atom(x.into()))
                                    .collect();
                                elems.insert(0, head);
                                Term::List(elems)
                            }
                        },
                        _ => {
                            Term::Compound(".".into(), vec![head, tail])
                        }
                    };
                    term_stack.push(list);
                }
                (HeapCellValueTag::Var | HeapCellValueTag::AttrVar | HeapCellValueTag::StackVar) => {
                    let var = var_names.get(&addr).map(|x| x.borrow().clone());
                    match var {
                        Some(Var::Named(name)) => term_stack.push(Term::Var(name)),
                        _ => {
                            let anon_name = loop {
                                // Generate a name for the anonymous variable
                                let anon_name = count_to_letter_code(anon_count);

                                // Find if this name is already being used
                                var_names.sort_by(|_, a, _, b| {
                                    var_ptr_cmp(a.borrow().clone(), b.borrow().clone())
                                });
                                let binary_result = var_names.binary_search_by(|_,a| {
                                    let var_ptr = Var::Named(anon_name.clone());
                                    var_ptr_cmp(a.borrow().clone(), var_ptr.clone())
                                });

                                match binary_result {
                                    Ok(_) => anon_count += 1, // Name already used
                                    Err(_) => {
                                        // Name not used, assign it to this variable
                                        let var_ptr = VarPtr::from(Var::Named(anon_name.clone()));
                                        var_names.insert(addr, var_ptr);
                                        break anon_name;
                                    },
                                }
                            };
                            term_stack.push(Term::Var(anon_name));
                        },
                    }
                }
                (HeapCellValueTag::F64, f) => {
                    term_stack.push(Term::Float((*f).into()));
                }
                (HeapCellValueTag::Char, c) => {
                    term_stack.push(Term::Atom(c.into()));
                }
                (HeapCellValueTag::Fixnum, n) => {
                    term_stack.push(Term::Integer(n.into()));
                }
                (HeapCellValueTag::Cons) => {
                    match Number::try_from(addr) {
                        Ok(Number::Integer(i)) => term_stack.push(Term::Integer((*i).clone())),
                        Ok(Number::Rational(r)) => term_stack.push(Term::Rational((*r).clone())),
                        _ => {}
                    }
                }
                (HeapCellValueTag::CStr, s) => {
                    term_stack.push(Term::String(s.as_str().to_string()));
                }
                (HeapCellValueTag::Atom, (name, arity)) => {
                    //let h = iter.focus().value() as usize;
                    //let mut arity = arity;

                    // Not sure why/if this is needed.
                    // Might find out with better testing later.
                    /*
                    if iter.heap.len() > h + arity + 1 {
                        let value = iter.heap[h + arity + 1];

                        if let Some(idx) = get_structure_index(value) {
                            // in the second condition, arity == 0,
                            // meaning idx cannot pertain to this atom
                            // if it is the direct subterm of a larger
                            // structure.
                            if arity > 0 || !iter.direct_subterm_of_str(h) {
                                term_stack.push(
                                    Term::Literal(Cell::default(), Literal::CodeIndex(idx))
                                );

                                arity += 1;
                            }
                        }
                    }
                    */

                    if arity == 0 {
                        let atom_name = name.as_str().to_string();
                        if atom_name == "[]" {
                            term_stack.push(Term::List(vec![]));
                        } else {
                            term_stack.push(Term::Atom(atom_name));
                        }
                    } else {
                        let subterms = term_stack
                            .drain(term_stack.len() - arity ..)
                            .collect();

                        term_stack.push(Term::Compound(name.as_str().to_string(), subterms));
                    }
                }
                (HeapCellValueTag::PStr, atom) => {
                    let tail = term_stack.pop().unwrap();

                    match tail {
                        Term::Atom(atom) => {
                            if atom == "[]" {
                                term_stack.push(Term::String(atom.as_str().to_string()));
                            }
                        },
                        Term::List(l) => {
                            let mut list: Vec<Term> = atom
                                .as_str()
                                .to_string()
                                .chars()
                                .map(|x| Term::Atom(x.to_string()))
                                .collect();
                            list.extend(l.into_iter());
                            term_stack.push(Term::List(list));
                        },
                        _ => {
                            let mut list: Vec<Term> = atom
                                .as_str()
                                .to_string()
                                .chars()
                                .map(|x| Term::Atom(x.to_string()))
                                .collect();

                            let mut partial_list = Term::Compound(
                                ".".into(),
                                vec![
                                    list.pop().unwrap(),
                                    tail,
                                ],
                            );

                            while let Some(last) = list.pop() {
                                partial_list = Term::Compound(
                                    ".".into(),
                                    vec![
                                        last,
                                        partial_list,
                                    ],
                                );
                            }

                            term_stack.push(partial_list);
                        }
                    }
                }
                // I dont know if this is needed here.
                /*
                (HeapCellValueTag::PStrLoc, h) => {
                    let atom = cell_as_atom_cell!(iter.heap[h]).get_name();
                    let tail = term_stack.pop().unwrap();

                    term_stack.push(Term::PartialString(
                        Cell::default(),
                        atom.as_str().to_owned(),
                        Box::new(tail),
                    ));
                }
                */
                _ => {
                }
            );
        }

        debug_assert_eq!(term_stack.len(), 1);
        term_stack.pop().unwrap()
    }
}

/// An iterator though the leaf answers of a query.
pub struct QueryState<'a> {
    machine: &'a mut Machine,
    term: TermWriteResult,
    stub_b: usize,
    var_names: IndexMap<HeapCellValue, VarPtr>,
    called: bool,
}

impl Drop for QueryState<'_> {
    fn drop(&mut self) {
        // This may be wrong if the iterator is not fully consumend, but from testing it seems
        // fine.
        self.machine.trust_me();
    }
}

impl Iterator for QueryState<'_> {
    type Item = Result<LeafAnswer, String>;

    fn next(&mut self) -> Option<Self::Item> {
        let var_names = &mut self.var_names;
        let term_write_result = &self.term;
        let machine = &mut self.machine;

        // No more choicepoints, end iteration
        if self.called && machine.machine_st.b <= self.stub_b {
            return None;
        }

        machine.dispatch_loop();

        self.called = true;

        if !machine.machine_st.ball.stub.is_empty() {
            // NOTE: this means an exception was thrown, at which
            // point we backtracked to the stub choice point.
            // this should halt the search for solutions as it
            // does in the Scryer top-level. the exception term is
            // contained in self.machine_st.ball.
            let error_string = self
                .machine
                .machine_st
                .ball
                .stub
                .iter()
                .filter(|h| {
                    matches!(
                        h.get_tag(),
                        HeapCellValueTag::Atom | HeapCellValueTag::Fixnum
                    )
                })
                .map(|h| match h.get_tag() {
                    HeapCellValueTag::Atom => {
                        let (name, _) = cell_as_atom_cell!(h).get_name_and_arity();
                        name.as_str().to_string()
                    }
                    HeapCellValueTag::Fixnum => h.get_value().clone().to_string(),
                    _ => unreachable!(),
                })
                .collect::<Vec<String>>()
                .join(" ");

            return Some(Err(error_string));
        }

        if machine.machine_st.p == LIB_QUERY_SUCCESS {
            if term_write_result.var_dict.is_empty() {
                self.machine.machine_st.backtrack();
                return Some(Ok(LeafAnswer::True));
            }
        } else if machine.machine_st.p == BREAK_FROM_DISPATCH_LOOP_LOC {
            return Some(Ok(LeafAnswer::False));
        }

        let mut bindings: BTreeMap<String, Term> = BTreeMap::new();

        let var_dict = &term_write_result.var_dict;

        for (var_key, term_to_be_printed) in var_dict.iter() {
            let mut var_name = var_key.to_string();
            if var_name.starts_with('_') {
                let should_print = var_names.values().any(|x| match x.borrow().clone() {
                    Var::Named(v) => v == var_name,
                    _ => false,
                });
                if !should_print {
                    continue;
                }
            }

            let mut term =
                Term::from_heapcell(machine, *term_to_be_printed, &mut var_names.clone());

            if let Term::Var(ref term_str) = term {
                if *term_str == var_name {
                    continue;
                }

                // Var dict is in the order things appear in the query. If var_name appears
                // after term in the query, switch their places.
                let var_name_idx = var_dict
                    .get_index_of(&VarKey::VarPtr(Var::Named(var_name.clone()).into()))
                    .unwrap();
                let term_idx =
                    var_dict.get_index_of(&VarKey::VarPtr(Var::Named(term_str.clone()).into()));
                if let Some(idx) = term_idx {
                    if idx < var_name_idx {
                        let new_term = Term::Var(var_name);
                        let new_var_name = term_str.into();
                        term = new_term;
                        var_name = new_var_name;
                    }
                }
            }

            bindings.insert(var_name, term);
        }

        // NOTE: there are outstanding choicepoints, backtrack
        // through them for further solutions. if
        // self.machine_st.b == stub_b we've backtracked to the stub
        // choice point, so we should break.
        self.machine.machine_st.backtrack();

        Some(Ok(LeafAnswer::LeafAnswer {
            bindings,
            residual_goals: vec![],
        }))
    }
}

impl Machine {
    /// Loads a module into the [`Machine`] from a string.
    pub fn load_module_string(&mut self, module_name: &str, program: impl Into<String>) {
        let stream = Stream::from_owned_string(program.into(), &mut self.machine_st.arena);
        self.load_file(module_name, stream);
    }

    /// Consults a module into the [`Machine`] from a string.
    pub fn consult_module_string(&mut self, module_name: &str, program: impl Into<String>) {
        let stream = Stream::from_owned_string(program.into(), &mut self.machine_st.arena);
        self.machine_st.registers[1] = stream_as_cell!(stream);
        self.machine_st.registers[2] = atom_as_cell!(&atom_table::AtomTable::build_with(
            &self.machine_st.atom_tbl,
            module_name
        ));

        self.run_module_predicate(atom!("loader"), (atom!("consult_stream"), 2));
    }

    fn allocate_stub_choice_point(&mut self) {
        // NOTE: create a choice point to terminate the dispatch_loop
        // if an exception is thrown.

        let stub_b = self.machine_st.stack.allocate_or_frame(0);
        let or_frame = self.machine_st.stack.index_or_frame_mut(stub_b);

        or_frame.prelude.num_cells = 0;
        or_frame.prelude.e = 0;
        or_frame.prelude.cp = 0;
        or_frame.prelude.b = 0;
        or_frame.prelude.bp = BREAK_FROM_DISPATCH_LOOP_LOC;
        or_frame.prelude.boip = 0;
        or_frame.prelude.biip = 0;
        or_frame.prelude.tr = 0;
        or_frame.prelude.h = 0;
        or_frame.prelude.b0 = 0;
        or_frame.prelude.attr_var_queue_len = 0;

        self.machine_st.b = stub_b;
        self.machine_st.hb = self.machine_st.heap.len();
        self.machine_st.block = stub_b;
    }

    /// Runs a query.
    pub fn run_query(&mut self, query: impl Into<String>) -> QueryState {
        let mut parser = Parser::new(
            Stream::from_owned_string(query.into(), &mut self.machine_st.arena),
            &mut self.machine_st,
        );
        let op_dir = CompositeOpDir::new(&self.indices.op_dir, None);
        let term = parser
            .read_term(&op_dir, Tokens::Default)
            .expect("Failed to parse query");

        self.allocate_stub_choice_point();

        // Write parsed term to heap
        let term_write_result =
            write_term_to_heap(&term, &mut self.machine_st.heap, &self.machine_st.atom_tbl)
                .expect("couldn't write term to heap");

        let var_names: IndexMap<_, _> = term_write_result
            .var_dict
            .iter()
            .map(|(var_key, cell)| match var_key {
                // NOTE: not the intention behind Var::InSitu here but
                // we can hijack it to store anonymous variables
                // without creating problems.
                VarKey::AnonVar(h) => (*cell, VarPtr::from(Var::InSitu(*h))),
                VarKey::VarPtr(var_ptr) => (*cell, var_ptr.clone()),
            })
            .collect();

        // Write term to heap
        self.machine_st.registers[1] = self.machine_st.heap[term_write_result.heap_loc];

        self.machine_st.cp = LIB_QUERY_SUCCESS; // BREAK_FROM_DISPATCH_LOOP_LOC;
        let call_index_p = self
            .indices
            .code_dir
            .get(&(atom!("call"), 1))
            .expect("couldn't get code index")
            .local()
            .unwrap();

        self.machine_st.execute_at_index(1, call_index_p);

        let stub_b = self.machine_st.b;
        QueryState {
            machine: self,
            term: term_write_result,
            stub_b,
            var_names,
            called: false,
        }
    }
}