use crate::eval::PrimitiveType;
use crate::ir::{Expression, Primitive, Program, TopLevel, Type, Value, ValueOrRef, Variable};
use crate::syntax::{ConstantType, Location};
use crate::util::scoped_map::ScopedMap;
use cranelift_codegen::ir::{
    self, entities, types, AbiParam, Function, GlobalValue, InstBuilder, Signature, UserFuncName,
};
use cranelift_codegen::isa::CallConv;
use cranelift_codegen::Context;
use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
use cranelift_module::{FuncId, Linkage, Module};
use internment::ArcIntern;

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

/// When we're talking about variables, it's handy to just have a table that points
/// from a variable to "what to do if you want to reference this variable", which is
/// agnostic about whether the variable is local, global, an argument, etc. Since
/// the type of that function is a little bit annoying, we summarize it here.
struct ReferenceBuilder {
    ir_type: ConstantType,
    cranelift_type: cranelift_codegen::ir::Type,
    local_data: GlobalValue,
}

impl ReferenceBuilder {
    fn refer_to(&self, builder: &mut FunctionBuilder) -> (entities::Value, ConstantType) {
        let value = builder
            .ins()
            .symbol_value(self.cranelift_type, self.local_data);
        (value, self.ir_type)
    }
}

impl<M: Module> Backend<M> {
    /// Translate the given IR type into an ABI parameter type for cranelift, as
    /// best as possible.
    fn translate_type(&self, t: &Type) -> AbiParam {
        let (value_type, extension) = match t {
            Type::Function(_, _) => (
                types::Type::triple_pointer_type(&self.platform),
                ir::ArgumentExtension::None,
            ),
            Type::Primitive(PrimitiveType::Void) => (types::I8, ir::ArgumentExtension::None), // FIXME?
            Type::Primitive(PrimitiveType::I8) => (types::I8, ir::ArgumentExtension::Sext),
            Type::Primitive(PrimitiveType::I16) => (types::I16, ir::ArgumentExtension::Sext),
            Type::Primitive(PrimitiveType::I32) => (types::I32, ir::ArgumentExtension::Sext),
            Type::Primitive(PrimitiveType::I64) => (types::I64, ir::ArgumentExtension::Sext),
            Type::Primitive(PrimitiveType::U8) => (types::I8, ir::ArgumentExtension::Uext),
            Type::Primitive(PrimitiveType::U16) => (types::I16, ir::ArgumentExtension::Uext),
            Type::Primitive(PrimitiveType::U32) => (types::I32, ir::ArgumentExtension::Uext),
            Type::Primitive(PrimitiveType::U64) => (types::I64, ir::ArgumentExtension::Uext),
        };

        AbiParam {
            value_type,
            purpose: ir::ArgumentPurpose::Normal,
            extension,
        }
    }

    /// Compile the given program.
    ///
    /// The returned value is a `FuncId` that represents a function that runs all the statements
    /// found in the program, which will be compiled using the given function name. (If there
    /// are no such statements, the function will do nothing.)
    pub fn compile_program(
        &mut self,
        function_name: &str,
        program: Program<Type>,
    ) -> Result<FuncId, BackendError> {
        let mut generated_body = vec![];

        for item in program.items {
            match item {
                TopLevel::Function(name, args, rettype, body) => {
                    self.compile_function(name.as_str(), &args, rettype, body)?;
                }

                TopLevel::Statement(stmt) => {
                    generated_body.push(stmt);
                }
            }
        }

        let void = Type::Primitive(PrimitiveType::Void);
        self.compile_function(
            function_name,
            &[],
            void.clone(),
            Expression::Block(Location::manufactured(), void, generated_body),
        )
    }

