use crate::ast::compilation_unit::CompilationUnit;
use crate::diagnostic::Diagnostics;
use crate::ir::ir_class::IrClass;
use crate::ir::ir_function::IrFunction;
use crate::parser::parse_compilation_unit;
use crate::symbol_table::SymbolTable;
use crate::symbol_table::util::try_insert_symbols_into;
use std::collections::HashMap;
use std::rc::Rc;

pub type Filename = Rc<str>;
pub type FileId = usize;

fn parse_compilation_units(
    inputs: &HashMap<FileId, &str>,
) -> Result<HashMap<FileId, CompilationUnit>, Diagnostics> {
    let mut parse_diagnostics = Vec::new();
    let mut compilation_units = HashMap::new();
    for (file_id, source) in inputs {
        let (compilation_unit, mut ds) = parse_compilation_unit(source, Some(*file_id));
        parse_diagnostics.append(&mut ds);
        compilation_units.insert(*file_id, compilation_unit);
    }
    if parse_diagnostics.is_empty() {
        Ok(compilation_units)
    } else {
        Err(parse_diagnostics)
    }
}

pub fn compile_compilation_units(
    inputs: &HashMap<FileId, &str>,
    symbol_table: &mut SymbolTable,
) -> Result<(Vec<IrClass>, Vec<IrFunction>), Diagnostics> {
    let mut compilation_units = parse_compilation_units(inputs)?;

    // init scopes
    for compilation_unit in compilation_units.values_mut() {
        compilation_unit.init_scopes(symbol_table);
    }

    // gather unordered symbols
    let all_symbols = compilation_units
        .values()
        .flat_map(|compilation_unit| compilation_unit.declared_symbols())
        .collect::<Vec<_>>();
    try_insert_symbols_into(all_symbols, symbol_table)?;

    // now we can just finish each compilation unit, since we have the symbols
    let mut ir_classes = Vec::new();
    let mut ir_functions = Vec::new();
    let mut diagnostics = Vec::new();
    for compilation_unit in compilation_units.values() {
        let (nodes_to_symbols, mut ds) = compilation_unit.resolve_names(symbol_table);

        // in the future, we'll ideally be able to *actually* continue with the following steps
        // instead of aborting here, but this needs to be tested :)
        if !ds.is_empty() {
            diagnostics.append(&mut ds);
            continue;
        }

        let (mut symbols_to_types, mut ds) = compilation_unit.declared_types(&nodes_to_symbols);
        if !ds.is_empty() {
            diagnostics.append(&mut ds);
            continue;
        }

        let (sts, nodes_to_types, mut ds) =
            compilation_unit.resolve_types(&nodes_to_symbols, &symbols_to_types);
        if !ds.is_empty() {
            diagnostics.append(&mut ds);
            continue;
        }

        // merge
        for (symbol, type_info) in sts {
            symbols_to_types.insert(symbol, type_info);
        }

        let (mut classes, mut functions) =
            compilation_unit.lower_to_ir(&nodes_to_symbols, &symbols_to_types, &nodes_to_types);
        ir_classes.append(&mut classes);
        ir_functions.append(&mut functions);
    }

    if diagnostics.is_empty() {
        Ok((ir_classes, ir_functions))
    } else {
        Err(diagnostics)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::symbol::class_symbol::ClassSymbol;

    fn prepare_symbol_table(symbol_table: &mut SymbolTable) {
        let global_scope = symbol_table.push_module_scope("global scope");
        let any_symbol = ClassSymbol::new(
            &"Any".into(),
            None,
            vec!["Any".into()],
            false,
            global_scope,
            Vec::new(),
            None,
            Vec::new(),
            Vec::new(),
        );
        symbol_table.insert_class_symbol(Rc::new(any_symbol));

        let void_symbol = ClassSymbol::new(
            &"Void".into(),
            None,
            vec!["Void".into()],
            false,
            global_scope,
            Vec::new(),
            None,
            Vec::new(),
            Vec::new(),
        );
        symbol_table.insert_class_symbol(Rc::new(void_symbol));
    }

    #[test]
    fn hello_world() -> Result<(), Diagnostics> {
        let input = "
            extern fn println(msg: Any) -> Void

            fn main()
                println(\"Hello, World!\")
            end
        ";
        let mut inputs = HashMap::new();
        inputs.insert(0, input);
        let mut symbol_table = SymbolTable::new();
        prepare_symbol_table(&mut symbol_table);
        let (ir_classes, ir_functions) = compile_compilation_units(&inputs, &mut symbol_table)?;
        assert_eq!(ir_classes.len(), 0);
        assert_eq!(ir_functions.len(), 1);
        Ok(())
    }
}