use crate::backend::error::BackendError;
use crate::backend::Backend;
use crate::eval::PrimitiveType;
use crate::ir::{Expression, Primitive, Program, TopLevel, Type, Value, ValueOrRef, Variable};
use crate::syntax::{ConstantType, Location};
use cranelift_codegen::ir::{
    self, entities, types, AbiParam, Function, GlobalValue, InstBuilder, MemFlags, Signature,
    UserFuncName,
};
use cranelift_codegen::isa::CallConv;
use cranelift_codegen::Context;
use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
use cranelift_module::{DataDescription, FuncId, Linkage, Module};
use internment::ArcIntern;
use std::collections::{hash_map, HashMap};

/// 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.
pub enum ReferenceBuilder {
    Global(ConstantType, GlobalValue),
    Local(ConstantType, cranelift_frontend::Variable),
    Argument(ConstantType, entities::Value),
}

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![];
        let mut variables = HashMap::new();

        for (top_level_name, top_level_type) in program.get_top_level_variables() {
            match top_level_type {
                Type::Function(argument_types, return_type) => {
                    let func_id = self.declare_function(
                        top_level_name.as_str(),
                        Linkage::Export,
                        argument_types,
                        *return_type,
                    )?;
                    self.defined_functions.insert(top_level_name, func_id);
                }

                Type::Primitive(pt) => {
                    let data_id = self.module.declare_data(
                        top_level_name.as_str(),
                        Linkage::Export,
                        true,
                        false,
                    )?;
                    println!("defining {} with primitive type {}", top_level_name, pt);
                    self.module.define_data(data_id, &pt.blank_data())?;
                    self.defined_symbols
                        .insert(top_level_name, (data_id, pt.into()));
                }
            }
        }

        let void = Type::Primitive(PrimitiveType::Void);
        let main_func_id =
            self.declare_function(function_name, Linkage::Export, vec![], void.clone())?;
        self.defined_functions
            .insert(ArcIntern::new(function_name.to_string()), main_func_id);

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

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

        self.compile_function(
            &mut variables,
            function_name,
            &[],
            void.clone(),
            Expression::Block(Location::manufactured(), void, generated_body),
        )
    }

    fn declare_function(
        &mut self,
        name: &str,
        linkage: Linkage,
        argument_types: Vec<Type>,
        return_type: Type,
    ) -> Result<FuncId, cranelift_module::ModuleError> {
        println!("Declaring {:?} function {}", linkage, name);
        let basic_signature = Signature {
            params: argument_types
                .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.
        self.module
            .declare_function(name, linkage, &basic_signature)
    }

    /// Compile the given function.
    pub fn compile_function(
        &mut self,
        variables: &mut HashMap<Variable, ReferenceBuilder>,
        function_name: &str,
        arguments: &[(Variable, Type)],
        return_type: Type,
        body: Expression<Type>,
    ) -> Result<FuncId, BackendError> {
        println!("Compiling function {}", function_name);
        {
            use pretty::{DocAllocator, Pretty};
            let allocator = pretty::BoxAllocator;
            allocator
                .text("Function body:")
                .append(allocator.hardline())
                .append(body.pretty(&allocator))
                .append(allocator.hardline())
                .1
                .render_colored(
                    70,
                    pretty::termcolor::StandardStream::stdout(pretty::termcolor::ColorChoice::Auto),
                )
                .expect("rendering works");
        }
        // reset the next variable counter. this value shouldn't matter; hopefully
        // we won't be using close to 2^32 variables!
        self.reset_local_variable_tracker();

        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 interned_name = ArcIntern::new(function_name.to_string());
        let func_id = match self.defined_functions.entry(interned_name) {
            hash_map::Entry::Occupied(entry) => *entry.get(),
            hash_map::Entry::Vacant(vac) => {
                let func_id = self.module.declare_function(
                    function_name,
                    Linkage::Export,
                    &basic_signature,
                )?;
                vac.insert(func_id);
                func_id
            }
        };

        // 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);

        // 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);
            variables.insert(name.clone(), ReferenceBuilder::Global(*ty, local_data));
        }

        // 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();
        // add the block parameters, which should be the function parameters
        for (name, ty) in arguments.iter() {
            let constant_type = ty
                .try_into()
                .map_err(|_| BackendError::NoFunctionArguments {
                    function_name: function_name.to_string(),
                    argument_name: name.to_string(),
                })?;
            let value = builder.append_block_param(main_block, ir::Type::from(constant_type));
            variables.insert(
                name.clone(),
                ReferenceBuilder::Argument(constant_type, value),
            );
        }

        builder.switch_to_block(main_block);

        let (value, _) = self.compile_expression(body, 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).
        if return_type == Type::Primitive(PrimitiveType::Void) {
            builder.ins().return_(&[]);
        } else {
            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 HashMap<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)),

                    (ConstantType::Void, Type::Primitive(PrimitiveType::Void)) => {
                        Ok((val, val_type))
                    }

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

                    _ => 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::I64, 0), ConstantType::Void)),
                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(_, 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 var_name = match var {
                    ValueOrRef::Ref(_, _, ref name) => name.as_ref(),
                    ValueOrRef::Value(_, _, _) => "<unknown>",
                };
                let string_data_id = self.string_reference(var_name)?;
                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);

                let (val, vtype) = self.compile_value_or_ref(var, variables, builder)?;

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

                let casted_val = match vtype {
                    ConstantType::U64 | ConstantType::I64 | ConstantType::Void => 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::I64, 0), ConstantType::Void))
            }

            Expression::Bind(_, name, _, expr) => {
                let (value, value_type) = self.compile_expression(*expr, variables, builder)?;
                let ir_type = ir::Type::from(value_type);
                let variable = self.generate_local();

                match variables.get(&name) {
                    Some(ReferenceBuilder::Global(_, global_value)) => {
                        let pointer = self.module.target_config().pointer_type();
                        let pointer_to = builder.ins().symbol_value(pointer, *global_value);
                        println!(
                            "STORE {}: cranelift_type {} origin type {:?}",
                            name, value, value_type
                        );
                        builder.ins().store(MemFlags::new(), value, pointer_to, 0);
                        Ok((builder.ins().iconst(types::I64, 0), ConstantType::Void))
                    }

                    Some(ReferenceBuilder::Argument(_, _)) => {
                        panic!("Attempt to mutate an argument {}", name)
                    }

                    Some(ReferenceBuilder::Local(_, _)) => {
                        panic!("Attempt to mutate local {}", name);
                    }

                    None => {
                        builder.declare_var(variable, ir_type);
                        builder.def_var(variable, value);
                        variables.insert(name, ReferenceBuilder::Local(value_type, variable));
                        Ok((builder.ins().iconst(types::I64, 0), ConstantType::Void))
                    }
                }
            }

            Expression::Call(_, _, function, args) => {
                let (arguments, _argument_types): (Vec<_>, Vec<_>) = args
                    .into_iter()
                    .map(|x| self.compile_value_or_ref(x, variables, builder))
                    .collect::<Result<Vec<(_, _)>, BackendError>>()?
                    .into_iter()
                    .unzip();

                match *function {
                    ValueOrRef::Value(_, _, _) => {
                        panic!("Can't use a value for a function")
                    }

                    ValueOrRef::Ref(_, result_type, name) => {
                        match self.defined_functions.get(&name) {
                            None => panic!("Couldn't find function {} to call", name),
                            Some(function) => {
                                let func_ref =
                                    self.module.declare_func_in_func(*function, builder.func);
                                let call = builder.ins().call(func_ref, &arguments);
                                let results = builder.inst_results(call);

                                match results {
                                    [] => Ok((
                                        builder.ins().iconst(types::I64, 0),
                                        ConstantType::Void,
                                    )),
                                    [result] => match result_type {
                                        Type::Primitive(ct) => Ok((*result, ct.into())),
                                        Type::Function(_, rt) => match *rt {
                                            Type::Function(_, _) => {
                                                panic!("function returns a function?")
                                            }
                                            Type::Primitive(ct) => Ok((*result, ct.into())),
                                        },
                                    },
                                    _ => panic!("don't support multi-value returns yet"),
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    /// 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: &HashMap<Variable, ReferenceBuilder>,
        builder: &mut FunctionBuilder,
    ) -> Result<(entities::Value, ConstantType), BackendError> {
        println!("compile_value_or_ref {:?}", valref);
        match valref {
            ValueOrRef::Value(_, _, val) => match val {
                Value::I8(_, v) => {
                    // Cranelift does a funny thing where it checks you aren't using bits in the I64
                    // we provide above the size of the type we provide. So, in this case, we can only
                    // set the low 8 bits of the i64. This restriction creates a bit of a problem when
                    // casting direction from i8 to i64, because Rust will (helpfully) sign-extend the
                    // negative number for us. Which sets the high bits, which makes Cranelift unhappy.
                    // So first we cast the i8 as u8, to get rid of the whole concept of sign extension,
                    // and *then* we cast to i64.
                    Ok((
                        builder.ins().iconst(types::I8, v as u8 as i64),
                        ConstantType::I8,
                    ))
                }
                Value::I16(_, v) => Ok((
                    // see above note for the "... as ... as"
                    builder.ins().iconst(types::I16, v as u16 as i64),
                    ConstantType::I16,
                )),
                Value::I32(_, v) => Ok((
                    // see above note for the "... as ... as"
                    builder.ins().iconst(types::I32, v as u32 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,
                )),
                Value::Void => Ok((builder.ins().iconst(types::I64, 0i64), ConstantType::Void)),
            },
            ValueOrRef::Ref(_, _, name) => match variables.get(&name) {
                None => Err(BackendError::VariableLookupFailure(name)),
                Some(ReferenceBuilder::Global(ty, gv)) => {
                    let pointer_to = self.module.target_config().pointer_type();
                    let pointer_value = builder.ins().symbol_value(pointer_to, *gv);
                    let cranelift_type = ir::Type::from(*ty);
                    println!(
                        "READ {}: cranelift_type {} origin type {:?}",
                        name, cranelift_type, ty
                    );
                    let value =
                        builder
                            .ins()
                            .load(cranelift_type, MemFlags::new(), pointer_value, 0);

                    Ok((value, *ty))
                }

                Some(ReferenceBuilder::Argument(ctype, val)) => Ok((*val, *ctype)),

                Some(ReferenceBuilder::Local(ctype, var)) => {
                    let value = builder.use_var(*var);
                    Ok((value, *ctype))
                }
            },
        }
    }
}

impl PrimitiveType {
    fn blank_data(&self) -> DataDescription {
        let (size, alignment) = match self {
            PrimitiveType::Void => (8, 8),
            PrimitiveType::U8 => (1, 1),
            PrimitiveType::U16 => (2, 2),
            PrimitiveType::U32 => (4, 4),
            PrimitiveType::U64 => (8, 8),
            PrimitiveType::I8 => (1, 1),
            PrimitiveType::I16 => (2, 2),
            PrimitiveType::I32 => (4, 4),
            PrimitiveType::I64 => (8, 8),
        };
        let mut result = DataDescription::new();
        result.define_zeroinit(size);
        result.set_align(alignment);
        result
    }
}