    /// Compile the given function.
    pub fn compile_function(
        &mut self,
        function_name: &str,
        arguments: &[(Variable, Type)],
        return_type: Type,
        body: Expression<Type>,
    ) -> Result<FuncId, BackendError> {
        let basic_signature = Signature {
            params: arguments
                .iter()
                .map(|(_, t)| self.translate_type(t))
                .collect(),
            returns: if return_type == Type::Primitive(PrimitiveType::Void) {
                vec![]
            } else {
                vec![self.translate_type(&return_type)]
            },
            call_conv: CallConv::triple_default(&self.platform),
        };

        // this generates the handle for the function that we'll eventually want to
        // return to the user. For now, we declare all functions defined by this
        // function as public/global/exported, although we may want to reconsider
        // this decision later.
        let func_id =
            self.module
                .declare_function(function_name, Linkage::Export, &basic_signature)?;

        // Next we have to generate the compilation context for the rest of this
        // function. Currently, we generate a fresh context for every function.
        // Since we're only generating one function per `Program`, this makes
        // complete sense. However, in the future, we may want to revisit this
        // decision.
        let mut ctx = Context::new();
        let user_func_name = UserFuncName::user(0, func_id.as_u32());
        ctx.func = Function::with_name_signature(user_func_name, basic_signature);

        // Let's start creating the variable table we'll use when we're dereferencing
        // them later. This table is a little interesting because instead of pointing
        // from data to data, we're going to point from data (the variable) to an
        // action to take if we encounter that variable at some later point. This
        // makes it nice and easy to have many different ways to access data, such
        // as globals, function arguments, etc.
        let mut variables: ScopedMap<ArcIntern<String>, ReferenceBuilder> = ScopedMap::new();

        // At the outer-most scope of things, we'll put global variables we've defined
        // elsewhere in the program.
        for (name, (data_id, ty)) in self.defined_symbols.iter() {
            let local_data = self.module.declare_data_in_func(*data_id, &mut ctx.func);
            let cranelift_type = ir::Type::from(*ty);
            variables.insert(
                name.clone(),
                ReferenceBuilder {
                    cranelift_type,
                    local_data,
                    ir_type: *ty,
                },
            );
        }

        // Once we have these, we're going to actually push a level of scope and
        // add our arguments. We push scope because if there happen to be any with
        // the same name (their shouldn't be, but just in case), we want the arguments
        // to win.
        variables.new_scope();

        // Finally (!), we generate the function builder that we're going to use to
        // make this function!
        let mut fctx = FunctionBuilderContext::new();
        let mut builder = FunctionBuilder::new(&mut ctx.func, &mut fctx);

        // Make the initial block to put instructions in. Later, when we have control
        // flow, we might add more blocks after this one. But, for now, we only have
        // the one block.
        let main_block = builder.create_block();
        builder.switch_to_block(main_block);

        let (value, _) = self.compile_expression(body, &mut variables, &mut builder)?;

        // Now that we're done, inject a return function (one with no actual value; basically
        // the equivalent of Rust's `return;`). We then seal the block (which lets Cranelift
        // know that the block is done), and then finalize the function (which lets Cranelift
        // know we're done with the function).
        builder.ins().return_(&[value]);
        builder.seal_block(main_block);
        builder.finalize();

        // This is a little odd. We want to tell the rest of Cranelift about this function,
        // so we register it using the function ID and our builder context. However, the
        // result of this function isn't actually super helpful. So we ignore it, unless
        // it's an error.
        self.module.define_function(func_id, &mut ctx)?;

        // done!
        Ok(func_id)
    }

