use crate::ast::NodeId;
use crate::ast::compilation_unit::CompilationUnit;
use crate::ast::extern_function::ExternFunction;
use crate::ast::function::Function;
use crate::ast::parameter::Parameter;
use crate::diagnostic::Diagnostics;
use crate::semantic_analysis::diagnostic_helpers::symbol_already_declared;
use crate::semantic_analysis::scope::{Scope, ScopeId};
use crate::semantic_analysis::symbol::{FunctionSymbol, ParameterSymbol, Symbol};
use std::collections::HashMap;

pub struct SymbolCollectionContext {
    scopes: Vec<Scope>,
    nodes_to_scopes: HashMap<NodeId, ScopeId>,
    symbols: Vec<Symbol>,
    diagnostics: Diagnostics,
}

impl SymbolCollectionContext {
    pub fn new(scopes: Vec<Scope>, nodes_to_scopes: HashMap<NodeId, ScopeId>) -> Self {
        Self {
            scopes,
            nodes_to_scopes,
            symbols: Vec::new(),
            diagnostics: Diagnostics::new(),
        }
    }

    pub fn find_symbol_in_scope_for(&self, name: &str, node_id: NodeId) -> Option<&Symbol> {
        let scope_id = self.nodes_to_scopes[&node_id];
        let scope = &self.scopes[scope_id];
        if let Some(symbol_id) = scope.symbols().get(name) {
            Some(&self.symbols[*symbol_id])
        } else {
            None
        }
    }

    pub fn insert_symbol(&mut self, symbol: Symbol, node_id: NodeId) {
        let declared_name = symbol.declared_name_owned();
        self.symbols.push(symbol);
        let symbol_id = self.symbols.len() - 1;
        let scope_id = self.nodes_to_scopes[&node_id];
        let scope = &mut self.scopes[scope_id];
        scope.symbols_mut().insert(declared_name, symbol_id);
    }

    pub fn diagnostics_mut(&mut self) -> &mut Diagnostics {
        &mut self.diagnostics
    }
}

pub fn collect_symbols(compilation_unit: &CompilationUnit, ctx: &mut SymbolCollectionContext) {
    for function in compilation_unit.functions() {
        collect_symbols_function(function, ctx);
    }

    for extern_function in compilation_unit.extern_functions() {
        collect_symbols_extern_function(extern_function, ctx);
    }
}

fn collect_symbols_function(function: &Function, ctx: &mut SymbolCollectionContext) {
    // function itself
    let function_symbol = Symbol::Function(FunctionSymbol::new(
        function.declared_name_owned(),
        Some(function.declared_name_source_range()),
        false,
    ));

    // insert
    if let Some(already_declared) =
        ctx.find_symbol_in_scope_for(function.declared_name(), function.node_id())
    {
        let diagnostic = symbol_already_declared(already_declared, &function_symbol);
        ctx.diagnostics_mut().push(diagnostic);
    } else {
        ctx.insert_symbol(function_symbol, function.node_id());
    }

    // parameters
    for parameter in function.parameters() {
        collect_symbols_parameter(parameter, ctx);
    }

    // n.b. do not do statements yet, because variables are declared and resolved in the resolution pass
}

fn collect_symbols_extern_function(
    extern_function: &ExternFunction,
    ctx: &mut SymbolCollectionContext,
) {
    // function itself
    let function_symbol = Symbol::Function(FunctionSymbol::new(
        extern_function.declared_name_owned(),
        Some(extern_function.declared_name_source_range()),
        true,
    ));

    // insert function symbol
    if let Some(already_declared) =
        ctx.find_symbol_in_scope_for(extern_function.declared_name(), extern_function.node_id())
    {
        let diagnostic = symbol_already_declared(already_declared, &function_symbol);
        ctx.diagnostics_mut().push(diagnostic);
    } else {
        ctx.insert_symbol(function_symbol, extern_function.node_id());
    }

    // parameters
    for parameter in extern_function.parameters() {
        collect_symbols_parameter(parameter, ctx);
    }
}

fn collect_symbols_parameter(parameter: &Parameter, ctx: &mut SymbolCollectionContext) {
    let parameter_symbol = ParameterSymbol::new(
        parameter.declared_name_owned(),
        Some(parameter.declared_name_source_range()),
    );
    ctx.insert_symbol(Symbol::Parameter(parameter_symbol), parameter.node_id());
}