use crate::alloc;
use alloc::Alloc;
use libc::c_char;
use log::*;
use solana_rbpf::{
    ebpf::{HelperContext, MM_HEAP_START},
    memory_region::{translate_addr, MemoryRegion},
    EbpfVm,
};
use std::alloc::Layout;
use std::ffi::CStr;
use std::io::{Error, ErrorKind};
use std::mem;
use std::slice::from_raw_parts;
use std::str::from_utf8;

/// Program heap allocators are intended to allocate/free from a given
/// chunk of memory.  The specific allocator implementation is
/// selectable at build-time.
/// Only one allocator is currently supported

/// Simple bump allocator, never frees
use crate::allocator_bump::BPFAllocator;

/// Default program heap size, allocators
/// are expected to enforce this
const DEFAULT_HEAP_SIZE: usize = 32 * 1024;

pub fn register_helpers(vm: &mut EbpfVm) -> Result<(MemoryRegion), Error> {
    vm.register_helper_ex("abort", helper_abort, None)?;
    vm.register_helper_ex("sol_panic", helper_sol_panic, None)?;
    vm.register_helper_ex("sol_panic_", helper_sol_panic, None)?;
    vm.register_helper_ex("sol_log", helper_sol_log, None)?;
    vm.register_helper_ex("sol_log_", helper_sol_log, None)?;
    vm.register_helper_ex("sol_log_64", helper_sol_log_u64, None)?;
    vm.register_helper_ex("sol_log_64_", helper_sol_log_u64, None)?;

    let heap = vec![0_u8; DEFAULT_HEAP_SIZE];
    let heap_region = MemoryRegion::new_from_slice(&heap, MM_HEAP_START);
    let context = Box::new(BPFAllocator::new(heap, MM_HEAP_START));
    vm.register_helper_ex("sol_alloc_free_", helper_sol_alloc_free, Some(context))?;

    Ok(heap_region)
}

/// Abort helper functions, called when the BPF program calls `abort()`
/// The verify function returns an error which will cause the BPF program
/// to be halted immediately
pub fn helper_abort(
    _arg1: u64,
    _arg2: u64,
    _arg3: u64,
    _arg4: u64,
    _arg5: u64,
    _context: &mut HelperContext,
    _ro_regions: &[MemoryRegion],
    _rw_regions: &[MemoryRegion],
) -> Result<(u64), Error> {
    Err(Error::new(
        ErrorKind::Other,
        "Error: BPF program called abort()!",
    ))
}

/// Panic helper functions, called when the BPF program calls 'sol_panic_()`
/// The verify function returns an error which will cause the BPF program
/// to be halted immediately
pub fn helper_sol_panic(
    file: u64,
    len: u64,
    line: u64,
    column: u64,
    _arg5: u64,
    _context: &mut HelperContext,
    ro_regions: &[MemoryRegion],
    _rw_regions: &[MemoryRegion],
) -> Result<(u64), Error> {
    if let Ok(host_addr) = translate_addr(file, len as usize, "Load", 0, ro_regions) {
        let c_buf: *const c_char = host_addr as *const c_char;
        let c_str: &CStr = unsafe { CStr::from_ptr(c_buf) };
        if let Ok(slice) = c_str.to_str() {
            return Err(Error::new(
                ErrorKind::Other,
                format!(
                    "Error: BPF program Panicked at {}, {}:{}",
                    slice, line, column
                ),
            ));
        }
    }
    Err(Error::new(ErrorKind::Other, "Error: BPF program Panicked"))
}

pub fn helper_sol_log(
    addr: u64,
    len: u64,
    _arg3: u64,
    _arg4: u64,
    _arg5: u64,
    _context: &mut HelperContext,
    ro_regions: &[MemoryRegion],
    _rw_regions: &[MemoryRegion],
) -> Result<(u64), Error> {
    let host_addr = translate_addr(addr, len as usize, "Load", 0, ro_regions)?;
    let c_buf: *const c_char = host_addr as *const c_char;
    unsafe {
        for i in 0..len {
            let c = std::ptr::read(c_buf.offset(i as isize));
            if i == len - 1 || c == 0 {
                let message =
                    from_utf8(from_raw_parts(host_addr as *const u8, len as usize)).unwrap();
                info!("info!: {}", message);
                return Ok(0);
            }
        }
    }
    Err(Error::new(
        ErrorKind::Other,
        "Error: Unterminated string logged",
    ))
}

pub fn helper_sol_log_u64(
    arg1: u64,
    arg2: u64,
    arg3: u64,
    arg4: u64,
    arg5: u64,
    _context: &mut HelperContext,
    _ro_regions: &[MemoryRegion],
    _rw_regions: &[MemoryRegion],
) -> Result<(u64), Error> {
    info!(
        "info!: {:#x}, {:#x}, {:#x}, {:#x}, {:#x}",
        arg1, arg2, arg3, arg4, arg5
    );
    Ok(0)
}

/// Dynamic memory allocation helper called when the BPF program calls
/// `sol_alloc_free_()`.  The allocator is expected to allocate/free
/// from/to a given chunk of memory and enforce size restrictions.  The
/// memory chunk is given to the allocator during allocator creation and
/// information about that memory (start address and size) is passed
/// to the VM to use for enforcement.
pub fn helper_sol_alloc_free(
    size: u64,
    free_addr: u64,
    _arg3: u64,
    _arg4: u64,
    _arg5: u64,
    context: &mut HelperContext,
    _ro_regions: &[MemoryRegion],
    _rw_regions: &[MemoryRegion],
) -> Result<(u64), Error> {
    if let Some(context) = context {
        if let Some(allocator) = context.downcast_mut::<BPFAllocator>() {
            return {
                let layout = Layout::from_size_align(size as usize, mem::align_of::<u8>()).unwrap();
                if free_addr == 0 {
                    match allocator.alloc(layout) {
                        Ok(addr) => Ok(addr as u64),
                        Err(_) => Ok(0),
                    }
                } else {
                    allocator.dealloc(free_addr, layout);
                    Ok(0)
                }
            };
        };
    }
    panic!("Failed to get alloc_free context");
}