use std::{cell::RefCell, collections::HashMap, hash::Hash, rc::Rc};

use super::{
    ast::Tree,
    typed::{GroundType, Type},
};

#[derive(Debug)]
pub struct Intern<T>(Rc<T>);

impl<T> PartialEq for Intern<T> {
    fn eq(&self, other: &Self) -> bool {
        Rc::ptr_eq(&self.0, &other.0)
    }
}

impl<T> Eq for Intern<T> {}

impl<T> Hash for Intern<T> {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        Rc::as_ptr(&self.0).hash(state)
    }
}

impl<T> From<Rc<T>> for Intern<T> {
    fn from(inner: Rc<T>) -> Self {
        Self(inner)
    }
}

#[derive(Debug)]
pub struct TreeInterner {
    epsilon: Option<Rc<Tree>>,
    bottom: Option<Rc<Tree>>,
    identity: Option<Rc<Tree>>,
    literal: HashMap<String, Rc<Tree>>,
    cat: HashMap<(Intern<Tree>, Intern<Tree>), Rc<Tree>>,
    alt: HashMap<(Intern<Tree>, Intern<Tree>), Rc<Tree>>,
    fix: HashMap<Intern<Tree>, Rc<Tree>>,
    call: HashMap<(Intern<Tree>, Intern<Tree>), Rc<Tree>>,
}

impl TreeInterner {
    pub fn epsilon(&mut self) -> Rc<Tree> {
        self.epsilon
            .get_or_insert_with(|| Rc::new(Tree::Epsilon))
            .clone()
    }

    pub fn bottom(&mut self) -> Rc<Tree> {
        self.bottom
            .get_or_insert_with(|| Rc::new(Tree::Bottom))
            .clone()
    }

    pub fn identity(&mut self) -> Rc<Tree> {
        self.identity
            .get_or_insert_with(|| Rc::new(Tree::Identity))
            .clone()
    }

    pub fn literal(&mut self, lit: String) -> Rc<Tree> {
        self.literal
            .entry(lit)
            .or_insert_with_key(|lit| Rc::new(Tree::Literal(lit.clone())))
            .clone()
    }

    pub fn cat(&mut self, front: Rc<Tree>, back: Rc<Tree>) -> Rc<Tree> {
        self.cat
            .entry((Intern(front), Intern(back)))
            .or_insert_with_key(|(front, back)| Rc::new(Tree::Cat(front.0.clone(), back.0.clone())))
            .clone()
    }

    pub fn alt(&mut self, left: Rc<Tree>, right: Rc<Tree>) -> Rc<Tree> {
        self.alt
            .entry((Intern(left), Intern(right)))
            .or_insert_with_key(|(left, right)| Rc::new(Tree::Alt(left.0.clone(), right.0.clone())))
            .clone()
    }

    pub fn fix(&mut self, inner: Rc<Tree>) -> Rc<Tree> {
        self.fix
            .entry(Intern(inner))
            .or_insert_with_key(|inner| Rc::new(Tree::Fix(inner.0.clone())))
            .clone()
    }

    pub fn call(&mut self, fun: Rc<Tree>, arg: Rc<Tree>) -> Rc<Tree> {
        self.call
            .entry((Intern(fun), Intern(arg)))
            .or_insert_with_key(|(fun, arg)| Rc::new(Tree::Call(fun.0.clone(), arg.0.clone())))
            .clone()
    }
}

#[derive(Debug)]
pub struct TypeInterner {
    epsilon: Option<Rc<Type>>,
    bottom: Option<Rc<Type>>,
    character: HashMap<char, Rc<Type>>,
    cat: HashMap<(Intern<GroundType>, Intern<GroundType>), Rc<Type>>,
    alt: HashMap<(Intern<GroundType>, Intern<GroundType>), Rc<Type>>,
    fix: HashMap<Intern<GroundType>, Rc<Type>>,
}

impl TypeInterner {
    pub fn epsilon(&mut self) -> Rc<Type> {
        self.epsilon
            .get_or_insert_with(|| Rc::new(GroundType::Epsilon.into()))
            .clone()
    }

    pub fn bottom(&mut self) -> Rc<Type> {
        self.bottom
            .get_or_insert_with(|| Rc::new(GroundType::Bottom.into()))
            .clone()
    }

    pub fn character(&mut self, c: char) -> Rc<Type> {
        self.character
            .entry(c)
            .or_insert_with(|| Rc::new(GroundType::Character(c).into()))
            .clone()
    }

    pub fn cat(&mut self, front: Rc<GroundType>, back: Rc<GroundType>) -> Rc<Type> {
        self.cat
            .entry((Intern(front), Intern(back)))
            .or_insert_with_key(|(front, back)| {
                Rc::new(
                    GroundType::Cat(RefCell::new(front.0.clone()), RefCell::new(back.0.clone()))
                        .into(),
                )
            })
            .clone()
    }

    pub fn alt(&mut self, left: Rc<GroundType>, right: Rc<GroundType>) -> Rc<Type> {
        self.alt
            .entry((Intern(left), Intern(right)))
            .or_insert_with_key(|(left, right)| {
                Rc::new(
                    GroundType::Alt(RefCell::new(left.0.clone()), RefCell::new(right.0.clone()))
                        .into(),
                )
            })
            .clone()
    }
}