use java_ast_parser::ast::{self, ClassCell, EnumCell, GetIdent, InterfaceCell, Root};
use orx_tree::{Bfs, Dyn, DynTree, NodeIdx, NodeRef};
use std::{
    collections::{HashMap, HashSet},
    ops::{Deref, DerefMut},
    sync::Arc,
};

#[derive(Debug, Clone)]
pub enum TreeNode {
    Root,
    Package(Arc<String>),
    Class(ClassCell),
    Enum(EnumCell),
    Interface(InterfaceCell),
}

#[derive(Debug, Clone)]
pub enum ResolvedType {
    Class(ClassCell),
    Interface(InterfaceCell),
}

impl ResolvedType {
    fn from_node(node: &TreeNode) -> Option<Self> {
        match node {
            TreeNode::Class(class_cell) => Some(Self::Class(class_cell.clone())),
            TreeNode::Interface(interface_cell) => Some(Self::Interface(interface_cell.clone())),
            _ => None,
        }
    }
}

impl TreeNode {
    pub fn ident(&self) -> Option<&'_ str> {
        match self {
            TreeNode::Root => None,
            TreeNode::Package(ident) => Some(ident.as_str()),
            TreeNode::Class(cell) => Some(cell.ident()),
            TreeNode::Enum(cell) => Some(cell.ident()),
            TreeNode::Interface(cell) => Some(cell.ident()),
        }
    }
}

impl std::hash::Hash for TreeNode {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        if let Some(ident) = self.ident() {
            ident.hash(state);
        } else {
            0.hash(state);
        }
    }
}

impl std::cmp::PartialEq for TreeNode {
    fn eq(&self, other: &Self) -> bool {
        match (self.ident(), other.ident()) {
            (None, None) => true,
            (Some(l), Some(r)) => l == r,
            _ => false,
        }
    }
}

impl std::cmp::Eq for TreeNode {}

impl From<&str> for TreeNode {
    fn from(value: &str) -> Self {
        Self::Package(Arc::from(value.to_string()))
    }
}

impl From<&ClassCell> for TreeNode {
    fn from(value: &ClassCell) -> Self {
        Self::Class(value.clone())
    }
}

impl From<&InterfaceCell> for TreeNode {
    fn from(value: &InterfaceCell) -> Self {
        Self::Interface(value.clone())
    }
}

impl From<&EnumCell> for TreeNode {
    fn from(value: &EnumCell) -> Self {
        Self::Enum(value.clone())
    }
}
pub type IndexTree = DynTree<TreeNode>;

#[derive(Debug, Clone)]
pub struct SharedLocalIndex {
    tree: Arc<IndexTree>,
    reverse_local: Arc<HashMap<ClassCell, NodeIdx<Dyn<TreeNode>>>>,
}

#[derive(Debug, Clone)]
pub struct PackageIndexTree {
    package: String,
    inner: IndexTree,
}

