use std::collections::HashMap;

use crate::ir::{Expression, Primitive, Program, Statement, Value, ValueOrRef};
use cranelift_codegen::entity::EntityRef;
use cranelift_codegen::ir::{
    entities, types, Function, GlobalValue, InstBuilder, MemFlags, Signature, UserFuncName,
};
use cranelift_codegen::isa::CallConv;
use cranelift_codegen::Context;
use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext, Variable};
use cranelift_module::{FuncId, Linkage, Module, ModuleError};
use internment::ArcIntern;

use crate::backend::error::BackendError;
use crate::backend::Backend;

type StringTable = HashMap<ArcIntern<String>, GlobalValue>;

impl<M: Module> Backend<M> {
    pub fn compile_function(
        &mut self,
        function_name: &str,
        mut program: Program,
    ) -> Result<FuncId, BackendError> {
        let basic_signature = Signature {
            params: vec![],
            returns: vec![],
            call_conv: CallConv::SystemV,
        };

        let func_id =
            self.module
                .declare_function(function_name, Linkage::Export, &basic_signature)?;
        let mut ctx = Context::new();
        ctx.func =
            Function::with_name_signature(UserFuncName::user(0, func_id.as_u32()), basic_signature);

        let string_table = self.build_string_table(&mut ctx.func, &program)?;
        let mut variable_table = HashMap::new();
        let mut next_var_num = 1;
        let print_func_ref = self.runtime_functions.include_runtime_function(
            "print",
            &mut self.module,
            &mut ctx.func,
        )?;
        let pre_defined_symbols: HashMap<String, GlobalValue> = self
            .defined_symbols
            .iter()
            .map(|(k, v)| {
                let local_data = self.module.declare_data_in_func(*v, &mut ctx.func);
                (k.clone(), local_data)
            })
            .collect();

        let mut fctx = FunctionBuilderContext::new();
        let mut builder = FunctionBuilder::new(&mut ctx.func, &mut fctx);
        let main_block = builder.create_block();
        builder.switch_to_block(main_block);

        for stmt in program.statements.drain(..) {
            match stmt {
                Statement::Print(ann, var) => {
                    let local_name_ref = string_table.get(&var).unwrap();
                    let name_ptr = builder.ins().symbol_value(types::I64, *local_name_ref);
                    let val = ValueOrRef::Ref(ann, var).into_cranelift(
                        &mut builder,
                        &variable_table,
                        &pre_defined_symbols,
                    )?;
                    builder.ins().call(print_func_ref, &[name_ptr, val]);
                }

                Statement::Binding(_, var_name, value) => {
                    let val = match value {
                        Expression::Value(_, Value::Number(_, v)) => {
                            builder.ins().iconst(types::I64, v)
                        }

                        Expression::Reference(_, name) => {
                            let value_var_num = variable_table.get(&name).unwrap();
                            builder.use_var(Variable::new(*value_var_num))
                        }

                        Expression::Primitive(_, prim, mut vals) => {
                            let right = vals.pop().unwrap().into_cranelift(
                                &mut builder,
                                &variable_table,
                                &pre_defined_symbols,
                            )?;
                            let left = vals.pop().unwrap().into_cranelift(
                                &mut builder,
                                &variable_table,
                                &pre_defined_symbols,
                            )?;

                            match prim {
                                Primitive::Plus => builder.ins().iadd(left, right),
                                Primitive::Minus => builder.ins().isub(left, right),
                                Primitive::Times => builder.ins().imul(left, right),
                                Primitive::Divide => builder.ins().sdiv(left, right),
                            }
                        }
                    };

                    if let Some(global_id) = pre_defined_symbols.get(var_name.as_str()) {
                        let val_ptr = builder.ins().symbol_value(types::I64, *global_id);
                        builder.ins().store(MemFlags::new(), val, val_ptr, 0);
                    } else {
                        let var = Variable::new(next_var_num);
                        variable_table.insert(var_name, next_var_num);
                        next_var_num += 1;
                        builder.declare_var(var, types::I64);
                        builder.def_var(var, val);
                    }
                }
            }
        }

        builder.ins().return_(&[]);
        builder.seal_block(main_block);
        builder.finalize();

        let _ = self.module.define_function(func_id, &mut ctx)?;

        Ok(func_id)
    }

    fn build_string_table(
        &mut self,
        func: &mut Function,
        program: &Program,
    ) -> Result<StringTable, BackendError> {
        let mut string_table = HashMap::new();

        for interned_value in program.strings().drain() {
            let global_id = match self.defined_strings.get(interned_value.as_str()) {
                Some(x) => *x,
                None => self.define_string(interned_value.as_str())?,
            };
            let local_data = self.module.declare_data_in_func(global_id, func);
            string_table.insert(interned_value, local_data);
        }

        Ok(string_table)
    }
}

impl ValueOrRef {
    fn into_cranelift(
        self,
        builder: &mut FunctionBuilder,
        local_variables: &HashMap<ArcIntern<String>, usize>,
        global_variables: &HashMap<String, GlobalValue>,
    ) -> Result<entities::Value, ModuleError> {
        match self {
            ValueOrRef::Value(_, value) => match value {
                Value::Number(_base, numval) => Ok(builder.ins().iconst(types::I64, numval)),
            },

            ValueOrRef::Ref(_, name) => {
                if let Some(local_num) = local_variables.get(&name) {
                    return Ok(builder.use_var(Variable::new(*local_num)));
                }

                if let Some(global_id) = global_variables.get(name.as_str()) {
                    let val_ptr = builder.ins().symbol_value(types::I64, *global_id);
                    return Ok(builder.ins().load(types::I64, MemFlags::new(), val_ptr, 0));
                }

                Err(ModuleError::Undeclared(name.to_string()))
            }
        }
    }
}