use std::{borrow::BorrowMut, cell::RefCell, rc::Rc};

use super::{
    ast::{Lambda, Tree},
    Context,
};

#[derive(Clone, Debug)]
pub enum Constraint {}

#[derive(Clone, Debug)]
pub enum CatCheckError {}

#[derive(Clone, Debug)]
pub enum TypeError {
    NeedGroundType(Rc<Type>),
    CatCheck(CatCheckError),
}

#[derive(Debug)]
pub struct AlgebraicType {}

#[derive(Debug)]
pub enum GroundType {
    Epsilon,
    Bottom,
    Character(char),
    Variable,
    Cat(RefCell<Rc<GroundType>>, RefCell<Rc<GroundType>>),
    Alt(RefCell<Rc<GroundType>>, RefCell<Rc<GroundType>>),
    Fix(RefCell<Rc<GroundType>>, RefCell<Rc<GroundType>>),
}

#[derive(Clone, Debug)]
pub struct UnificationError(pub Rc<Type>, pub Rc<Type>);

impl GroundType {
    pub fn is_variable(&self) -> bool {
        matches!(self, Self::Variable)
    }

    pub fn substitute(self: &mut Rc<Self>, from: &Rc<Self>, into: &Rc<Self>) {
        if Rc::ptr_eq(self, from) {
            *self = into.clone();
            return;
        }

        match self.as_ref() {
            GroundType::Epsilon => {}
            GroundType::Bottom => {}
            GroundType::Character(_) => {}
            GroundType::Variable => {}
            GroundType::Cat(front, back) => {
                front.borrow_mut().substitute(from, into);
                back.borrow_mut().substitute(from, into);
            }
            GroundType::Alt(left, right) => {
                left.borrow_mut().substitute(from, into);
                right.borrow_mut().substitute(from, into);
            }
            GroundType::Fix(_, inner) => {
                inner.borrow_mut().substitute(from, into);
            }
        }
    }

    pub fn occurs_in(self: &Rc<Self>, other: &Rc<Self>) -> bool {
        if Rc::ptr_eq(self, other) {
            return true;
        }

        match other.as_ref() {
            GroundType::Epsilon => false,
            GroundType::Bottom => false,
            GroundType::Character(_) => false,
            GroundType::Variable => false,
            GroundType::Cat(front, back) => {
                self.occurs_in(&front.borrow()) || self.occurs_in(&back.borrow())
            }
            GroundType::Alt(left, right) => {
                self.occurs_in(&left.borrow()) || self.occurs_in(&right.borrow())
            }
            GroundType::Fix(arg, out) => {
                self.occurs_in(&arg.borrow()) || self.occurs_in(&out.borrow())
            }
        }
    }

    pub fn unify<'a>(
        self: &'a mut Rc<Self>,
        other: &'a mut Rc<Self>,
    ) -> Result<(), UnificationError> {
        if Rc::ptr_eq(self, other) {
            return Ok(());
        }

        match (self.as_ref(), other.as_ref()) {
            (GroundType::Epsilon, GroundType::Epsilon) => Ok(()),
            (GroundType::Bottom, GroundType::Bottom) => Ok(()),
            (GroundType::Character(c), GroundType::Character(d)) if c == d => Ok(()),
            (GroundType::Cat(front1, back1), GroundType::Cat(front2, back2)) => {
                front1.borrow_mut().unify(&mut front2.borrow_mut())?;
                back1.borrow_mut().unify(&mut back2.borrow_mut())
            }
            (GroundType::Alt(left1, right1), GroundType::Alt(left2, right2)) => {
                left1.borrow_mut().unify(&mut left2.borrow_mut())?;
                right1.borrow_mut().unify(&mut right2.borrow_mut())
            }
            (GroundType::Fix(arg1, out1), GroundType::Fix(arg2, out2)) => {
                let new_arg = Rc::new(GroundType::Variable);
                out1.borrow_mut().substitute(&arg1.borrow(), &new_arg);
                out2.borrow_mut().substitute(&arg2.borrow(), &new_arg);
                arg1.replace(new_arg.clone());
                arg2.replace(new_arg.clone());
                out1.borrow_mut().unify(&mut out2.borrow_mut())
            }
            (GroundType::Variable, GroundType::Variable) => {
                self = other;
                Ok(())
            }
            (GroundType::Variable, _) if !self.occurs_in(other) => {
                self = other;
                Ok(())
            }
            (_, GroundType::Variable) if !other.occurs_in(self) => {
                other = self;
                Ok(())
            }
            _ => Err(UnificationError(
                Rc::new(self.clone().into()),
                Rc::new(other.clone().into()),
            )),
        }
    }
}

/// Function type where the formal parameter is guarded in the body.
/// The actual parameter is checked in an unguarded context.
#[derive(Debug)]
pub struct GuardedFunction(pub RefCell<Rc<GroundType>>, pub RefCell<Rc<Type>>);

