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

use parking_lot::{Mutex, RwLock};

use nac3core::{
    codegen::{CodeGenContext, CodeGenerator},
    inkwell::{module::Linkage, values::BasicValue},
    nac3parser::ast::{self, StrRef},
    symbol_resolver::{SymbolResolver, SymbolValue, ValueEnum},
    toplevel::{DefinitionId, TopLevelDef},
    typecheck::{
        type_inferencer::PrimitiveStore,
        typedef::{Type, Unifier},
    },
};

pub struct ResolverInternal {
    pub id_to_type: Mutex<HashMap<StrRef, Type>>,
    pub id_to_def: Mutex<HashMap<StrRef, DefinitionId>>,
    pub module_globals: Mutex<HashMap<StrRef, SymbolValue>>,
    pub str_store: Mutex<HashMap<String, i32>>,
}

impl ResolverInternal {
    pub fn add_id_def(&self, id: StrRef, def: DefinitionId) {
        self.id_to_def.lock().insert(id, def);
    }

    pub fn add_id_type(&self, id: StrRef, ty: Type) {
        self.id_to_type.lock().insert(id, ty);
    }

    pub fn add_module_global(&self, id: StrRef, val: SymbolValue) {
        self.module_globals.lock().insert(id, val);
    }
}

pub struct Resolver(pub Arc<ResolverInternal>);

impl SymbolResolver for Resolver {
    fn get_default_param_value(&self, expr: &ast::Expr) -> Option<SymbolValue> {
        match &expr.node {
            ast::ExprKind::Name { id, .. } => self.0.module_globals.lock().get(id).cloned(),
            _ => unimplemented!("other type of expr not supported at {}", expr.location),
        }
    }

    fn get_symbol_type(
        &self,
        unifier: &mut Unifier,
        _: &[Arc<RwLock<TopLevelDef>>],
        primitives: &PrimitiveStore,
        str: StrRef,
    ) -> Result<Type, String> {
        self.0
            .id_to_type
            .lock()
            .get(&str)
            .copied()
            .or_else(|| {
                self.0
                    .module_globals
                    .lock()
                    .get(&str)
                    .cloned()
                    .map(|v| v.get_type(primitives, unifier))
            })
            .ok_or(format!("cannot get type of {str}"))
    }

    fn get_symbol_value<'ctx>(
        &self,
        str: StrRef,
        ctx: &mut CodeGenContext<'ctx, '_>,
        generator: &mut dyn CodeGenerator,
    ) -> Option<ValueEnum<'ctx>> {
        self.0.module_globals.lock().get(&str).cloned().map(|v| {
            ctx.module
                .get_global(&str.to_string())
                .unwrap_or_else(|| {
                    let ty = v.get_type(&ctx.primitives, &mut ctx.unifier);

                    let init_val = ctx.gen_symbol_val(generator, &v, ty);
                    let llvm_ty = init_val.get_type();

                    let global = ctx.module.add_global(llvm_ty, None, &str.to_string());
                    global.set_linkage(Linkage::LinkOnceAny);
                    global.set_initializer(&init_val);

                    global
                })
                .as_basic_value_enum()
                .into()
        })
    }

    fn get_identifier_def(&self, id: StrRef) -> Result<DefinitionId, HashSet<String>> {
        self.0
            .id_to_def
            .lock()
            .get(&id)
            .copied()
            .ok_or_else(|| HashSet::from([format!("Undefined identifier `{id}`")]))
    }

    fn get_string_id(&self, s: &str) -> i32 {
        let mut str_store = self.0.str_store.lock();
        if let Some(id) = str_store.get(s) {
            *id
        } else {
            let id = i32::try_from(str_store.len())
                .expect("Symbol resolver string store size exceeds max capacity (i32::MAX)");
            str_store.insert(s.to_string(), id);
            id
        }
    }

    fn get_exception_id(&self, _: usize) -> usize {
        unimplemented!()
    }
}