//! AST -> `ItemTree` lowering code.

use std::cell::OnceCell;

use base_db::{Crate, FxIndexSet, SourceDatabase};
use cfg::CfgOptions;
use either::Either;
use hir_expand::{HirFileId, mod_path::PathKind, name::AsName, span_map::SpanMap};
use la_arena::Arena;
use rustc_hash::FxHashSet;
use span::{AstIdMap, FileAstId, SyntaxContext};
use syntax::{
    AstNode,
    ast::{self, HasModuleItem, HasName},
};

use crate::item_tree::{
    BigModItem, Const, Enum, ExternBlock, ExternCrate, Function, Impl, ImportAlias, Interned,
    ItemTree, Macro2, MacroCall, MacroRules, Mod, ModItemId, ModKind, ModPath, RawVisibility,
    RawVisibilityId, SmallModItem, Static, Struct, StructKind, StructValueNsCtor, Trait, TypeAlias,
    Union, Use, UseTree, UseTreeKind, VisibilityExplicitness, attrs::AttrsOrCfg,
};

pub(super) struct Ctx<'db> {
    pub(super) db: &'db dyn SourceDatabase,
    tree: ItemTree,
    source_ast_id_map: &'db AstIdMap,
    span_map: OnceCell<SpanMap<'db>>,
    file: HirFileId,
    cfg_options: OnceCell<&'db CfgOptions>,
    krate: Crate,
    top_level: Vec<ModItemId>,
    visibilities: FxIndexSet<RawVisibility>,
}

impl<'db> Ctx<'db> {
    pub(super) fn new(db: &'db dyn SourceDatabase, file: HirFileId, krate: Crate) -> Self {
        Self {
            db,
            tree: ItemTree::default(),
            source_ast_id_map: file.ast_id_map(db),
            file,
            cfg_options: OnceCell::new(),
            span_map: OnceCell::new(),
            visibilities: FxIndexSet::default(),
            top_level: Vec::new(),
            krate,
        }
    }

