iowo/crates/lang/src/ast/mod_tree.rs

use std::{
    collections::{HashMap, HashSet},
    path::PathBuf,
    sync::Arc,
};

use rowan::ast::{AstNode, AstPtr};

use crate::{lst_parser::syntax_kind::SyntaxKind, SyntaxNode};

use super::{
    error::Error,
    nodes::{self, Mod, ModBody, ModName, Root},
    source_file::SourceFile,
    Loc,
};

pub struct ModuleTree {
    modules: HashMap<String, Arc<Module>>,
}

impl ModuleTree {
    pub fn parse_from_main(
        path: &PathBuf,
        main_file: &SourceFile,
    ) -> (Self, HashMap<PathBuf, SourceFile>, Vec<Error>) {
        let tree = main_file.tree();
        let mut files = HashMap::new();
        let mut errors = Vec::new();

        let entry_path = path.parent().unwrap().to_owned();
        let modules = main_file
            .top_level_modules()
            .into_iter()
            .filter_map(|m| {
                let module =
                    Module::parse_mod(m.to_node(tree.syntax()), Vec::new(), &path, &entry_path);
                match module {
                    Ok(module) => {
                        files.extend(module.1);
                        errors.extend(module.2);
                        Some((module.0.name(), Arc::new(module.0)))
                    }
                    Err(err) => {
                        errors.push(err);
                        None
                    }
                }
            })
            .collect::<HashMap<String, Arc<Module>>>();

        (Self { modules }, files, errors)
    }

    pub fn print_tree(&self, lst: &Root) {
        let name = "main";
        print_tree(&name, &self.modules, 0)
    }
}

pub struct Module {
    path: Vec<String>,
    name: String,
    kind: ModuleKind,
    children: HashMap<String, Arc<Module>>,
    parent: Option<Arc<Module>>,
}

impl Module {
    fn parse_mod(
        module: Mod,
        cur_path: Vec<String>,
        cur_file: &PathBuf,
        entry_path: &PathBuf,
    ) -> Result<(Self, HashMap<PathBuf, SourceFile>, Vec<Error>), Error> {
        dbg!(cur_file);
        dbg!(entry_path);
        let children = module
            .syntax()
            .children()
            // .map(|n| n.kind())
            .collect::<Vec<_>>();

        if children.len() == 1 {
            let name = &children[0];
            assert_eq!(name.kind(), SyntaxKind::MODULE_NAME);
            return Self::parse_file_mod(
                name.text().to_string(),
                Loc::new(cur_file.clone(), &module),
                cur_path,
                entry_path,
            );
        } else if children.len() == 2 {
            let name = &children[0];
            assert_eq!(name.kind(), SyntaxKind::MODULE_NAME);
            let body = &children[1];
            assert_eq!(body.kind(), SyntaxKind::MODULE_BODY);
            return Ok(Self::parse_inline_mod(
                module, cur_path, cur_file, entry_path,
            ));
        }
        todo!()
    }

    fn parse_file_mod(
        name: String,
        decl: Loc<Mod>,
        mut cur_path: Vec<String>,
        entry_path: &PathBuf,
    ) -> Result<(Self, HashMap<PathBuf, SourceFile>, Vec<Error>), Error> {
        let mut mod_file_path = entry_path.to_owned();

        for el in &cur_path {
            mod_file_path.push(format!("{el}/"));
        }
        mod_file_path.push(format!("{name}.owo"));

        let mut files = HashMap::new();
        let mut errors = Vec::new();

        let source = match std::fs::read_to_string(dbg!(&mod_file_path)) {
            Ok(f) => f,
            Err(e) => return Err(Error::FailedToOpenFileMod(mod_file_path, e)),
        };

        let (source_file, file_errors) = SourceFile::parse_from(mod_file_path.clone(), source);
        errors.extend(file_errors);

        let tree = source_file.tree();
        let old_path = cur_path.clone();
        cur_path.push(name.clone());
        let children = source_file
            .top_level_modules()
            .into_iter()
            .filter_map(|m| {
                let module = Module::parse_mod(
                    m.to_node(tree.syntax()),
                    cur_path.clone(),
                    &mod_file_path,
                    &entry_path,
                );
                match module {
                    Ok(module) => {
                        files.extend(module.1);
                        errors.extend(module.2);
                        Some((module.0.name(), Arc::new(module.0)))
                    }
                    Err(err) => {
                        errors.push(err);
                        None
                    }
                }
            })
            .collect::<HashMap<String, Arc<Module>>>();

        files.insert(mod_file_path.clone(), source_file);

        Ok((
            Self {
                path: old_path,
                name,
                kind: ModuleKind::File {
                    declaration: decl,
                    file_id: mod_file_path,
                },
                children,
                parent: None,
            },
            files,
            errors,
        ))
    }

    fn parse_inline_mod(
        module: Mod,
        mut cur_path: Vec<String>,
        cur_file: &PathBuf,
        entry_path: &PathBuf,
    ) -> (Self, HashMap<PathBuf, SourceFile>, Vec<Error>) {
        let mut children = module.syntax().children().collect::<Vec<_>>();

        let body = ModBody::cast(children.pop().unwrap()).unwrap();
        let name = ModName::cast(children.pop().unwrap()).unwrap();

        let mut files = HashMap::new();
        let mut errors = Vec::new();

        let old_path = cur_path.clone();
        cur_path.push(name.syntax().to_string());
        let children = body
            .syntax()
            .children()
            .filter_map(|node| Mod::cast(node))
            .filter_map(|m| {
                let m = Self::parse_mod(m, cur_path.clone(), cur_file, entry_path);
                match m {
                    Ok(module) => {
                        files.extend(module.1);
                        errors.extend(module.2);
                        Some((module.0.name(), Arc::new(module.0)))
                    }
                    Err(err) => {
                        errors.push(err);
                        None
                    }
                }
            })
            .collect::<HashMap<String, Arc<Module>>>();

        (
            Self {
                name: name.syntax().text().to_string(),
                kind: ModuleKind::Inline(Loc::new(cur_file.to_owned(), &module)),
                children,
                parent: None,
                path: old_path,
            },
            files,
            errors,
        )
    }

    pub fn name(&self) -> String {
        // self.name.to_node(lst.syntax()).syntax().text().to_string()
        self.name.clone()
    }
}

fn print_tree(name: &str, children: &HashMap<String, Arc<Module>>, level: u32) {
    const INDENT_STR: &str = "    ";

    for _ in 0..level {
        print!("{}", INDENT_STR);
    }
    print!("{name}\n");

    for (name, module) in children {
        print_tree(name, &module.children, level + 1);
    }
}

enum ModuleKind {
    Inline(Loc<nodes::Mod>),
    File {
        declaration: Loc<nodes::Mod>,
        file_id: PathBuf,
    },
}