#[derive(Debug)]
pub enum Type {
    Ground(RefCell<Rc<GroundType>>),
    /// Function type where the formal parameter is unguarded in the body.
    /// The actual parameter is checked in a guarded context.
    UnguardedFunction(RefCell<Rc<GroundType>>, RefCell<Rc<Type>>),
    GuardedFunction(GuardedFunction),
    Variable,
}

impl Type {
    pub fn try_as_ground<'a>(
        self: &'a Rc<Self>,
    ) -> Result<&'a RefCell<Rc<GroundType>>, &'a Rc<Self>> {
        if let Type::Ground(g) = self.as_ref() {
            Ok(g)
        } else {
            Err(self)
        }
    }
    pub fn unify<'a>(
        self: &'a mut Rc<Self>,
        other: &'a mut Rc<Self>,
    ) -> Result<(), UnificationError> {
        if Rc::ptr_eq(self, other) {
            return Ok(());
        }

        match self.as_ref() {
            Type::Ground(g) => {
                if let Type::Ground(o) = other.as_ref() {
                    g.borrow_mut().unify(&mut o.borrow_mut())
                } else {
                    Err(UnificationError(self.clone(), other.clone()))
                }
            }
            Type::UnguardedFunction(arg, out) => {
                if let Type::UnguardedFunction(o_arg, o_out) = other.as_ref() {
                    arg.borrow_mut().unify(&mut o_arg.borrow_mut())?;
                    out.borrow_mut().unify(&mut o_out.borrow_mut())
                } else {
                    Err(UnificationError(self.clone(), other.clone()))
                }
            }
            Type::GuardedFunction(GuardedFunction(arg, out)) => {
                if let Type::GuardedFunction(GuardedFunction(o_arg, o_out)) = other.as_ref() {
                    arg.borrow_mut().unify(&mut o_arg.borrow_mut())?;
                    out.borrow_mut().unify(&mut o_out.borrow_mut())
                } else {
                    Err(UnificationError(self.clone(), other.clone()))
                }
            }
            Type::Variable => {
                self = other;
                Ok(())
            }
        }
    }
}

impl From<GroundType> for Type {
    fn from(ty: GroundType) -> Self {
        Rc::new(ty).into()
    }
}

impl From<Rc<GroundType>> for Type {
    fn from(ty: Rc<GroundType>) -> Self {
        Self::Ground(RefCell::new(ty))
    }
}

pub fn type_of(ctx: &mut Context, tree: &Rc<Tree>) -> Result<Rc<Type>, TypeError> {
    match tree.as_ref() {
        Tree::Epsilon => Ok(ctx.type_interner.epsilon()),
        Tree::Bottom => Ok(ctx.type_interner.bottom()),
        Tree::Identity => {
            let arg = Rc::new(GroundType::Variable);
            Ok(Rc::new(Type::UnguardedFunction(
                RefCell::new(arg),
                RefCell::new(Rc::new(arg.into())),
            )))
        }
        Tree::Literal(l) => {
            if let Some(c) = l.chars().next() {
                Ok(ctx.type_interner.character(c))
            } else {
                Ok(ctx.type_interner.epsilon())
            }
        }
        Tree::Variable => todo!(), // ctx.lookup_type_variable(tree),
        Tree::Cat(front, back) => {
            let mut front_ty: Rc<Type> = Rc::new(GroundType::Variable.into());
            let mut back_ty: Rc<Type> = Rc::new(GroundType::Variable.into());
            ctx.type_of(front)?
                .unify(&mut front_ty)
                .map_err(|_| todo!())?;
            ctx.type_of(back)?
                .unify(&mut back_ty)
                .map_err(|_| todo!())?;
            let front_ty = front_ty.try_as_ground().unwrap();
            let back_ty = back_ty.try_as_ground().unwrap();
            Ok(ctx
                .type_interner
                .cat(front_ty.borrow().clone(), back_ty.borrow().clone()))
        }
        Tree::Alt(left, right) => {
            let mut left_ty: Rc<Type> = Rc::new(GroundType::Variable.into());
            let mut right_ty: Rc<Type> = Rc::new(GroundType::Variable.into());
            ctx.type_of(left)?
                .unify(&mut left_ty)
                .map_err(|_| todo!())?;
            ctx.type_of(right)?
                .unify(&mut right_ty)
                .map_err(|_| todo!())?;
            let left_ty = left_ty.try_as_ground().unwrap();
            let right_ty = right_ty.try_as_ground().unwrap();
            Ok(ctx
               .type_interner
               .alt(left_ty.borrow().clone(), right_ty.borrow().clone()))
        }
        Tree::Fix(inner) => {
            let shape = Rc::new(Type::GuardedFunction(GuardedFunction(RefCell::new(Rc::new(GroundType::Variable)), RefCell::new(Rc::new(GroundType::Variable.into())))));
            ctx.type_of(inner)?
                .unify(&mut shape)
                .map_err(|_| todo!())?;
            todo!()
        }
        Tree::Call(_, _) => todo!(),
        Tree::Lambda(_) => todo!(),
        Tree::Let(_, _, _) => todo!(),
    }
}