use cranelift_codegen::ir::{types, AbiParam, FuncRef, Function, Signature};
use cranelift_codegen::isa::CallConv;
use cranelift_jit::JITBuilder;
use cranelift_module::{FuncId, Linkage, Module, ModuleResult};
use std::alloc::Layout;
use std::collections::HashMap;
use std::ffi::CStr;
use std::fmt::Write;
use target_lexicon::Triple;
use thiserror::Error;

use crate::syntax::ConstantType;

/// An object for querying / using functions built into the runtime.
///
/// Right now, this is a quite a bit of boilerplate for very nebulous
/// value. However, as the number of built-in functions gets large, it's
/// nice to have a single point to register and query them, so here we
/// go.
pub struct RuntimeFunctions {
    builtin_functions: HashMap<String, FuncId>,
}

#[derive(Debug, Error, PartialEq)]
pub enum RuntimeFunctionError {
    #[error("Could not find runtime function named '{0}'")]
    CannotFindRuntimeFunction(String),
}

impl RuntimeFunctions {
    /// Generate a new runtime function table for the given platform, and
    /// declare them within the provided Cranelift module.
    ///
    /// Note that this is very conservative: it assumes that your module
    /// will want to use every runtime function. Unless the Cranelift object
    /// builder is smart, this might inject a bunch of references (and thus
    /// linker requirements) that aren't actually needed by your program.
    ///
    /// Then again, right now there's exactly one runtime function, so ...
    /// not a big deal.
    pub fn new<M: Module>(platform: &Triple, module: &mut M) -> ModuleResult<RuntimeFunctions> {
        let mut builtin_functions = HashMap::new();

        let string_param = AbiParam::new(types::I64);
        let int64_param = AbiParam::new(types::I64);

        // declare print for Cranelift; it's something we're going to import
        // into the current module (it's compiled separately), and takes two
        // strings and an integer. (Which ... turn out to all be the same
        // underlying type, which is weird but the way it is.)
        let print_id = module.declare_function(
            "print",
            Linkage::Import,
            &Signature {
                params: vec![string_param, string_param, int64_param, int64_param],
                returns: vec![],
                call_conv: CallConv::triple_default(platform),
            },
        )?;

        // Toss this function in our internal dictionary, as well.
        builtin_functions.insert("print".to_string(), print_id);

        Ok(RuntimeFunctions { builtin_functions })
    }

    /// Include the named runtime function into the current Function context.
    ///
    /// This is necessary for every runtime function reference within each
    /// function. The returned `FuncRef` can be used in `call` invocations.
    /// The only reason for this function to error is if you pass a name that
    /// the runtime isn't familiar with.
    pub fn include_runtime_function<M: Module>(
        &self,
        name: &str,
        module: &mut M,
        func: &mut Function,
    ) -> Result<FuncRef, RuntimeFunctionError> {
        match self.builtin_functions.get(name) {
            None => Err(RuntimeFunctionError::CannotFindRuntimeFunction(
                name.to_string(),
            )),
            Some(func_id) => Ok(module.declare_func_in_func(*func_id, func)),
        }
    }

    /// Register live, local versions of the runtime functions into the JIT.
    ///
    /// Note that these implementations are *not* the same as the ones defined
    /// in `CARGO_MANIFEST_DIR/runtime/`, for ... reasons. It might be a good
    /// change, in the future, to find a way to unify these implementations into
    /// one; both to reduce the chance that they deviate, and to reduce overall
    /// maintenance burden.
    pub fn register_jit_implementations(builder: &mut JITBuilder) {
        let allocation_pointer = unsafe {
            std::alloc::alloc_zeroed(
                Layout::from_size_align(1024 * 1024, 1024 * 1024)
                    .expect("reasonable layout is reasonable"),
            )
        };
        let allocation_pointer_pointer = unsafe {
            let res = std::alloc::alloc(Layout::for_value(&allocation_pointer)) as *mut *mut u8;
            *res = allocation_pointer;
            res as *const u8
        };
        builder.symbol("print", runtime_print as *const u8);
        builder.symbol("__global_allocation_pointer__", allocation_pointer_pointer);
    }
}

// Print! This implementation is used in the JIT compiler, to actually print data. We
// use the `output_buffer` argument as an aid for testing; if it's non-NULL, it's a string
// we extend with the output, so that multiple JIT'd `Program`s can run concurrently
// without stomping over each other's output. If `output_buffer` is NULL, we just print
// to stdout.
extern "C" fn runtime_print(
    output_buffer: *mut String,
    name: *const i8,
    vtype_repr: i64,
    value: i64,
) {
    let cstr = unsafe { CStr::from_ptr(name) };
    let reconstituted = cstr.to_string_lossy();

    let output = match vtype_repr.try_into() {
        Ok(ConstantType::Void) => format!("{} = <void>", reconstituted),
        Ok(ConstantType::I8) => format!("{} = {}i8", reconstituted, value as i8),
        Ok(ConstantType::I16) => format!("{} = {}i16", reconstituted, value as i16),
        Ok(ConstantType::I32) => format!("{} = {}i32", reconstituted, value as i32),
        Ok(ConstantType::I64) => format!("{} = {}i64", reconstituted, value),
        Ok(ConstantType::U8) => format!("{} = {}u8", reconstituted, value as u8),
        Ok(ConstantType::U16) => format!("{} = {}u16", reconstituted, value as u16),
        Ok(ConstantType::U32) => format!("{} = {}u32", reconstituted, value as u32),
        Ok(ConstantType::U64) => format!("{} = {}u64", reconstituted, value as u64),
        Err(_) => format!("{} = {}<unknown type {}>", reconstituted, value, vtype_repr),
    };

    if let Some(output_buffer) = unsafe { output_buffer.as_mut() } {
        writeln!(output_buffer, "{}", output).unwrap();
    } else {
        println!("{}", output);
    }
}