    /// Compile an expression, returning the Cranelift Value for the expression and
    /// its type.
    fn compile_expression(
        &mut self,
        expr: Expression<Type>,
        variables: &mut ScopedMap<Variable, ReferenceBuilder>,
        builder: &mut FunctionBuilder,
    ) -> Result<(entities::Value, ConstantType), BackendError> {
        match expr {
            Expression::Atomic(x) => self.compile_value_or_ref(x, variables, builder),
            Expression::Cast(_, target_type, valref) => {
                let (val, val_type) = self.compile_value_or_ref(valref, variables, builder)?;

                match (val_type, &target_type) {
                    (ConstantType::I8, Type::Primitive(PrimitiveType::I8)) => Ok((val, val_type)),
                    (ConstantType::I8, Type::Primitive(PrimitiveType::I16)) => {
                        Ok((builder.ins().sextend(types::I16, val), ConstantType::I16))
                    }
                    (ConstantType::I8, Type::Primitive(PrimitiveType::I32)) => {
                        Ok((builder.ins().sextend(types::I32, val), ConstantType::I32))
                    }
                    (ConstantType::I8, Type::Primitive(PrimitiveType::I64)) => {
                        Ok((builder.ins().sextend(types::I64, val), ConstantType::I64))
                    }

                    (ConstantType::I16, Type::Primitive(PrimitiveType::I16)) => Ok((val, val_type)),
                    (ConstantType::I16, Type::Primitive(PrimitiveType::I32)) => {
                        Ok((builder.ins().sextend(types::I32, val), ConstantType::I32))
                    }
                    (ConstantType::I16, Type::Primitive(PrimitiveType::I64)) => {
                        Ok((builder.ins().sextend(types::I64, val), ConstantType::I64))
                    }

                    (ConstantType::I32, Type::Primitive(PrimitiveType::I32)) => Ok((val, val_type)),
                    (ConstantType::I32, Type::Primitive(PrimitiveType::I64)) => {
                        Ok((builder.ins().sextend(types::I64, val), ConstantType::I64))
                    }

                    (ConstantType::I64, Type::Primitive(PrimitiveType::I64)) => Ok((val, val_type)),

                    (ConstantType::U8, Type::Primitive(PrimitiveType::U8)) => Ok((val, val_type)),
                    (ConstantType::U8, Type::Primitive(PrimitiveType::U16)) => {
                        Ok((builder.ins().uextend(types::I16, val), ConstantType::U16))
                    }
                    (ConstantType::U8, Type::Primitive(PrimitiveType::U32)) => {
                        Ok((builder.ins().uextend(types::I32, val), ConstantType::U32))
                    }
                    (ConstantType::U8, Type::Primitive(PrimitiveType::U64)) => {
                        Ok((builder.ins().uextend(types::I64, val), ConstantType::U64))
                    }

                    (ConstantType::U16, Type::Primitive(PrimitiveType::U16)) => Ok((val, val_type)),
                    (ConstantType::U16, Type::Primitive(PrimitiveType::U32)) => {
                        Ok((builder.ins().uextend(types::I32, val), ConstantType::U32))
                    }
                    (ConstantType::U16, Type::Primitive(PrimitiveType::U64)) => {
                        Ok((builder.ins().uextend(types::I64, val), ConstantType::U64))
                    }

                    (ConstantType::U32, Type::Primitive(PrimitiveType::U32)) => Ok((val, val_type)),
                    (ConstantType::U32, Type::Primitive(PrimitiveType::U64)) => {
                        Ok((builder.ins().uextend(types::I64, val), ConstantType::U64))
                    }

                    (ConstantType::U64, Type::Primitive(PrimitiveType::U64)) => Ok((val, val_type)),

                    _ => Err(BackendError::InvalidTypeCast {
                        from: val_type.into(),
                        to: target_type,
                    }),
                }
            }

            Expression::Primitive(_, _, prim, mut vals) => {
                let mut values = vec![];
                let mut first_type = None;

                for val in vals.drain(..) {
                    let (compiled, compiled_type) =
                        self.compile_value_or_ref(val, variables, builder)?;

                    if let Some(leftmost_type) = first_type {
                        assert_eq!(leftmost_type, compiled_type);
                    } else {
                        first_type = Some(compiled_type);
                    }

                    values.push(compiled);
                }

                let first_type = first_type.expect("primitive op has at least one argument");

                // then we just need to tell Cranelift how to do each of our primitives! Much
                // like Statements, above, we probably want to eventually shuffle this off into
                // a separate function (maybe something off `Primitive`), but for now it's simple
                // enough that we just do the `match` here.
                match prim {
                    Primitive::Plus => Ok((builder.ins().iadd(values[0], values[1]), first_type)),
                    Primitive::Minus if values.len() == 2 => {
                        Ok((builder.ins().isub(values[0], values[1]), first_type))
                    }
                    Primitive::Minus => Ok((builder.ins().ineg(values[0]), first_type)),
                    Primitive::Times => Ok((builder.ins().imul(values[0], values[1]), first_type)),
                    Primitive::Divide if first_type.is_signed() => {
                        Ok((builder.ins().sdiv(values[0], values[1]), first_type))
                    }
                    Primitive::Divide => Ok((builder.ins().udiv(values[0], values[1]), first_type)),
                }
            }

            Expression::Block(_, _, mut exprs) => match exprs.pop() {
                None => Ok((builder.ins().iconst(types::I8, 0), ConstantType::I8)),
                Some(last) => {
                    for inner in exprs {
                        // we can ignore all of these return values and such, because we
                        // don't actually use them anywhere
                        self.compile_expression(inner, variables, builder)?;
                    }
                    // instead, we just return the last one
                    self.compile_expression(last, variables, builder)
                }
            },

            Expression::Print(ann, var) => {
                // Get the output buffer (or null) from our general compilation context.
                let buffer_ptr = self.output_buffer_ptr();
                let buffer_ptr = builder.ins().iconst(types::I64, buffer_ptr as i64);

                // Get a reference to the string we want to print.
                let string_data_id = self
                    .defined_strings
                    .get(var.as_ref())
                    .ok_or_else(|| BackendError::UnknownString(var.clone()))?;
                let local_name_ref = self
                    .module
                    .declare_data_in_func(*string_data_id, builder.func);
                let name_ptr = builder.ins().symbol_value(types::I64, local_name_ref);

                // Look up the value for the variable. Because this might be a
                // global variable (and that requires special logic), we just turn
                // this into an `Expression` and re-use the logic in that implementation.
                let fake_ref = ValueOrRef::Ref(ann, Type::Primitive(PrimitiveType::U8), var);
                let (val, vtype) = self.compile_value_or_ref(fake_ref, variables, builder)?;

                let vtype_repr = builder.ins().iconst(types::I64, vtype as i64);

                let casted_val = match vtype {
                    ConstantType::U64 | ConstantType::I64 => val,
                    ConstantType::I8 | ConstantType::I16 | ConstantType::I32 => {
                        builder.ins().sextend(types::I64, val)
                    }
                    ConstantType::U8 | ConstantType::U16 | ConstantType::U32 => {
                        builder.ins().uextend(types::I64, val)
                    }
                };

                // Finally, we can generate the call to print.
                let print_func_ref = self.runtime_functions.include_runtime_function(
                    "print",
                    &mut self.module,
                    builder.func,
                )?;
                builder.ins().call(
                    print_func_ref,
                    &[buffer_ptr, name_ptr, vtype_repr, casted_val],
                );
                Ok((builder.ins().iconst(types::I8, 0), ConstantType::I8))
            }

            Expression::Bind(_, _, _, _) => unimplemented!(),
        }
    }

