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::collections::HashMap;
use std::ffi::CStr;
use target_lexicon::Triple;
use thiserror::Error;

pub struct RuntimeFunctions {
    builtin_functions: HashMap<String, FuncId>,
    _referenced_functions: Vec<String>,
}

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

extern "C" fn runtime_print(name: *const i8, value: u64) {
    let cstr = unsafe { CStr::from_ptr(name) };
    let reconstituted = cstr.to_string_lossy();
    println!("{} = {}", reconstituted, value);
}

impl RuntimeFunctions {
    pub fn new<M: Module>(platform: &Triple, module: &mut M) -> ModuleResult<RuntimeFunctions> {
        let mut builtin_functions = HashMap::new();
        let _referenced_functions = Vec::new();

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

        let print_id = module.declare_function(
            "print",
            Linkage::Import,
            &Signature {
                params: vec![string_param, int64_param],
                returns: vec![],
                call_conv: CallConv::triple_default(platform),
            },
        )?;
        builtin_functions.insert("print".to_string(), print_id);

        Ok(RuntimeFunctions {
            builtin_functions,
            _referenced_functions,
        })
    }

    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)),
        }
    }

    pub fn register_jit_implementations(builder: &mut JITBuilder) {
        builder.symbol("print", runtime_print as *const u8);
    }
}