impl Deref for PackageIndexTree {
    type Target = DynTree<TreeNode>;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl DerefMut for PackageIndexTree {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

pub fn merge_index_trees(
    target_tree: &mut DynTree<TreeNode>,
    target_node_idx: NodeIdx<Dyn<TreeNode>>,
    source_tree: &DynTree<TreeNode>,
    source_node_idx: NodeIdx<Dyn<TreeNode>>,
) {
    for source_child in source_tree.node(source_node_idx).children() {
        let same_target_idx =
            target_tree
                .node(target_node_idx)
                .children()
                .find_map(|target_child| {
                    if target_child.data() == source_child.data() {
                        Some(target_child.idx())
                    } else {
                        None
                    }
                });

        if let Some(same_target_idx) = same_target_idx {
            merge_index_trees(
                target_tree,
                same_target_idx,
                source_tree,
                source_child.idx(),
            );
        } else {
            target_tree
                .node_mut(target_node_idx)
                .push_child_tree(source_child.as_cloned_subtree());
        }
    }
}

impl PackageIndexTree {
    pub fn from_ast(ast: &Root) -> Self {
        fn walk_interface(
            tree: &mut DynTree<TreeNode>,
            parent_idx: NodeIdx<Dyn<TreeNode>>,
            self_cell: &InterfaceCell,
        ) {
            let self_idx = {
                let mut parent_mut = tree.node_mut(parent_idx);
                parent_mut.push_child(TreeNode::from(self_cell))
            };

            let self_ref = self_cell.borrow();

            for class_cell in &self_ref.classes {
                walk_class(tree, self_idx, class_cell);
            }

            for interface_cell in &self_ref.interfaces {
                walk_interface(tree, self_idx, interface_cell);
            }

            for enum_cell in &self_ref.enums {
                walk_enum(tree, self_idx, enum_cell);
            }
        }

        fn walk_class(
            tree: &mut DynTree<TreeNode>,
            parent_idx: NodeIdx<Dyn<TreeNode>>,
            self_cell: &ClassCell,
        ) {
            let self_idx = {
                let mut parent_mut = tree.node_mut(parent_idx);
                parent_mut.push_child(TreeNode::from(self_cell))
            };

            let self_ref = self_cell.borrow();

            for class_cell in &self_ref.classes {
                walk_class(tree, self_idx, class_cell);
            }

            for interface_cell in &self_ref.interfaces {
                walk_interface(tree, self_idx, interface_cell);
            }

            for enum_cell in &self_ref.enums {
                walk_enum(tree, self_idx, enum_cell);
            }
        }

        fn walk_enum(
            tree: &mut DynTree<TreeNode>,
            parent_idx: NodeIdx<Dyn<TreeNode>>,
            self_cell: &EnumCell,
        ) {
            let self_idx = {
                let mut parent_mut = tree.node_mut(parent_idx);
                parent_mut.push_child(TreeNode::from(self_cell))
            };

            let self_ref = self_cell.borrow();

            for class_cell in &self_ref.classes {
                walk_class(tree, self_idx, class_cell);
            }

            for interface_cell in &self_ref.interfaces {
                walk_interface(tree, self_idx, interface_cell);
            }

            for enum_cell in &self_ref.enums {
                walk_enum(tree, self_idx, enum_cell);
            }
        }

        let mut tree = DynTree::new(TreeNode::Root);
        let root_idx = tree.root().idx();

        for interface_cell in &ast.interfaces {
            walk_interface(&mut tree, root_idx, interface_cell);
        }

        for class_cell in &ast.classes {
            walk_class(&mut tree, root_idx, class_cell);
        }

        for enum_cell in &ast.enums {
            walk_enum(&mut tree, root_idx, enum_cell);
        }

        Self {
            package: ast.package.to_string(),
            inner: tree,
        }
    }

    pub fn package(&self) -> &str {
        &self.package
    }

    pub fn merge_with(&mut self, other: &Self) {
        let self_idx = self.inner.root().idx();
        let other_idx = other.inner.root().idx();

        merge_index_trees(&mut self.inner, self_idx, &other.inner, other_idx);
    }

    pub fn shared_local_index(&self) -> SharedLocalIndex {
        let tree = Arc::new(self.inner.clone());
        let mut reverse_local = HashMap::new();

        for idx in tree.root().indices::<Bfs>() {
            let TreeNode::Class(class_cell) = tree.node(idx).data() else {
                continue;
            };

            reverse_local.insert(class_cell.clone(), idx);
        }

        SharedLocalIndex {
            tree,
            reverse_local: Arc::new(reverse_local),
        }
    }
}

#[derive(Debug, Clone)]
pub struct GlobalIndexTree(IndexTree);

impl Deref for GlobalIndexTree {
    type Target = DynTree<TreeNode>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for GlobalIndexTree {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl<'a> FromIterator<&'a PackageIndexTree> for GlobalIndexTree {
    fn from_iter<IT: IntoIterator<Item = &'a PackageIndexTree>>(iter: IT) -> Self {
        let mut tree = IndexTree::new(TreeNode::Root);
        let root_idx = tree.root().idx();

        for package_index_tree in iter {
            let package_idx = {
                let mut current_idx = root_idx;

                for node in package_index_tree.package().split('.').map(TreeNode::from) {
                    let node_idx = tree.node(current_idx).children().find_map(|x| {
                        if x.data() == &node {
                            Some(x.idx())
                        } else {
                            None
                        }
                    });

                    current_idx =
                        node_idx.unwrap_or_else(|| tree.node_mut(current_idx).push_child(node));
                }

                current_idx
            };

            let mut package_mut = tree.node_mut(package_idx);

            for child in package_index_tree.root().children() {
                package_mut.push_child_tree(child.as_cloned_subtree());
            }
        }

        Self(tree)
    }
}

impl GlobalIndexTree {
    pub fn search_path<'a, I>(&self, query: I) -> Option<ResolvedType>
    where
        I: IntoIterator<Item = &'a str>,
    {
        let mut current_idx = self.0.root().idx();

        for ident in query {
            let node_idx = self.0.node(current_idx).children().find_map(|x| {
                if x.data().ident().is_some_and(|x| x == ident) {
                    Some(x.idx())
                } else {
                    None
                }
            })?;

            current_idx = node_idx;
        }

        ResolvedType::from_node(self.0.node(current_idx).data())
    }

    pub fn search(&self, query: &ast::QualifiedType) -> Option<ResolvedType> {
        let parts = query
            .iter()
            .map(|query_part| {
                let ast::TypeName::Ident(ident) = &query_part.name else {
                    return None;
                };

                Some(ident.as_str())
            })
            .collect::<Option<Vec<_>>>()?;

        self.search_path(parts)
    }
}

#[derive(Debug, Clone)]
pub struct ImportedIndexTree {
    global: Arc<GlobalIndexTree>,
    imports: Box<[String]>,
}

impl ImportedIndexTree {
    pub fn from_imports<'a, I>(import_iter: I, global_index_tree: Arc<GlobalIndexTree>) -> Self
    where
        I: IntoIterator<Item = &'a str>,
    {
        let mut seen = HashSet::new();
        let mut imports = Vec::new();
        for import in import_iter {
            let import = import.to_string();
            if seen.insert(import.clone()) {
                imports.push(import);
            }
        }

        Self {
            global: global_index_tree,
            imports: imports.into_boxed_slice(),
        }
    }

    pub fn search(&self, query: &ast::QualifiedType) -> Option<ResolvedType> {
        let query_parts = query
            .iter()
            .map(|query_part| {
                let ast::TypeName::Ident(ident) = &query_part.name else {
                    return None;
                };

                Some(ident.as_str())
            })
            .collect::<Option<Vec<_>>>()?;

        for import in &self.imports {
            if let Some(prefix) = import.strip_suffix(".*") {
                let resolved = self
                    .global
                    .search_path(prefix.split('.').chain(query_parts.iter().copied()));

                if resolved.is_some() {
                    return resolved;
                }
                continue;
            }

            let import_parts = import.split('.').collect::<Vec<_>>();
            if import_parts.last().copied() != query_parts.first().copied() {
                continue;
            }

            let resolved = self.global.search_path(
                import_parts[..import_parts.len().saturating_sub(1)]
                    .iter()
                    .copied()
                    .chain(query_parts.iter().copied()),
            );

            if resolved.is_some() {
                return resolved;
            }
        }

        None
    }
}

#[derive(Debug, Clone)]
pub struct LocalIndexTree {
    global: Arc<GlobalIndexTree>,
    imported: ImportedIndexTree,
    local: Arc<IndexTree>,
    reverse_local: Arc<HashMap<ClassCell, NodeIdx<Dyn<TreeNode>>>>,
}

impl LocalIndexTree {
    pub fn new(
        global: Arc<GlobalIndexTree>,
        imported: ImportedIndexTree,
        shared_local: SharedLocalIndex,
    ) -> Self {
        Self {
            global,
            imported,
            local: shared_local.tree,
            reverse_local: shared_local.reverse_local,
        }
    }

    pub fn search_global(&self, query: &ast::QualifiedType) -> Option<ResolvedType> {
        self.global.search(query)
    }

    pub fn search_imported(&self, query: &ast::QualifiedType) -> Option<ResolvedType> {
        self.imported.search(query)
    }

    pub fn search_local(
        &self,
        scope: Option<&ClassCell>,
        query: &ast::QualifiedType,
    ) -> Option<ResolvedType> {
        let root_idx = scope
            .and_then(|x| self.reverse_local.get(x))
            .cloned()
            .unwrap_or_else(|| self.local.root().idx());

        let try_parent = || {
            scope?;

            if let Some(parent_node) = self.local.node(root_idx).parent() {
                let scope = if let TreeNode::Class(class_cell) = parent_node.data() {
                    Some(class_cell)
                } else {
                    None
                };

                self.search_local(scope, query)
            } else {
                None
            }
        };

        let mut current_idx = root_idx;

        for query_part in query.iter() {
            let ast::TypeName::Ident(ident) = &query_part.name else {
                return None;
            };

            let Some(node_idx) = self.local.node(current_idx).children().find_map(|x| {
                if x.data().ident().is_some_and(|x| x == ident) {
                    Some(x.idx())
                } else {
                    None
                }
            }) else {
                return try_parent();
            };

            current_idx = node_idx;
        }

        ResolvedType::from_node(self.local.node(current_idx).data()).or_else(try_parent)
    }

    pub fn search(
        &self,
        scope: Option<&ClassCell>,
        query: &ast::QualifiedType,
    ) -> Option<ResolvedType> {
        self.search_local(scope, query)
            .or_else(|| self.search_imported(query))
            .or_else(|| self.search_global(query))
    }
}
