use std::rc::Rc;

use super::Context;

#[derive(Debug)]
pub struct Lambda(pub Rc<Tree>, pub Rc<Tree>);

#[derive(Debug)]
pub enum Tree {
    Epsilon,
    Bottom,
    Identity,
    Literal(String),
    Variable,
    Cat(Rc<Tree>, Rc<Tree>),
    Alt(Rc<Tree>, Rc<Tree>),
    Fix(Rc<Tree>),
    Call(Rc<Tree>, Rc<Tree>),
    Lambda(Lambda),
    Let(Rc<Tree>, Rc<Tree>, Rc<Tree>),
}

impl Tree {
    pub fn try_as_lambda(&self) -> Option<&Lambda> {
        if let Self::Lambda(v) = self {
            Some(v)
        } else {
            None
        }
    }
}

#[derive(Clone, Debug)]
pub enum TranslateError {
    CallFunNotLambda(Rc<Tree>),
}

pub(crate) fn substitute(
    ctx: &mut Context,
    tree: &Rc<Tree>,
    from: &Rc<Tree>,
    into: &Rc<Tree>,
) -> Option<Rc<Tree>> {
    if Rc::ptr_eq(tree, from) {
        return Some(into.clone());
    }

    match tree.as_ref() {
        Tree::Epsilon | Tree::Bottom | Tree::Identity | Tree::Literal(_) | Tree::Variable => None,
        Tree::Cat(front, back) => {
            let new_front = ctx.substitute(front, from, into);
            let new_back = ctx.substitute(back, from, into);
            let (front, back) = match (new_front, new_back) {
                (None, None) => return None,
                (None, Some(back)) => (front.clone(), back),
                (Some(front), None) => (front, back.clone()),
                (Some(front), Some(back)) => (front, back),
            };
            Some(ctx.tree_interner.cat(front, back))
        }
        Tree::Alt(left, right) => {
            let new_left = ctx.substitute(left, from, into);
            let new_right = ctx.substitute(right, from, into);
            let (left, right) = match (new_left, new_right) {
                (None, None) => return None,
                (None, Some(right)) => (left.clone(), right),
                (Some(left), None) => (left, right.clone()),
                (Some(left), Some(right)) => (left, right),
            };
            Some(ctx.tree_interner.alt(left, right))
        }
        Tree::Fix(inner) => {
            let inner = ctx.substitute(inner, from, into)?;
            Some(ctx.tree_interner.fix(inner))
        }
        Tree::Call(fun, arg) => {
            let new_fun = ctx.substitute(fun, from, into);
            let new_arg = ctx.substitute(arg, from, into);
            let (fun, arg) = match (new_fun, new_arg) {
                (None, None) => return None,
                (None, Some(arg)) => (fun.clone(), arg),
                (Some(fun), None) => (fun, arg.clone()),
                (Some(fun), Some(arg)) => (fun, arg),
            };
            Some(ctx.tree_interner.call(fun, arg))
        }
        Tree::Lambda(Lambda(arg, inner)) => {
            let inner = ctx.substitute(inner, from, into)?;
            Some(Rc::new(Tree::Lambda(Lambda(arg.clone(), inner))))
        }
        Tree::Let(bind, arg, inner) => {
            let new_bind = ctx.substitute(bind, from, into);
            let new_inner = ctx.substitute(inner, from, into);
            let (bind, inner) = match (new_bind, new_inner) {
                (None, None) => return None,
                (None, Some(inner)) => (bind.clone(), inner),
                (Some(bind), None) => (bind, inner.clone()),
                (Some(bind), Some(inner)) => (bind, inner),
            };
            Some(Rc::new(Tree::Let(bind, arg.clone(), inner)))
        }
    }
}

pub(crate) fn translate(
    ctx: &mut Context,
    tree: &Rc<Tree>,
) -> Result<Option<Rc<Tree>>, TranslateError> {
    match tree.as_ref() {
        Tree::Epsilon
        | Tree::Bottom
        | Tree::Identity
        | Tree::Literal(_)
        | Tree::Variable
        | Tree::Cat(_, _)
        | Tree::Alt(_, _)
        | Tree::Fix(_)
        | Tree::Lambda(_) => Ok(None),
        Tree::Call(fun, arg) => {
            let fun = ctx.translate(fun)?.unwrap_or_else(|| fun.clone());
            let lambda = fun
                .try_as_lambda()
                .ok_or_else(|| TranslateError::CallFunNotLambda(fun.clone()))?;
            let mut out = ctx
                .substitute(&lambda.1, &lambda.0, arg)
                .unwrap_or_else(|| lambda.1.clone());
            loop {
                match ctx.translate(&out)? {
                    Some(tree) => out = tree,
                    None => return Ok(Some(out)),
                }
            }
        }
        Tree::Let(bind, arg, inner) => {
            let mut out = ctx
                .substitute(inner, arg, bind)
                .unwrap_or_else(|| inner.clone());
            loop {
                match ctx.translate(&out)? {
                    Some(tree) => out = tree,
                    None => return Ok(Some(out)),
                }
            }
        }
    }
}