    /// Compile a value or reference into Cranelift, returning the Cranelift Value for
    /// the expression and its type.
    fn compile_value_or_ref(
        &self,
        valref: ValueOrRef<Type>,
        variables: &ScopedMap<Variable, ReferenceBuilder>,
        builder: &mut FunctionBuilder,
    ) -> Result<(entities::Value, ConstantType), BackendError> {
        match valref {
            ValueOrRef::Value(_, _, val) => match val {
                Value::I8(_, v) => {
                    Ok((builder.ins().iconst(types::I8, v as i64), ConstantType::I8))
                }
                Value::I16(_, v) => Ok((
                    builder.ins().iconst(types::I16, v as i64),
                    ConstantType::I16,
                )),
                Value::I32(_, v) => Ok((
                    builder.ins().iconst(types::I32, v as i64),
                    ConstantType::I32,
                )),
                Value::I64(_, v) => Ok((builder.ins().iconst(types::I64, v), ConstantType::I64)),
                Value::U8(_, v) => {
                    Ok((builder.ins().iconst(types::I8, v as i64), ConstantType::U8))
                }
                Value::U16(_, v) => Ok((
                    builder.ins().iconst(types::I16, v as i64),
                    ConstantType::U16,
                )),
                Value::U32(_, v) => Ok((
                    builder.ins().iconst(types::I32, v as i64),
                    ConstantType::U32,
                )),
                Value::U64(_, v) => Ok((
                    builder.ins().iconst(types::I64, v as i64),
                    ConstantType::U64,
                )),
            },
            ValueOrRef::Ref(_, _, name) => match variables.get(&name) {
                None => Err(BackendError::VariableLookupFailure(name)),
                Some(x) => Ok(x.refer_to(builder)),
            },
        }
    }

