use std::{fmt::Display, str};

pub type Ident = String;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PrimitiveType {
    Nat,
    Bool,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TypeTag {
    Num,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TaggedType {
    Tagged(TypeTag, Ident, Box<TaggedType>),
    Concrete(Type),
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Type {
    Generic(Ident),              // a
    Primitive(PrimitiveType),    // Bool | Nat
    Arrow(Box<Type>, Box<Type>), // 0 -> 1
}

impl From<Type> for TaggedType {
    fn from(value: Type) -> Self {
        Self::Concrete(value)
    }
}

impl TaggedType {
    pub fn r#type(&self) -> &Type {
        match self {
            TaggedType::Tagged(_, _, tagged_type) => tagged_type.r#type(),
            TaggedType::Concrete(t) => t,
        }
    }

    pub fn type_mut(&mut self) -> &mut Type {
        match self {
            TaggedType::Tagged(_, _, tagged_type) => tagged_type.type_mut(),
            TaggedType::Concrete(t) => t,
        }
    }

    pub fn to_concrete(self) -> Option<Type> {
        match self {
            TaggedType::Tagged(type_tag, _, tagged_type) => None,
            TaggedType::Concrete(t) => {
                if t.is_concrete() {
                    Some(t)
                } else {
                    None
                }
            }
        }
    }

    pub fn matches(&self, typ: &Type) -> (bool, Option<Ident>) {
        match self {
            TaggedType::Tagged(type_tag, ident, tagged_type) => {
                if typ.has_tag(type_tag) {
                    (true, Some(ident.clone()))
                } else {
                    tagged_type.matches(typ)
                }
            }
            TaggedType::Concrete(c) => (c == typ, None),
        }
    }

    pub fn arrow(self) -> Option<(TaggedType, TaggedType)> {
        match self {
            TaggedType::Tagged(type_tag, ident, tagged_type) => {
                let (lhs, rhs) = tagged_type.arrow()?;
                Some((
                    TaggedType::Tagged(type_tag.clone(), ident.clone(), Box::new(lhs))
                        .clear_unused_names(),
                    TaggedType::Tagged(type_tag, ident, Box::new(rhs)).clear_unused_names(),
                ))
            }
            TaggedType::Concrete(Type::Arrow(lhs, rhs)) => {
                Some((TaggedType::Concrete(*lhs), TaggedType::Concrete(*rhs)))
            }
            TaggedType::Concrete(_) => None,
        }
    }

    pub fn specialize(self, ident: &str, typ: &Type) -> TaggedType {
        match self {
            TaggedType::Tagged(_, i, tagged_type) if i == ident => {
                tagged_type.specialize(ident, typ)
            }
            TaggedType::Tagged(t, i, tagged_type) => {
                TaggedType::Tagged(t, i, Box::new(tagged_type.specialize(ident, typ)))
            }
            TaggedType::Concrete(c) => TaggedType::Concrete(c.specialize(ident, typ)),
        }
    }

    fn name_used(&self, ident: &str) -> bool {
        match self {
            TaggedType::Tagged(_, _, tagged_type) => tagged_type.name_used(ident),
            TaggedType::Concrete(c) => c.name_used(ident),
        }
    }

    fn clear_unused_names(self) -> TaggedType {
        match self {
            TaggedType::Tagged(type_tag, ident, tagged_type) => {
                if tagged_type.name_used(&ident) {
                    TaggedType::Tagged(type_tag, ident, tagged_type)
                } else {
                    *tagged_type
                }
            }
            t => t,
        }
    }
}

impl From<PrimitiveType> for Type {
    fn from(value: PrimitiveType) -> Self {
        Self::Primitive(value)
    }
}

impl Type {
    pub fn arrow<T1: Into<Type>, T2: Into<Type>>(t1: T1, t2: T2) -> Self {
        Self::Arrow(Box::new(t1.into()), Box::new(t2.into()))
    }

    pub fn is_concrete(&self) -> bool {
        match self {
            Type::Generic(_) => false,
            Type::Primitive(primitive_type) => true,
            Type::Arrow(t1, t2) => t1.is_concrete() && t2.is_concrete(),
        }
    }

    pub fn has_tag(&self, tag: &TypeTag) -> bool {
        match self {
            Type::Generic(_) => false,
            Type::Primitive(primitive_type) => match (primitive_type, tag) {
                (PrimitiveType::Nat, TypeTag::Num) => true,
                _ => false,
            },
            Type::Arrow(_, _) => false,
        }
    }

    fn name_used(&self, ident: &str) -> bool {
        match self {
            Type::Generic(i) if *i == ident => true,
            _ => false,
        }
    }

    fn specialize(self, ident: &str, typ: &Type) -> Type {
        match self {
            Type::Generic(i) if i == ident => typ.clone(),
            Type::Arrow(lhs, rhs) => Type::Arrow(
                Box::new(lhs.specialize(ident, typ)),
                Box::new(rhs.specialize(ident, typ)),
            ),
            t => t,
        }
    }
}

impl Display for TypeTag {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TypeTag::Num => write!(f, "Num"),
        }
    }
}

impl Display for TaggedType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TaggedType::Tagged(type_tag, ident, tagged_type) => {
                write!(f, "{type_tag} {ident} => {tagged_type}")
            }
            TaggedType::Concrete(t) => write!(f, "{t}"),
        }
    }
}

impl Display for Type {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Type::Generic(i) => write!(f, "{i}"),
            Type::Primitive(primitive_type) => write!(f, "{primitive_type}"),
            Type::Arrow(t1, t2) => write!(f, "({t1} -> {t2})"),
        }
    }
}

impl Display for PrimitiveType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            PrimitiveType::Nat => write!(f, "Nat"),
            PrimitiveType::Bool => write!(f, "Bool"),
        }
    }
}