use std::cell::SyncUnsafeCell;

use jni::{objects::JByteBuffer, JNIEnv};
use unicorn_engine::{uc_error, Permission, Unicorn};
use zerocopy::{FromBytes, FromZeros, Immutable, IntoBytes, KnownLayout};

/// A simple replacement for [std::cell::Cell] in Sync contexts which explicitly does not try to synchronize the contents.
/// This is used for contexts where unsynchronized access is explicitly fine.
/// ## Dropping:
/// This type does not drop the value stored inside it if you write a new value to it. Generally the intention of the type is to store primitives and slices,
/// so this shouldn't be a problem for anyone, but **be warned**.
#[derive(Default, FromBytes, IntoBytes, KnownLayout)]
#[repr(transparent)]
pub struct UnsyncCell<T: ?Sized>(SyncUnsafeCell<T>);

impl<T> UnsyncCell<T> {
  pub fn as_mut_ptr(&self) -> *mut T {
    self.0.get()
  }
}

impl<T: Sized + Copy> UnsyncCell<T> {
  pub fn new(value: T) -> Self {
    Self(SyncUnsafeCell::new(value))
  }

  pub fn write(&self, value: T) {
    unsafe { self.as_mut_ptr().write(value) };
  }

  pub fn read(&self) -> T {
    unsafe { self.as_mut_ptr().read() }
  }
}

impl<T: Sized + Copy + IntoBytes> UnsyncCell<T> {
  // N must be the same as the size of U
  pub fn read_into<U: FromBytes>(&self, offset: usize) -> U
  where
    [u8; size_of::<U>()]:,
  {
    let mut buffer = [0; size_of::<U>()];
    let end = (offset + size_of::<U>()).min(size_of::<T>());
    assert!(end - offset <= size_of::<U>());
    // Safety: both pointers are valid for 0..(end - offset)
    unsafe { (self.as_mut_ptr() as *const u8).copy_to(buffer.as_mut_ptr(), end - offset) };
    U::read_from_bytes(&buffer).expect("N was not the same as the size of U")
  }
}

impl<T: Sized + Copy + FromBytes> UnsyncCell<T> {
  pub fn write_from<U: IntoBytes + Immutable>(&self, offset: usize, value: &U) {
    let end = (offset + size_of::<U>()).min(size_of::<T>());
    assert!(end - offset <= size_of::<U>());
    unsafe { (self.as_mut_ptr() as *mut u8).copy_from(value.as_bytes().as_ptr(), end - offset) };
  }
}

#[macro_export]
macro_rules! unsync_read {
  ($value: expr, $T: ty, $field: ident) => {
    $crate::unsync_cell::UnsyncCell::<$T>::read_into(&($value), std::mem::offset_of!($T, $field))
  };
}

impl<T: Sized + Copy> UnsyncCell<[T]> {
  pub fn new_zeroed_box(size: usize) -> Box<Self> {
    let mut memory = Box::new_uninit_slice(size);
    memory.zero();

    // Safety: UnsafeCell transparently wraps the slice, meaning the two types are semantically identical
    unsafe { std::mem::transmute::<Box<[T]>, Box<Self>>(memory.assume_init()) }
  }

  fn as_mut_ptr(&self) -> *mut [T] {
    self.0.get()
  }

  pub fn len(&self) -> usize {
    self.as_mut_ptr().len()
  }

  pub fn update_from_slice(&self, bytes: &[u8]) -> Result<(), usize> {
    if self.len() < bytes.len() {
      return Err(bytes.len() - self.len());
    }

    // Safety: our slice is valid for bytes.len(), both are u8 pointers and therefore aligned
    unsafe {
      (self.as_mut_ptr() as *mut u8).copy_from(bytes.as_ptr(), bytes.len());
    };

    Ok(())
  }

  pub fn update_from_jni(&self, env: &mut JNIEnv, byte_buffer: &JByteBuffer) -> Result<(), usize> {
    let size = env.get_direct_buffer_capacity(&byte_buffer).expect("failed to get byte buffer size");
    if self.len() < size {
      return Err(size - self.len());
    }
    let address = env.get_direct_buffer_address(&byte_buffer).expect("failed to get byte buffer address");

    // Safety: our slice is valid for bytes.len(), jni gives us a valid pointer for bytes.len(), both are u8 pointers and therefore aligned
    unsafe {
      (self.as_mut_ptr() as *mut u8).copy_from(address, size);
    };

    Ok(())
  }

  pub fn map<D>(
    &self,
    vcpu: &mut Unicorn<D>,
    address: u64,
    size: Option<usize>,
    permissions: Permission,
  ) -> Result<(), uc_error> {
    let size = size.unwrap_or(self.len());
    if size > self.len() {
      return Err(uc_error::ARG);
    }
    println!("mapping {permissions:?} at {address:08X} for {size:08X}");
    unsafe { vcpu.mem_map_ptr(address, size, permissions, self.as_mut_ptr() as *mut _) }
  }
}

impl<T: Clone + Copy> Clone for UnsyncCell<T> {
  fn clone(&self) -> Self {
    Self(SyncUnsafeCell::new(unsafe { self.0.get().read() }))
  }
}

impl<T: Sized> From<T> for UnsyncCell<T> {
  fn from(value: T) -> Self {
    Self(value.into())
  }
}