    // Compiling a function is just compiling each of the statements in order.
    // At the moment, we do the pattern match for statements here, and then
    // directly compile the statements. If/when we add more statement forms,
    // this is likely to become more cumbersome, and we'll want to separate
    // these off. But for now, given the amount of tables we keep around to track
    // state, it's easier to just include them.
    //        for item in program.items.drain(..) {
    //            match item {
    //                TopLevel::Function(_, _, _) => unimplemented!(),
    //
    //                // Print statements are fairly easy to compile: we just lookup the
    //                // output buffer, the address of the string to print, and the value
    //                // of whatever variable we're printing. Then we just call print.
    //                TopLevel::Statement(Statement::Print(ann, t, var)) => {
    //                    // Get the output buffer (or null) from our general compilation context.
    //                    let buffer_ptr = self.output_buffer_ptr();
    //                    let buffer_ptr = builder.ins().iconst(types::I64, buffer_ptr as i64);
    //
    //                    // Get a reference to the string we want to print.
    //                    let local_name_ref = string_table.get(&var).unwrap();
    //                    let name_ptr = builder.ins().symbol_value(types::I64, *local_name_ref);
    //
    //                    // Look up the value for the variable. Because this might be a
    //                    // global variable (and that requires special logic), we just turn
    //                    // this into an `Expression` and re-use the logic in that implementation.
    //                    let (val, vtype) = ValueOrRef::Ref(ann, t, var).into_crane(
    //                        &mut builder,
    //                        &variable_table,
    //                        &pre_defined_symbols,
    //                    )?;
    //
    //                    let vtype_repr = builder.ins().iconst(types::I64, vtype as i64);
    //
    //                    let casted_val = match vtype {
    //                        ConstantType::U64 | ConstantType::I64 => val,
    //                        ConstantType::I8 | ConstantType::I16 | ConstantType::I32 => {
    //                            builder.ins().sextend(types::I64, val)
    //                        }
    //                        ConstantType::U8 | ConstantType::U16 | ConstantType::U32 => {
    //                            builder.ins().uextend(types::I64, val)
    //                        }
    //                    };
    //
    //                    // Finally, we can generate the call to print.
    //                    builder.ins().call(
    //                        print_func_ref,
    //                        &[buffer_ptr, name_ptr, vtype_repr, casted_val],
    //                    );
    //                }
    //
    //                // Variable binding is a little more con
    //                TopLevel::Statement(Statement::Binding(_, var_name, _, value)) => {
    //                    // Kick off to the `Expression` implementation to see what value we're going
    //                    // to bind to this variable.
    //                    let (val, etype) =
    //                        value.into_crane(&mut builder, &variable_table, &pre_defined_symbols)?;
    //
    //                    // Now the question is: is this a local variable, or a global one?
    //                    if let Some((global_id, ctype)) = pre_defined_symbols.get(var_name.as_str()) {
    //                        // It's a global variable! In this case, we assume that someone has already
    //                        // dedicated some space in memory to store this value. We look this location
    //                        // up, and then tell Cranelift to store the value there.
    //                        assert_eq!(etype, *ctype);
    //                        let val_ptr = builder
    //                            .ins()
    //                            .symbol_value(ir::Type::from(*ctype), *global_id);
    //                        builder.ins().store(MemFlags::new(), val, val_ptr, 0);
    //                    } else {
    //                        // It's a local variable! In this case, we need to allocate a new Cranelift
    //                        // `Variable` for this variable, which we do using our `next_var_num` counter.
    //                        // (While we're doing this, we also increment `next_var_num`, so that we get
    //                        // a fresh `Variable` next time. This is one of those very narrow cases in which
    //                        // I wish Rust had an increment expression.)
    //                        let var = Variable::new(next_var_num);
    //                        next_var_num += 1;
    //
    //                        // We can add the variable directly to our local variable map; it's `Copy`.
    //                        variable_table.insert(var_name, (var, etype));
    //
    //                        // Now we tell Cranelift about our new variable!
    //                        builder.declare_var(var, ir::Type::from(etype));
    //                        builder.def_var(var, val);
    //                    }
    //                }
    //            }
    //        }

