mod literal;
mod printer;

pub use crate::generated::ir::*;
#[cfg(feature = "std")]
pub use printer::{print, print_no_color, print_ty};

use crate::{boxed, string, vec};

impl Program {
    pub fn get_output_ty(&self) -> &Ty {
        let main_ty = self.main.ty.kind.as_function().unwrap();
        &main_ty.body
    }

    pub fn get_input_ty(&self) -> &Ty {
        let main_ty = self.main.ty.kind.as_function().unwrap();
        assert_eq!(main_ty.params.len(), 1);
        &main_ty.params[0]
    }
}

impl Expr {
    pub fn new(kind: impl Into<ExprKind>, ty: Ty) -> Expr {
        Expr {
            kind: kind.into(),
            ty,
        }
    }
    pub fn new_lit_bool(value: bool) -> Self {
        Expr {
            kind: ExprKind::Literal(Literal::bool(value)),
            ty: Ty::new(TyPrimitive::bool),
        }
    }
}
impl From<ParameterPtr> for ExprKind {
    fn from(ptr: ParameterPtr) -> Self {
        ExprKind::Pointer(Pointer::Parameter(ptr))
    }
}
impl From<ExternalPtr> for ExprKind {
    fn from(ptr: ExternalPtr) -> Self {
        ExprKind::Pointer(Pointer::External(ptr))
    }
}
impl From<TupleLookup> for ExprKind {
    fn from(v: TupleLookup) -> Self {
        ExprKind::TupleLookup(boxed::Box::new(v))
    }
}
impl From<Binding> for ExprKind {
    fn from(v: Binding) -> Self {
        ExprKind::Binding(boxed::Box::new(v))
    }
}
impl From<Call> for ExprKind {
    fn from(v: Call) -> Self {
        ExprKind::Call(boxed::Box::new(v))
    }
}
impl From<Function> for ExprKind {
    fn from(v: Function) -> Self {
        ExprKind::Function(boxed::Box::new(v))
    }
}
impl From<EnumEq> for ExprKind {
    fn from(eq: EnumEq) -> Self {
        ExprKind::EnumEq(boxed::Box::new(eq))
    }
}

impl PartialEq for Ty {
    fn eq(&self, other: &Self) -> bool {
        self.kind == other.kind
    }
}

impl Eq for Ty {}

impl Ty {
    pub fn new(kind: impl Into<TyKind>) -> Self {
        Ty {
            kind: kind.into(),
            layout: None,
            name: None,
            variants_recursive: vec![],
        }
    }
    pub fn new_unit() -> Self {
        Ty {
            kind: TyKind::Tuple(vec![]),
            layout: Some(TyLayout {
                head_size: 0,
                body_ptrs: vec![],
            }),
            name: None,
            variants_recursive: vec![],
        }
    }
    pub fn is_unit(&self) -> bool {
        self.kind.is_unit()
    }
}

impl TyKind {
    pub fn is_unit(&self) -> bool {
        self.as_tuple().is_some_and(|f| f.is_empty())
    }

    pub fn as_option(&self) -> Option<&Ty> {
        self.as_enum()
            .filter(|v| v.len() == 2 && v[0].ty.is_unit() && !v[1].ty.is_unit())
            .map(|v| &v[1].ty)
    }
}

impl From<TyPrimitive> for TyKind {
    fn from(value: TyPrimitive) -> Self {
        TyKind::Primitive(value)
    }
}
impl From<vec::Vec<TyTupleField>> for TyKind {
    fn from(value: vec::Vec<TyTupleField>) -> Self {
        TyKind::Tuple(value)
    }
}
impl From<TyFunction> for TyKind {
    fn from(value: TyFunction) -> Self {
        TyKind::Function(boxed::Box::new(value))
    }
}
impl From<Path> for TyKind {
    fn from(value: Path) -> Self {
        TyKind::Ident(value)
    }
}

impl Module {
    pub fn insert(&mut self, path: &[string::String], decl: Decl) {
        if path.is_empty() {
            panic!();
        }

        if path.len() == 1 {
            self.decls.retain(|d| d.name != path[0]);
            self.decls.push(ModuledeclsItems {
                name: path[0].clone(),
                decl,
            });
        } else {
            let exists = self.decls.iter().any(|d| d.name == path[0]);
            if !exists {
                self.decls.push(ModuledeclsItems {
                    name: path[0].clone(),
                    decl: Decl::Module(boxed::Box::new(Module {
                        decls: vec::Vec::new(),
                    })),
                });
            }

            let sub_module = self.decls.iter_mut().find(|d| d.name == path[0]);
            let Decl::Module(sub_module) = &mut sub_module.unwrap().decl else {
                panic!()
            };
            sub_module.insert(&path[1..], decl)
        }
    }

    pub fn iter_defs_re(&self) -> impl Iterator<Item = (Path, &Decl)> {
        self.decls.iter().flat_map(|item| match &item.decl {
            Decl::Module(sub_module) => sub_module
                .iter_defs_re()
                .map(|(mut p, d)| {
                    p.0.insert(0, item.name.clone());
                    (p, d)
                })
                .collect::<vec::Vec<_>>(),
            _ => {
                vec![(Path(vec![item.name.clone()]), &item.decl)]
            }
        })
    }

    pub fn iter_types_re(&self) -> impl Iterator<Item = (Path, &Ty)> {
        self.iter_defs_re().filter_map(|(p, d)| {
            if let Decl::Type(ty) = d {
                Some((p, ty))
            } else {
                None
            }
        })
    }
}
