use std::{collections::HashMap, rc::Rc};

use intern::Intern;

use self::ast::Lambda;

mod ast;
mod cst;
mod intern;
mod typed;

type SubstTriple = (Intern<ast::Tree>, Intern<ast::Tree>, Intern<ast::Tree>);

#[derive(Debug)]
pub struct Context {
    // Abstract tree interning
    var_stack: Vec<Rc<ast::Tree>>,
    pub tree_interner: intern::TreeInterner,
    pub type_interner: intern::TypeInterner,

    // Caches
    as_tree: HashMap<Intern<cst::Shape>, Rc<ast::Tree>>,
    substitute: HashMap<SubstTriple, Option<Rc<ast::Tree>>>,
    translate: HashMap<Intern<ast::Tree>, Result<Option<Rc<ast::Tree>>, ast::TranslateError>>,
    type_of: HashMap<
        Intern<ast::Tree>,
        Result<typed::ConstrainedType, typed::TypeError>,
    >,
}

impl Context {
    /// Converts a concrete syntax tree into an abstract one.
    pub fn as_tree(&mut self, shape: &Rc<cst::Shape>) -> Rc<ast::Tree> {
        if let Some(tree) = self.as_tree.get(&shape.clone().into()) {
            return tree.clone();
        }
        let tree = cst::as_tree(self, shape);
        self.as_tree.insert(shape.clone().into(), tree.clone());
        tree
    }

    /// Creates a copy of `tree` that replaces variables pointing to `from` with `into`.
    pub fn substitute(
        &mut self,
        tree: &Rc<ast::Tree>,
        from: &Rc<ast::Tree>,
        into: &Rc<ast::Tree>,
    ) -> Option<Rc<ast::Tree>> {
        if let Some(tree) = self.substitute.get(&(
            tree.clone().into(),
            from.clone().into(),
            into.clone().into(),
        )) {
            return tree.clone();
        }
        let out = ast::substitute(self, tree, from, into);
        self.substitute.insert(
            (
                tree.clone().into(),
                from.clone().into(),
                into.clone().into(),
            ),
            out.clone(),
        );
        out
    }

    /// Expands let bindings and function calls until the top level construct is a base term.
    pub fn translate(
        &mut self,
        tree: &Rc<ast::Tree>,
    ) -> Result<Option<Rc<ast::Tree>>, ast::TranslateError> {
        if let Some(tree) = self.translate.get(&tree.clone().into()) {
            return tree.clone();
        }
        let out = ast::translate(self, tree);
        self.translate.insert(tree.clone().into(), out.clone());
        out
    }

    /// Computes set of unsolvable type constraints for a term. Returns an error
    /// if a solvable constraint is unsatisfiable, or the constraints cannot be inferred.
    pub fn type_of(
        &mut self,
        tree: &Rc<ast::Tree>,
    ) -> Result<typed::ConstrainedType, typed::TypeError> {
        if let Some(ty) = self.type_of.get(&tree.clone().into()) {
            return ty.clone();
        }
        let ty = typed::type_of(self, tree);
        self.type_of.insert(tree.clone().into(), ty.clone());
        ty
    }

    fn lookup_tree_variable(&self, idx: usize) -> Rc<ast::Tree> {
        self.var_stack[idx].clone()
    }

    fn new_lambda<F: FnOnce(&mut Self) -> Rc<ast::Tree>>(&mut self, f: F) -> Rc<ast::Tree> {
        // Create variable node
        let arg = Rc::new(ast::Tree::Variable);

        // Apply inner with the variable
        self.var_stack.push(arg.clone());
        let inner = f(self);
        self.var_stack.pop();

        // Create the lambda
        Rc::new(ast::Tree::Lambda(Lambda(arg, inner)))
    }

    fn new_let<F: FnOnce(&mut Self) -> Rc<ast::Tree>>(
        &mut self,
        bind: Rc<ast::Tree>,
        inner: F,
    ) -> Rc<ast::Tree> {
        // Create variable node
        let arg = Rc::new(ast::Tree::Variable);

        // Apply inner with the variable
        self.var_stack.push(arg.clone());
        let inner = inner(self);
        self.var_stack.pop();

        // Create the let
        Rc::new(ast::Tree::Let(bind, arg, inner))
    }
}