use rustc_hash::FxHashMap;
use swc_core::atoms::{atom, Atom, Wtf8Atom};
use swc_core::common::comments::Comments;
use swc_core::common::Span;
use swc_core::ecma::ast::*;
use swc_core::ecma::visit::{noop_visit_mut_type, visit_mut_pass, VisitMut, VisitMutWith};

#[cfg(test)]
mod tests;

/// This pass adds a /*#__PURE__#/ annotation to calls to known pure top-level
/// Inferno methods, so that terser and other minifiers can safely remove them
/// during dead code elimination.
pub fn pure_annotations<C>(comments: Option<C>) -> impl Pass
where
    C: Comments,
{
    visit_mut_pass(PureAnnotations {
        imports: Default::default(),
        comments,
    })
}

struct PureAnnotations<C>
where
    C: Comments,
{
    imports: FxHashMap<Id, (Wtf8Atom, Atom)>,
    comments: Option<C>,
}

impl<C> VisitMut for PureAnnotations<C>
where
    C: Comments,
{
    noop_visit_mut_type!();

    fn visit_mut_module(&mut self, module: &mut Module) {
        // Pass 1: collect imports
        for item in &module.body {
            if let ModuleItem::ModuleDecl(ModuleDecl::Import(import)) = item {
                let Some(src_str) = import.src.value.as_str() else {
                    continue;
                };
                if src_str != "inferno" {
                    continue;
                }

                for specifier in &import.specifiers {
                    let src = import.src.value.clone();
                    match specifier {
                        ImportSpecifier::Named(named) => {
                            let imported: Atom = match &named.imported {
                                Some(ModuleExportName::Ident(imported)) => imported.sym.clone(),
                                Some(ModuleExportName::Str(s)) => {
                                    s.value.to_atom_lossy().into_owned()
                                }
                                None => named.local.sym.clone(),
                                #[cfg(swc_ast_unknown)]
                                Some(_) => continue,
                            };
                            self.imports.insert(named.local.to_id(), (src, imported));
                        }
                        ImportSpecifier::Default(default) => {
                            self.imports
                                .insert(default.local.to_id(), (src, atom!("default")));
                        }
                        ImportSpecifier::Namespace(ns) => {
                            self.imports.insert(ns.local.to_id(), (src, atom!("*")));
                        }
                        #[cfg(swc_ast_unknown)]
                        _ => (),
                    }
                }
            }
        }

        if self.imports.is_empty() {
            return;
        }

        // Pass 2: add pure annotations.
        module.visit_mut_children_with(self);
    }

    fn visit_mut_call_expr(&mut self, call: &mut CallExpr) {
        let is_inferno_call = match &call.callee {
            Callee::Expr(expr) => match &**expr {
                Expr::Ident(ident) => {
                    if let Some((src, specifier)) = self.imports.get(&ident.to_id()) {
                        is_pure(src, specifier)
                    } else {
                        false
                    }
                }
                Expr::Member(member) => match &*member.obj {
                    Expr::Ident(ident) => {
                        if let Some((src, specifier)) = self.imports.get(&ident.to_id()) {
                            if &**specifier == "default" || &**specifier == "*" {
                                match &member.prop {
                                    MemberProp::Ident(ident) => is_pure(src, &ident.sym),
                                    _ => false,
                                }
                            } else {
                                false
                            }
                        } else {
                            false
                        }
                    }
                    _ => false,
                },
                _ => false,
            },
            _ => false,
        };

        if is_inferno_call {
            if let Some(comments) = &self.comments {
                if call.span.lo.is_dummy() {
                    call.span.lo = Span::dummy_with_cmt().lo;
                }

                comments.add_pure_comment(call.span.lo);
            }
        }

        call.visit_mut_children_with(self);
    }
}

fn is_pure(src: &Wtf8Atom, specifier: &Atom) -> bool {
    let Some(src) = src.as_str() else {
        return false;
    };

    match src {
        "inferno" => matches!(
            &**specifier,
            "createComponentVNode"
                | "createFragment"
                | "createPortal"
                | "createRef"
                | "createRenderer"
                | "createTextVNode"
                | "createVNode"
                | "forwardRef"
                | "directClone"
                | "findDOMFromVNode"
                | "getFlagsForElementVnode"
                | "linkEvent"
                | "normalizeProps"
                | "createElement"
                | "createClass"
        ),
        _ => false,
    }
}
