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

use crate::{
    parse::Constant,
    types::{Ident, PrimitiveType, TaggedType, Type, TypeTag},
};

use super::DeBrujinAst;

#[derive(Clone)]
pub enum DeBrujinBuiltInAst {
    Abstraction(Ident, Type, Box<DeBrujinBuiltInAst>), // \:1.2
    Application(Box<DeBrujinBuiltInAst>, Box<DeBrujinBuiltInAst>), // 0 1
    FreeVariable(String),                              // x
    BoundVariable(usize),                              // 1
    Constant(Constant),                                // true | false | n
    Builtin(Rc<dyn Builtin>),
}

impl DeBrujinAst {
    pub fn resolve_builtins(
        self,
        builtins: &HashMap<String, Rc<dyn Builtin>>,
    ) -> DeBrujinBuiltInAst {
        match self {
            DeBrujinAst::Abstraction(i, t, ast) => {
                DeBrujinBuiltInAst::Abstraction(i, t, Box::new(ast.resolve_builtins(builtins)))
            }
            DeBrujinAst::Application(lhs, rhs) => DeBrujinBuiltInAst::Application(
                Box::new(lhs.resolve_builtins(builtins)),
                Box::new(rhs.resolve_builtins(builtins)),
            ),
            DeBrujinAst::FreeVariable(x) => {
                if let Some(b) = builtins.get(&x) {
                    DeBrujinBuiltInAst::Builtin(b.clone())
                } else {
                    DeBrujinBuiltInAst::FreeVariable(x)
                }
            }
            DeBrujinAst::BoundVariable(b) => DeBrujinBuiltInAst::BoundVariable(b),
            DeBrujinAst::Constant(c) => DeBrujinBuiltInAst::Constant(c),
        }
    }
}

pub trait Builtin {
    fn name(&self) -> String;

    fn to_ast(&self) -> DeBrujinAst {
        DeBrujinAst::FreeVariable(self.name())
    }

    fn r#type(&self) -> TaggedType;

    fn apply(&self, rhs: DeBrujinBuiltInAst) -> Option<DeBrujinBuiltInAst>;
}

impl DeBrujinAst {
    pub fn reduce_builtins(self, builtins: &HashMap<String, Rc<dyn Builtin>>) -> DeBrujinAst {
        self.resolve_builtins(builtins).reduce_builtins().into()
    }
}

impl DeBrujinBuiltInAst {
    fn reduce_builtins(self) -> DeBrujinBuiltInAst {
        match self {
            DeBrujinBuiltInAst::Abstraction(i, t, ast) => {
                DeBrujinBuiltInAst::Abstraction(i, t, Box::new(ast.reduce_builtins()))
            }
            DeBrujinBuiltInAst::Application(lhs, rhs) => match *lhs {
                DeBrujinBuiltInAst::Builtin(builtin) => builtin
                    .apply(*rhs)
                    .expect("the type checker should make sure we can apply")
                    .reduce_builtins(),
                lhs => DeBrujinBuiltInAst::Application(
                    Box::new(lhs.reduce_builtins()),
                    Box::new(rhs.reduce_builtins()),
                )
                .reduce_builtins(),
            },
            a => a,
        }
    }
}

impl Into<DeBrujinAst> for DeBrujinBuiltInAst {
    fn into(self) -> DeBrujinAst {
        match self {
            DeBrujinBuiltInAst::Abstraction(i, t, ast) => {
                DeBrujinAst::Abstraction(i, t, Box::new((*ast).into()))
            }
            DeBrujinBuiltInAst::Application(lhs, rhs) => {
                DeBrujinAst::Application(Box::new((*lhs).into()), Box::new((*rhs).into()))
            }
            DeBrujinBuiltInAst::FreeVariable(x) => DeBrujinAst::FreeVariable(x),
            DeBrujinBuiltInAst::BoundVariable(i) => DeBrujinAst::BoundVariable(i),
            DeBrujinBuiltInAst::Constant(constant) => DeBrujinAst::Constant(constant),
            DeBrujinBuiltInAst::Builtin(builtin) => DeBrujinAst::FreeVariable(builtin.name()),
        }
    }
}