    #[inline]
    pub(super) fn cfg_options(&self) -> &'db CfgOptions {
        self.cfg_options.get_or_init(|| self.krate.cfg_options(self.db))
    }

    pub(super) fn span_map(&self) -> SpanMap<'_> {
        *self.span_map.get_or_init(|| self.file.span_map(self.db))
    }

    pub(super) fn lower_module_items(mut self, item_owner: &dyn HasModuleItem) -> ItemTree {
        self.top_level = item_owner.items().filter_map(|item| self.lower_mod_item(&item)).collect();
        self.tree.vis.arena = self.visibilities.into_iter().collect();
        self.tree.top_level = self.top_level.into_boxed_slice();
        self.tree
    }

    pub(super) fn lower_macro_stmts(mut self, stmts: ast::MacroStmts) -> ItemTree {
        self.top_level = stmts
            .statements()
            .filter_map(|stmt| {
                match stmt {
                    ast::Stmt::Item(item) => Some(item),
                    // Macro calls can be both items and expressions. The syntax library always treats
                    // them as expressions here, so we undo that.
                    ast::Stmt::ExprStmt(es) => match es.expr()? {
                        ast::Expr::MacroExpr(expr) => {
                            cov_mark::hit!(macro_call_in_macro_stmts_is_added_to_item_tree);
                            Some(expr.macro_call()?.into())
                        }
                        _ => None,
                    },
                    _ => None,
                }
            })
            .filter_map(|item| self.lower_mod_item(&item))
            .collect();

        if let Some(ast::Expr::MacroExpr(tail_macro)) = stmts.expr()
            && let Some(call) = tail_macro.macro_call()
        {
            cov_mark::hit!(macro_stmt_with_trailing_macro_expr);
            if let Some(mod_item) = self.lower_mod_item(&call.into()) {
                self.top_level.push(mod_item);
            }
        }

        self.tree.vis.arena = self.visibilities.into_iter().collect();
        self.tree.top_level = self.top_level.into_boxed_slice();
        self.tree
    }

    pub(super) fn lower_block(mut self, block: &ast::BlockExpr) -> ItemTree {
        self.add_attrs(ModItemId::TOP_OWNER, self.lower_attrs(block));
        self.top_level = block
            .statements()
            .filter_map(|stmt| match stmt {
                ast::Stmt::Item(item) => self.lower_mod_item(&item),
                // Macro calls can be both items and expressions. The syntax library always treats
                // them as expressions here, so we undo that.
                ast::Stmt::ExprStmt(es) => match es.expr()? {
                    ast::Expr::MacroExpr(expr) => self.lower_mod_item(&expr.macro_call()?.into()),
                    _ => None,
                },
                _ => None,
            })
            .collect();
        if let Some(ast::Expr::MacroExpr(expr)) = block.tail_expr()
            && let Some(call) = expr.macro_call()
            && let Some(mod_item) = self.lower_mod_item(&call.into())
        {
            self.top_level.push(mod_item);
        }
        self.tree.vis.arena = self.visibilities.into_iter().collect();
        self.tree.top_level = self.top_level.into_boxed_slice();
        self.tree
    }

    fn lower_mod_item(&mut self, item: &ast::Item) -> Option<ModItemId> {
        let mod_item = match item {
            ast::Item::Struct(ast) => self.lower_struct(ast)?,
            ast::Item::Union(ast) => self.lower_union(ast)?,
            ast::Item::Enum(ast) => self.lower_enum(ast)?,
            ast::Item::Fn(ast) => self.lower_function(ast)?,
            ast::Item::TypeAlias(ast) => self.lower_type_alias(ast)?,
            ast::Item::Static(ast) => self.lower_static(ast)?,
            ast::Item::Const(ast) => self.lower_const(ast),
            ast::Item::Module(ast) => self.lower_module(ast)?,
            ast::Item::Trait(ast) => self.lower_trait(ast)?,
            ast::Item::Impl(ast) => self.lower_impl(ast),
            ast::Item::Use(ast) => self.lower_use(ast)?,
            ast::Item::ExternCrate(ast) => self.lower_extern_crate(ast)?,
            ast::Item::MacroCall(ast) => self.lower_macro_call(ast)?,
            ast::Item::MacroRules(ast) => self.lower_macro_rules(ast)?,
            ast::Item::MacroDef(ast) => self.lower_macro_def(ast)?,
            ast::Item::ExternBlock(ast) => self.lower_extern_block(ast),
            // FIXME: Handle `global_asm!()`.
            ast::Item::AsmExpr(_) => return None,
        };
        let attrs = self.lower_attrs(item);
        self.add_attrs(mod_item, attrs);

        Some(mod_item)
    }

    pub(super) fn add_attrs(&mut self, item: ModItemId, attrs: AttrsOrCfg) {
        if !attrs.is_empty() {
            self.tree.attrs.push((item, attrs));
        }
    }

    fn alloc_small(&mut self, ast_id: FileAstId<ast::Item>, item: SmallModItem) -> ModItemId {
        let index = self.tree.small_data.len() as u32;
        self.tree.small_data.push((ast_id, item));
        ModItemId::new_small(index)
    }

    fn alloc_big(&mut self, ast_id: FileAstId<ast::Item>, item: BigModItem) -> ModItemId {
        let index = self.tree.big_data.len() as u32;
        self.tree.big_data.push((ast_id, item));
        ModItemId::new_big(index)
    }

    fn lower_struct(&mut self, strukt: &ast::Struct) -> Option<ModItemId> {
        let visibility = self.lower_visibility_from_ast(strukt);
        let visibility_id = self.alloc_visibility(visibility.clone());
        let name = strukt.name()?.as_name();
        let ast_id = self.source_ast_id_map.ast_id(strukt);
        let value_ns_ctor = match strukt.kind() {
            StructKind::Record(_) => StructValueNsCtor::NoValueNsCtor,
            StructKind::Unit => StructValueNsCtor::ValueNsCtorWithVis(visibility_id),
            StructKind::Tuple(fields) => self.tuple_struct_ctor_vis(visibility, fields),
        };
        let res = Struct { name, visibility: visibility_id, value_ns_ctor };
        Some(self.alloc_big(ast_id.upcast(), BigModItem::Struct(res)))
    }

    fn tuple_struct_ctor_vis(
        &mut self,
        struct_vis: RawVisibility,
        fields: ast::TupleFieldList,
    ) -> StructValueNsCtor {
        if let RawVisibility::PubSelf(_) = struct_vis {
            // `pub(self)` <= anything.
            return StructValueNsCtor::ValueNsCtorWithVis(self.alloc_visibility(struct_vis));
        }

        let mut value_ns_ctor_vis = Either::Left(struct_vis);
        for field in fields.fields() {
            let field_vis = self.lower_visibility_from_ast(&field);
            match (&field_vis, &mut value_ns_ctor_vis) {
                (RawVisibility::PubSelf(_), _) => {
                    // If there is a field private, the minimal visibility is private.
                    value_ns_ctor_vis = Either::Left(field_vis);
                    break;
                }
                (_, Either::Left(RawVisibility::Public)) => {
                    // anything <= `pub`.
                    value_ns_ctor_vis = Either::Left(field_vis);
                }
                (RawVisibility::Module(..), Either::Left(RawVisibility::PubCrate)) => {
                    // `pub(in module)` <= `pub(crate)`.
                    value_ns_ctor_vis = Either::Left(field_vis);
                }
                (
                    RawVisibility::Public | RawVisibility::PubCrate,
                    Either::Right(_)
                    | Either::Left(RawVisibility::Module(..) | RawVisibility::PubCrate),
                ) => {
                    // `pub(in module) | pub(crate)` <= `pub(crate) | pub`.
                }
                (_, Either::Left(RawVisibility::PubSelf(_))) => {
                    unreachable!("we early-exit on `PubSelf`");
                }
                (RawVisibility::Module(..), Either::Left(RawVisibility::Module(..))) => {
                    let old_vis = match std::mem::replace(
                        &mut value_ns_ctor_vis,
                        Either::Left(RawVisibility::Public),
                    ) {
                        Either::Left(it) => it,
                        _ => unreachable!(),
                    };
                    let mut multiple = FxHashSet::default();
                    multiple.insert(old_vis);
                    multiple.insert(field_vis);
                    value_ns_ctor_vis = Either::Right(multiple);
                }
                (RawVisibility::Module(..), Either::Right(multiple)) => {
                    multiple.insert(field_vis);
                }
            }
        }

        match value_ns_ctor_vis {
            Either::Left(vis) => StructValueNsCtor::ValueNsCtorWithVis(self.alloc_visibility(vis)),
            Either::Right(multiple) => {
                if multiple.len() == 1 {
                    StructValueNsCtor::ValueNsCtorWithVis(
                        self.alloc_visibility(multiple.into_iter().next().unwrap()),
                    )
                } else {
                    StructValueNsCtor::ValueNsCtorWithMinVis(
                        multiple.into_iter().map(|vis| self.alloc_visibility(vis)).collect(),
                    )
                }
            }
        }
    }

    fn lower_union(&mut self, union: &ast::Union) -> Option<ModItemId> {
        let visibility = self.lower_visibility(union);
        let name = union.name()?.as_name();
        let ast_id = self.source_ast_id_map.ast_id(union);
        let res = Union { name, visibility };
        Some(self.alloc_small(ast_id.upcast(), SmallModItem::Union(res)))
    }

    fn lower_enum(&mut self, enum_: &ast::Enum) -> Option<ModItemId> {
        let visibility = self.lower_visibility(enum_);
        let name = enum_.name()?.as_name();
        let ast_id = self.source_ast_id_map.ast_id(enum_);
        let res = Enum { name, visibility };
        Some(self.alloc_small(ast_id.upcast(), SmallModItem::Enum(res)))
    }

    fn lower_function(&mut self, func: &ast::Fn) -> Option<ModItemId> {
        let visibility = self.lower_visibility(func);
        let name = func.name()?.as_name();

        let ast_id = self.source_ast_id_map.ast_id(func);

        let res = Function { name, visibility };

        Some(self.alloc_small(ast_id.upcast(), SmallModItem::Function(res)))
    }

    fn lower_type_alias(&mut self, type_alias: &ast::TypeAlias) -> Option<ModItemId> {
        let name = type_alias.name()?.as_name();
        let visibility = self.lower_visibility(type_alias);
        let ast_id = self.source_ast_id_map.ast_id(type_alias);
        let res = TypeAlias { name, visibility };
        Some(self.alloc_small(ast_id.upcast(), SmallModItem::TypeAlias(res)))
    }

    fn lower_static(&mut self, static_: &ast::Static) -> Option<ModItemId> {
        let name = static_.name()?.as_name();
        let visibility = self.lower_visibility(static_);
        let ast_id = self.source_ast_id_map.ast_id(static_);
        let res = Static { name, visibility };
        Some(self.alloc_small(ast_id.upcast(), SmallModItem::Static(res)))
    }

    fn lower_const(&mut self, konst: &ast::Const) -> ModItemId {
        let name = konst.name().map(|it| it.as_name());
        let visibility = self.lower_visibility(konst);
        let ast_id = self.source_ast_id_map.ast_id(konst);
        let res = Const { name, visibility };
        self.alloc_small(ast_id.upcast(), SmallModItem::Const(res))
    }

    fn lower_module(&mut self, module: &ast::Module) -> Option<ModItemId> {
        let name = module.name()?.as_name();
        let visibility = self.lower_visibility(module);
        let kind = if module.semicolon_token().is_some() {
            ModKind::Outline
        } else {
            ModKind::Inline {
                items: module
                    .item_list()
                    .map(|list| {
                        list.items().filter_map(|item| self.lower_mod_item(&item)).collect()
                    })
                    .unwrap_or_else(|| {
                        cov_mark::hit!(name_res_works_for_broken_modules);
                        Box::new([]) as Box<[_]>
                    }),
            }
        };
        let ast_id = self.source_ast_id_map.ast_id(module);
        let res = Mod { name, visibility, kind };
        Some(self.alloc_big(ast_id.upcast(), BigModItem::Mod(res)))
    }

    fn lower_trait(&mut self, trait_def: &ast::Trait) -> Option<ModItemId> {
        let name = trait_def.name()?.as_name();
        let visibility = self.lower_visibility(trait_def);
        let ast_id = self.source_ast_id_map.ast_id(trait_def);

        let def = Trait { name, visibility };
        Some(self.alloc_small(ast_id.upcast(), SmallModItem::Trait(def)))
    }

    fn lower_impl(&mut self, impl_def: &ast::Impl) -> ModItemId {
        let ast_id = self.source_ast_id_map.ast_id(impl_def);
        // Note that trait impls don't get implicit `Self` unlike traits, because here they are a
        // type alias rather than a type parameter, so this is handled by the resolver.
        let res = Impl { is_trait_impl: impl_def.trait_().is_some() };
        self.alloc_small(ast_id.upcast(), SmallModItem::Impl(res))
    }

    fn lower_use(&mut self, use_item: &ast::Use) -> Option<ModItemId> {
        let visibility = self.lower_visibility(use_item);
        let ast_id = self.source_ast_id_map.ast_id(use_item);
        let (use_tree, _) = lower_use_tree(self.db, use_item.use_tree()?, &mut |range| {
            self.span_map().span_for_range(range).ctx
        })?;

        let res = Use { visibility, use_tree };
        Some(self.alloc_big(ast_id.upcast(), BigModItem::Use(res)))
    }

    fn lower_extern_crate(&mut self, extern_crate: &ast::ExternCrate) -> Option<ModItemId> {
        let name = extern_crate.name_ref()?.as_name();
        let alias = extern_crate.rename().map(|a| {
            a.name().map(|it| it.as_name()).map_or(ImportAlias::Underscore, ImportAlias::Alias)
        });
        let visibility = self.lower_visibility(extern_crate);
        let ast_id = self.source_ast_id_map.ast_id(extern_crate);

        let res = ExternCrate { name, alias, visibility };
        Some(self.alloc_big(ast_id.upcast(), BigModItem::ExternCrate(res)))
    }

    fn lower_macro_call(&mut self, m: &ast::MacroCall) -> Option<ModItemId> {
        let span_map = self.span_map();
        let path = m.path()?;
        let range = path.syntax().text_range();
        let path = Interned::new(ModPath::from_src(self.db, path, &mut |range| {
            span_map.span_for_range(range).ctx
        })?);
        let ast_id = self.source_ast_id_map.ast_id(m);
        let expand_to = hir_expand::ExpandTo::from_call_site(m);
        let res = MacroCall { path, expand_to, ctxt: span_map.span_for_range(range).ctx };
        Some(self.alloc_small(ast_id.upcast(), SmallModItem::MacroCall(res)))
    }

    fn lower_macro_rules(&mut self, m: &ast::MacroRules) -> Option<ModItemId> {
        let name = m.name()?;
        let ast_id = self.source_ast_id_map.ast_id(m);

        let res = MacroRules { name: name.as_name() };
        Some(self.alloc_small(ast_id.upcast(), SmallModItem::MacroRules(res)))
    }

    fn lower_macro_def(&mut self, m: &ast::MacroDef) -> Option<ModItemId> {
        let name = m.name()?;

        let ast_id = self.source_ast_id_map.ast_id(m);
        let visibility = self.lower_visibility(m);

        let res = Macro2 { name: name.as_name(), visibility };
        Some(self.alloc_small(ast_id.upcast(), SmallModItem::Macro2(res)))
    }

    fn lower_extern_block(&mut self, block: &ast::ExternBlock) -> ModItemId {
        let ast_id = self.source_ast_id_map.ast_id(block);
        let children: Box<[_]> = block.extern_item_list().map_or(Box::new([]), |list| {
            list.extern_items()
                .filter_map(|item| {
                    // Note: All items in an `extern` block need to be lowered as if they're outside of one
                    // (in other words, the knowledge that they're in an extern block must not be used).
                    // This is because an extern block can contain macros whose ItemTree's top-level items
                    // should be considered to be in an extern block too.
                    let mod_item = match &item {
                        ast::ExternItem::Fn(ast) => self.lower_function(ast)?,
                        ast::ExternItem::Static(ast) => self.lower_static(ast)?,
                        ast::ExternItem::TypeAlias(ty) => self.lower_type_alias(ty)?,
                        ast::ExternItem::MacroCall(call) => self.lower_macro_call(call)?,
                    };
                    let attrs = self.lower_attrs(&item);
                    self.add_attrs(mod_item, attrs);
                    Some(mod_item)
                })
                .collect()
        });

        let res = ExternBlock { children };
        self.alloc_small(ast_id.upcast(), SmallModItem::ExternBlock(res))
    }

    fn lower_visibility(&mut self, item: &dyn ast::HasVisibility) -> RawVisibilityId {
        let vis = self.lower_visibility_from_ast(item);
        self.alloc_visibility(vis)
    }

    fn lower_visibility_from_ast(&mut self, item: &dyn ast::HasVisibility) -> RawVisibility {
        visibility_from_ast(self.db, item.visibility(), &mut |range| {
            self.span_map().span_for_range(range).ctx
        })
    }

    fn alloc_visibility(&mut self, vis: RawVisibility) -> RawVisibilityId {
        match &vis {
            RawVisibility::Public => RawVisibilityId::PUB,
            RawVisibility::PubSelf(VisibilityExplicitness::Explicit) => {
                RawVisibilityId::PRIV_EXPLICIT
            }
            RawVisibility::PubSelf(VisibilityExplicitness::Implicit) => {
                RawVisibilityId::PRIV_IMPLICIT
            }
            RawVisibility::PubCrate => RawVisibilityId::PUB_CRATE,
            _ => RawVisibilityId(self.visibilities.insert_full(vis).0 as u32),
        }
    }
}