    // Build the string table for use in referencing strings later.
    //
    // This function is slightly smart, in that it only puts strings in the table that
    // are used by the `Program`. (Thanks to `Progam::strings()`!) If the strings have
    // been declared globally, via `Backend::define_string()`, we will re-use that data.
    // Otherwise, this will define the string for you.
    //    fn build_string_table(
    //        &mut self,
    //        func: &mut Function,
    //        program: &Expression<Type>,
    //    ) -> 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 Expression {
//    fn into_crane(
//        self,
//        builder: &mut FunctionBuilder,
//        local_variables: &HashMap<ArcIntern<String>, (Variable, ConstantType)>,
//        global_variables: &HashMap<String, (GlobalValue, ConstantType)>,
//    ) -> Result<(entities::Value, ConstantType), BackendError> {
//        match self {
//            Expression::Atomic(x) => x.into_crane(builder, local_variables, global_variables),
//
//            Expression::Cast(_, target_type, expr) => {
//                let (val, val_type) =
//                    expr.into_crane(builder, local_variables, global_variables)?;
//
//                match (val_type, &target_type) {
//                }
//            }
//
//            Expression::Primitive(_, _, prim, mut vals) => {
//            }
//        }
//    }
//}
//
//// Just to avoid duplication, this just leverages the `From<ValueOrRef>` trait implementation
//// for `ValueOrRef` to compile this via the `Expression` logic, above.
//impl ValueOrRef {
//    fn into_crane(
//        self,
//        builder: &mut FunctionBuilder,
//        local_variables: &HashMap<ArcIntern<String>, (Variable, ConstantType)>,
//        global_variables: &HashMap<String, (GlobalValue, ConstantType)>,
//    ) -> Result<(entities::Value, ConstantType), BackendError> {
//        match self {
//            // Values are pretty straightforward to compile, mostly because we only
//            // have one type of variable, and it's an integer type.
//            ValueOrRef::Value(_, _, val) => match val {
//                Value::I8(_, v) => {
//                    Ok((builder.ins().iconst(types::I8, v as i64), ConstantType::I8))
//                }
//                Value::I16(_, v) => Ok((
//                    builder.ins().iconst(types::I16, v as i64),
//                    ConstantType::I16,
//                )),
//                Value::I32(_, v) => Ok((
//                    builder.ins().iconst(types::I32, v as i64),
//                    ConstantType::I32,
//                )),
//                Value::I64(_, v) => Ok((builder.ins().iconst(types::I64, v), ConstantType::I64)),
//                Value::U8(_, v) => {
//                    Ok((builder.ins().iconst(types::I8, v as i64), ConstantType::U8))
//                }
//                Value::U16(_, v) => Ok((
//                    builder.ins().iconst(types::I16, v as i64),
//                    ConstantType::U16,
//                )),
//                Value::U32(_, v) => Ok((
//                    builder.ins().iconst(types::I32, v as i64),
//                    ConstantType::U32,
//                )),
//                Value::U64(_, v) => Ok((
//                    builder.ins().iconst(types::I64, v as i64),
//                    ConstantType::U64,
//                )),
//            },
//
//            ValueOrRef::Ref(_, _, name) => {
//                // first we see if this is a local variable (which is nicer, from an
//                // optimization point of view.)
//                if let Some((local_var, etype)) = local_variables.get(&name) {
//                    return Ok((builder.use_var(*local_var), *etype));
//                }
//
//                // then we check to see if this is a global reference, which requires us to
//                // first lookup where the value is stored, and then load it.
//                if let Some((global_var, etype)) = global_variables.get(name.as_ref()) {
//                    let cranelift_type = ir::Type::from(*etype);
//                    let val_ptr = builder.ins().symbol_value(cranelift_type, *global_var);
//                    return Ok((
//                        builder
//                            .ins()
//                            .load(cranelift_type, MemFlags::new(), val_ptr, 0),
//                        *etype,
//                    ));
//                }
//
//                // this should never happen, because we should have made sure that there are
//                // no unbound variables a long time before this. but still ...
//                Err(BackendError::VariableLookupFailure(name))
//            }
//        }
//    }
//}
//