struct UseTreeLowering<'a> {
    db: &'a dyn SourceDatabase,
    mapping: Arena<ast::UseTree>,
}

impl UseTreeLowering<'_> {
    fn lower_use_tree(
        &mut self,
        tree: ast::UseTree,
        span_for_range: &mut dyn FnMut(::tt::TextRange) -> SyntaxContext,
    ) -> Option<UseTree> {
        if let Some(use_tree_list) = tree.use_tree_list() {
            let prefix = match tree.path() {
                // E.g. use something::{{{inner}}};
                None => None,
                // E.g. `use something::{inner}` (prefix is `None`, path is `something`)
                // or `use something::{path::{inner::{innerer}}}` (prefix is `something::path`, path is `inner`)
                Some(path) => {
                    match ModPath::from_src(self.db, path, span_for_range) {
                        Some(it) => Some(it),
                        None => {
                            // FIXME: report errors somewhere
                            return None;
                        }
                    }
                }
            };

            self.mapping.alloc(tree.clone());
            let list = use_tree_list
                .use_trees()
                .filter_map(|tree| self.lower_use_tree(tree, span_for_range))
                .collect();

            Some(UseTree {
                kind: UseTreeKind::Prefixed { prefix: prefix.map(Interned::new), list },
            })
        } else {
            let is_glob = tree.star_token().is_some();
            let path = match tree.path() {
                Some(path) => Some(ModPath::from_src(self.db, path, span_for_range)?),
                None => None,
            };
            let alias = tree.rename().map(|a| {
                a.name().map(|it| it.as_name()).map_or(ImportAlias::Underscore, ImportAlias::Alias)
            });
            if alias.is_some() && is_glob {
                return None;
            }

            match (path, alias, is_glob) {
                (path, None, true) => {
                    if path.is_none() {
                        cov_mark::hit!(glob_enum_group);
                    }
                    self.mapping.alloc(tree.clone());
                    Some(UseTree { kind: UseTreeKind::Glob { path: path.map(Interned::new) } })
                }
                // Globs can't be renamed
                (_, Some(_), true) | (None, None, false) => None,
                // `bla::{ as Name}` is invalid
                (None, Some(_), false) => None,
                (Some(path), alias, false) => {
                    self.mapping.alloc(tree.clone());
                    Some(UseTree { kind: UseTreeKind::Single { path: Interned::new(path), alias } })
                }
            }
        }
    }
}

pub(crate) fn lower_use_tree(
    db: &dyn SourceDatabase,
    tree: ast::UseTree,
    span_for_range: &mut dyn FnMut(::tt::TextRange) -> SyntaxContext,
) -> Option<(UseTree, Arena<ast::UseTree>)> {
    let mut lowering = UseTreeLowering { db, mapping: Arena::new() };
    let tree = lowering.lower_use_tree(tree, span_for_range)?;
    Some((tree, lowering.mapping))
}

fn private_vis() -> RawVisibility {
    RawVisibility::PubSelf(VisibilityExplicitness::Implicit)
}

pub(crate) fn visibility_from_ast(
    db: &dyn SourceDatabase,
    node: Option<ast::Visibility>,
    span_for_range: &mut dyn FnMut(::tt::TextRange) -> SyntaxContext,
) -> RawVisibility {
    let Some(node) = node else { return private_vis() };
    let path = match node.kind() {
        ast::VisibilityKind::In(path) => {
            let path = ModPath::from_src(db, path, span_for_range);
            match path {
                None => return private_vis(),
                Some(path) => path,
            }
        }
        ast::VisibilityKind::PubCrate => return RawVisibility::PubCrate,
        ast::VisibilityKind::PubSuper => ModPath::from_kind(PathKind::Super(1)),
        ast::VisibilityKind::PubSelf => {
            return RawVisibility::PubSelf(VisibilityExplicitness::Explicit);
        }
        ast::VisibilityKind::Pub => return RawVisibility::Public,
    };
    RawVisibility::Module(Interned::new(path), VisibilityExplicitness::Explicit)
}
