use anyhow::{bail, Result};
use std::borrow::Cow;
use std::cell::UnsafeCell;
use std::fmt;
use std::mem;
use std::ops::Range;
use std::str;

pub use wiggle_macro::{async_trait, from_witx};

pub use anyhow;
pub use wiggle_macro::wasmtime_integration;

pub use bitflags;

#[cfg(feature = "wiggle_metadata")]
pub use witx;

mod error;
mod guest_type;
mod region;

pub use tracing;

pub use error::GuestError;
pub use guest_type::{GuestErrorType, GuestType, GuestTypeTransparent};
pub use region::Region;

pub mod async_trait_crate {
    pub use async_trait::*;
}

#[cfg(feature = "wasmtime")]
pub mod wasmtime_crate {
    pub use wasmtime::*;
}

/// Representation of guest memory for `wiggle`-generated trait methods.
///
/// Guest memory is represented as an array of bytes. Memories are either
/// "unshared" or "shared". Unshared means that the host has exclusive access to
/// the entire array of memory. This allows safe borrows into wasm linear
/// memory. Shared memories can be modified at any time and are represented as
/// an array of `UnsafeCell<u8>`.
///
/// This is generated by the `wiggle` bindings macros.
pub enum GuestMemory<'a> {
    Unshared(&'a mut [u8]),
    Shared(&'a [UnsafeCell<u8>]),
}

// manual impls are needed because of the `UnsafeCell` in the `Shared` branch
// but this otherwise upholds send/sync invariants.
unsafe impl Send for GuestMemory<'_> {}
unsafe impl Sync for GuestMemory<'_> {}

impl<'a> GuestMemory<'a> {
    /// Read a value from the provided pointer.
    ///
    /// This method will delegate to `T`'s implementation of `read` which will
    /// read a value from the `ptr` provided.
    ///
    /// # Errors
    ///
    /// An error is returned if `ptr` is out of bounds, misaligned, or otherwise
    /// not valid to read from.
    pub fn read<T>(&self, ptr: GuestPtr<T>) -> Result<T, GuestError>
    where
        T: GuestType,
    {
        T::read(self, ptr)
    }

    /// Writes the `val` provided to the `ptr` provided.
    ///
    /// This commit will write a `val` into a guest's linear memory. This will
    /// delegate to `T`'s implementation of `write`.
    ///
    /// # Errors
    ///
    /// An error is returned if `ptr` is out of bounds, misaligned, or otherwise
    /// not valid to read from.
    pub fn write<T>(&mut self, ptr: GuestPtr<T>, val: T) -> Result<(), GuestError>
    where
        T: GuestType,
    {
        T::write(self, ptr, val)
    }

    /// Acquires a slice or owned copy of the memory pointed to by `ptr`.
    ///
    /// This method will attempt to borrow `ptr` directly from linear memory. If
    /// memory is shared and cannot be borrowed directly then an owned copy is
    /// returned instead.
    ///
    /// # Errors
    ///
    /// An error is returned if `ptr` is out of bounds, misaligned, or otherwise
    /// not valid to read from.
    pub fn as_cow(&self, ptr: GuestPtr<[u8]>) -> Result<Cow<'_, [u8]>, GuestError> {
        match self {
            GuestMemory::Unshared(_) => match self.as_slice(ptr)? {
                Some(slice) => Ok(Cow::Borrowed(slice)),
                None => unreachable!(),
            },
            GuestMemory::Shared(_) => Ok(Cow::Owned(self.to_vec(ptr)?)),
        }
    }

    /// Same as [`GuestMemory::as_cow`] but for strings.
    ///
    /// # Errors
    ///
    /// An error is returned if `ptr` is out of bounds, misaligned, or otherwise
    /// not valid to read from.
    pub fn as_cow_str(&self, ptr: GuestPtr<str>) -> Result<Cow<'_, str>, GuestError> {
        match self.as_cow(ptr.cast::<[u8]>())? {
            Cow::Owned(bytes) => Ok(Cow::Owned(
                String::from_utf8(bytes).map_err(|e| e.utf8_error())?,
            )),
            Cow::Borrowed(bytes) => Ok(Cow::Borrowed(std::str::from_utf8(bytes)?)),
        }
    }

    /// Attempts to borrow a raw guest slice of memory pointed to by `ptr`.
    ///
    /// This method will attempt to return a raw pointer into guest memory. This
    /// can only be done for `Unshared` memories. A `Shared` memory will return
    /// `Ok(None)` here.
    ///
    /// # Errors
    ///
    /// An error is returned if `ptr` is out of bounds, misaligned, or otherwise
    /// not valid to read from.
    pub fn as_slice(&self, ptr: GuestPtr<[u8]>) -> Result<Option<&[u8]>, GuestError> {
        let range = self.validate_range::<u8>(ptr.pointer.0, ptr.pointer.1)?;
        match self {
            GuestMemory::Unshared(slice) => Ok(Some(&slice[range])),
            GuestMemory::Shared(_) => Ok(None),
        }
    }

    /// Same as [`GuestMemory::as_slice`] but for strings.
    pub fn as_str(&self, ptr: GuestPtr<str>) -> Result<Option<&str>, GuestError> {
        match self.as_slice(ptr.cast())? {
            Some(bytes) => Ok(Some(std::str::from_utf8(bytes)?)),
            None => Ok(None),
        }
    }

    /// Attempts return `ptr` as a raw slice of mutable bytes in wasm linear
    /// memory.
    ///
    /// Like [`GuestMemory::as_slice`] this only works for `Unshared` memories
    /// and will not work for `Shared` memories.
    pub fn as_slice_mut(&mut self, ptr: GuestPtr<[u8]>) -> Result<Option<&mut [u8]>, GuestError> {
        let range = self.validate_range::<u8>(ptr.pointer.0, ptr.pointer.1)?;
        match self {
            GuestMemory::Unshared(slice) => Ok(Some(&mut slice[range])),
            GuestMemory::Shared(_) => Ok(None),
        }
    }

    /// Copies the data in the guest region into a [`Vec`].
    ///
    /// This is useful when one cannot use [`GuestMemory::as_slice`], e.g., when
    /// pointing to a region of WebAssembly shared memory.
    pub fn to_vec<T>(&self, ptr: GuestPtr<[T]>) -> Result<Vec<T>, GuestError>
    where
        T: GuestTypeTransparent + Copy,
    {
        let guest = self.validate_size_align::<T>(ptr.pointer.0, ptr.pointer.1)?;
        let mut host = Vec::with_capacity(guest.len());

        // SAFETY: The `guest_slice` variable is already a valid pointer into
        // the guest's memory, and it may or may not be a pointer into shared
        // memory. We can't naively use `.to_vec(..)` which could introduce data
        // races but all that needs to happen is to copy data into our local
        // `vec` as all the data is `Copy` and transparent anyway. For this
        // purpose the `ptr::copy` function should be sufficient for copying
        // over all the data.
        //
        // TODO: audit that this use of `std::ptr::copy` is safe with shared
        // memory (https://github.com/bytecodealliance/wasmtime/issues/4203)
        unsafe {
            std::ptr::copy(guest.as_ptr().cast(), host.as_mut_ptr(), guest.len());
            host.set_len(guest.len());
        }
        Ok(host)
    }

    /// Copies the data pointed to by `slice` into this guest region.
    ///
    /// This method is a *safe* method to copy data from the host to the guest.
    /// This requires that `self` and `slice` have the same length. The pointee
    /// type `T` requires the [`GuestTypeTransparent`] trait which is an
    /// assertion that the representation on the host and on the guest is the
    /// same.
    ///
    /// # Errors
    ///
    /// Returns an error if this guest pointer is out of bounds or if the length
    /// of this guest pointer is not equal to the length of the slice provided.
    pub fn copy_from_slice<T>(&mut self, slice: &[T], ptr: GuestPtr<[T]>) -> Result<(), GuestError>
    where
        T: GuestTypeTransparent + Copy,
    {
        if usize::try_from(ptr.len())? != slice.len() {
            return Err(GuestError::SliceLengthsDiffer);
        }
        if slice.is_empty() {
            return Ok(());
        }

        let guest = self.validate_size_align::<T>(ptr.pointer.0, ptr.pointer.1)?;

        // SAFETY: in the shared memory case, we copy and accept that
        // the guest data may be concurrently modified. TODO: audit that
        // this use of `std::ptr::copy` is safe with shared memory
        // (https://github.com/bytecodealliance/wasmtime/issues/4203)
        //
        // Also note that the validity of `guest_slice` has already been
        // determined by the `as_unsafe_slice_mut` call above.
        assert_eq!(guest.len(), slice.len());
        unsafe {
            let guest: &[UnsafeCell<T>] = guest;
            let guest: *const UnsafeCell<T> = guest.as_ptr();
            let guest = guest.cast_mut().cast::<T>();
            std::ptr::copy(slice.as_ptr(), guest, slice.len());
        }
        Ok(())
    }

    /// Validates a guest-relative pointer given various attributes, and returns
    /// the corresponding host pointer.
    ///
    /// * `mem` - this is the guest memory being accessed.
    /// * `offset` - this is the guest-relative pointer, an offset from the
    ///   base.
    /// * `len` - this is the number of length, in units of `T`, to return
    ///   in the resulting slice.
    ///
    /// If the parameters are valid then this function will return a slice into
    /// `mem` for units of `T`, assuming everything is in-bounds and properly
    /// aligned. Additionally the byte-based `Region` is returned, used for borrows
    /// later on.
    fn validate_size_align<T>(&self, offset: u32, len: u32) -> Result<&[UnsafeCell<T>], GuestError>
    where
        T: GuestTypeTransparent,
    {
        let range = self.validate_range::<T>(offset, len)?;
        let cells = match self {
            GuestMemory::Unshared(s) => {
                let s: &[u8] = s;
                unsafe { &*(s as *const [u8] as *const [UnsafeCell<u8>]) }
            }
            GuestMemory::Shared(s) => s,
        };
        let memory = &cells[range.clone()];

        // ... and then align it to `T`, failing if either the head or tail slices
        // are nonzero in length. This `unsafe` here is from the standard library
        // and should be ok since the input slice is `UnsafeCell<u8>` and the output
        // slice is `UnsafeCell<T>`, meaning the only guarantee of the output is
        // that it's valid addressable memory, still unsafe to actually access.
        assert!(mem::align_of::<T>() <= T::guest_align());
        let (start, mid, end) = unsafe { memory.align_to() };
        if start.len() > 0 || end.len() > 0 {
            let region = Region {
                start: range.start as u32,
                len: range.len() as u32,
            };
            return Err(GuestError::PtrNotAligned(region, T::guest_align() as u32));
        }
        Ok(mid)
    }

    fn validate_range<T>(&self, offset: u32, len: u32) -> Result<Range<usize>, GuestError>
    where
        T: GuestTypeTransparent,
    {
        let byte_len = len
            .checked_mul(T::guest_size())
            .ok_or(GuestError::PtrOverflow)?;
        let region = Region {
            start: offset,
            len: byte_len,
        };
        let offset = usize::try_from(offset)?;
        let byte_len = usize::try_from(byte_len)?;

        let range = offset..offset + byte_len;
        let oob = match self {
            GuestMemory::Unshared(b) => b.get(range.clone()).is_none(),
            GuestMemory::Shared(b) => b.get(range.clone()).is_none(),
        };
        if oob {
            Err(GuestError::PtrOutOfBounds(region))
        } else {
            Ok(range)
        }
    }

    /// Returns whether this is a shared memory or not.
    pub fn is_shared_memory(&self) -> bool {
        match self {
            GuestMemory::Shared(_) => true,
            GuestMemory::Unshared(_) => false,
        }
    }
}

/// A *guest* pointer.
///
/// This type represents a pointer from the guest that points into host memory.
/// Internally a `GuestPtr` the offset into the memory that the pointer is
/// pointing at. At this time this is always a 32-bit offset so this is not
/// suitable for bindings where wasm has 64-bit addresses.
///
/// Presence of a [`GuestPtr`] does not imply any form of validity. Pointers can
/// be out-of-bounds, misaligned, etc. It is safe to construct a `GuestPtr` with
/// any offset at any time. Consider a `GuestPtr<T>` roughly equivalent to `*mut
/// T`.
///
/// ## Slices and Strings
///
/// Note that the type parameter does not need to implement the `Sized` trait,
/// so you can implement types such as this:
///
/// * `GuestPtr<str>` - a pointer to a guest string.
/// * `GuestPtr<[T]>` - a pointer to a guest array.
///
/// Note that generated bindings won't use these types so you'll have to
/// otherwise construct the types with `.cast()` or `.as_array()`. Unsized types
/// track both the pointer and length in guest memory.
///
/// ## Type parameter and pointee
///
/// The `T` type parameter is largely intended for more static safety in Rust as
/// well as having a better handle on what we're pointing to. A `GuestPtr<T>`,
/// however, does not necessarily literally imply a guest pointer pointing to
/// type `T`. Instead the [`GuestType`] trait is a layer of abstraction where
/// `GuestPtr<T>` may actually be a pointer to `U` in guest memory, but you can
/// construct a `T` from a `U`.
///
/// For example `GuestPtr<GuestPtr<T>>` is a valid type, but this is actually
/// more equivalent to `GuestPtr<u32>` because guest pointers are always
/// 32-bits. That being said you can create a `GuestPtr<T>` from a `u32`.
///
/// Additionally `GuestPtr<MyEnum>` will actually delegate, typically, to and
/// implementation which loads the underlying data as `GuestPtr<u8>` (or
/// similar) and then the bytes loaded are validated to fit within the
/// definition of `MyEnum` before `MyEnum` is returned.
///
/// For more information see the [`GuestMemory::read`] and
/// [`GuestMemory::write`] methods. In general though be extremely careful about
/// writing `unsafe` code when working with a `GuestPtr` if you're not using one
/// of the already-attached helper methods.
#[repr(transparent)]
pub struct GuestPtr<T: ?Sized + Pointee> {
    pointer: T::Pointer,
}

impl<T: ?Sized + Pointee> GuestPtr<T> {
    /// Creates a new `GuestPtr` from the given `mem` and `pointer` values.
    ///
    /// Note that for sized types like `u32`, `GuestPtr<T>`, etc, the `pointer`
    /// value is a `u32` offset into guest memory. For slices and strings,
    /// `pointer` is a `(u32, u32)` offset/length pair.
    pub fn new(pointer: T::Pointer) -> GuestPtr<T> {
        GuestPtr { pointer }
    }

    /// Returns the offset of this pointer in guest memory.
    ///
    /// Note that for sized types this returns a `u32`, but for slices and
    /// strings it returns a `(u32, u32)` pointer/length pair.
    pub fn offset(&self) -> T::Pointer {
        self.pointer
    }

    /// Casts this `GuestPtr` type to a different type.
    ///
    /// This is a safe method which is useful for simply reinterpreting the type
    /// parameter on this `GuestPtr`. Note that this is a safe method, where
    /// again there's no guarantees about alignment, validity, in-bounds-ness,
    /// etc of the returned pointer.
    pub fn cast<U>(&self) -> GuestPtr<U>
    where
        U: Pointee<Pointer = T::Pointer> + ?Sized,
    {
        GuestPtr::new(self.pointer)
    }

    /// Performs pointer arithmetic on this pointer, moving the pointer forward
    /// `amt` slots.
    ///
    /// This will either return the resulting pointer or `Err` if the pointer
    /// arithmetic calculation would overflow around the end of the address
    /// space.
    pub fn add(&self, amt: u32) -> Result<GuestPtr<T>, GuestError>
    where
        T: GuestType + Pointee<Pointer = u32>,
    {
        let offset = amt
            .checked_mul(T::guest_size())
            .and_then(|o| self.pointer.checked_add(o));
        let offset = match offset {
            Some(o) => o,
            None => return Err(GuestError::PtrOverflow),
        };
        Ok(GuestPtr::new(offset))
    }

    /// Returns a `GuestPtr` for an array of `T`s using this pointer as the
    /// base.
    pub fn as_array(&self, elems: u32) -> GuestPtr<[T]>
    where
        T: GuestType + Pointee<Pointer = u32>,
    {
        GuestPtr::new((self.pointer, elems))
    }
}

impl<T> GuestPtr<[T]> {
    /// For slices, specifically returns the relative pointer to the base of the
    /// array.
    ///
    /// This is similar to `<[T]>::as_ptr()`
    pub fn offset_base(&self) -> u32 {
        self.pointer.0
    }

    /// For slices, returns the length of the slice, in elements.
    pub fn len(&self) -> u32 {
        self.pointer.1
    }

    /// Returns an iterator over interior pointers.
    ///
    /// Each item is a `Result` indicating whether it overflowed past the end of
    /// the address space or not.
    pub fn iter(&self) -> impl ExactSizeIterator<Item = Result<GuestPtr<T>, GuestError>> + '_
    where
        T: GuestType,
    {
        let base = self.as_ptr();
        (0..self.len()).map(move |i| base.add(i))
    }

    /// Returns a `GuestPtr` pointing to the base of the array for the interior
    /// type `T`.
    pub fn as_ptr(&self) -> GuestPtr<T> {
        GuestPtr::new(self.offset_base())
    }

    pub fn get(&self, index: u32) -> Option<GuestPtr<T>>
    where
        T: GuestType,
    {
        if index < self.len() {
            Some(
                self.as_ptr()
                    .add(index)
                    .expect("just performed bounds check"),
            )
        } else {
            None
        }
    }

    pub fn get_range(&self, r: std::ops::Range<u32>) -> Option<GuestPtr<[T]>>
    where
        T: GuestType,
    {
        if r.end < r.start {
            return None;
        }
        let range_length = r.end - r.start;
        if r.start <= self.len() && r.end <= self.len() {
            Some(
                self.as_ptr()
                    .add(r.start)
                    .expect("just performed bounds check")
                    .as_array(range_length),
            )
        } else {
            None
        }
    }
}

impl GuestPtr<str> {
    /// For strings, returns the relative pointer to the base of the string
    /// allocation.
    pub fn offset_base(&self) -> u32 {
        self.pointer.0
    }

    /// Returns the length, in bytes, of the string.
    pub fn len(&self) -> u32 {
        self.pointer.1
    }

    /// Returns a raw pointer for the underlying slice of bytes that this
    /// pointer points to.
    pub fn as_bytes(&self) -> GuestPtr<[u8]> {
        GuestPtr::new(self.pointer)
    }
}

impl<T: ?Sized + Pointee> Clone for GuestPtr<T> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<T: ?Sized + Pointee> Copy for GuestPtr<T> {}

impl<T: ?Sized + Pointee> fmt::Debug for GuestPtr<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        T::debug(self.pointer, f)
    }
}

impl<T: ?Sized + Pointee> PartialEq for GuestPtr<T> {
    fn eq(&self, other: &Self) -> bool {
        self.pointer == other.pointer
    }
}

mod private {
    pub trait Sealed {}
    impl<T> Sealed for T {}
    impl<T> Sealed for [T] {}
    impl Sealed for str {}
}

/// Types that can be pointed to by `GuestPtr<T>`.
///
/// In essence everything can, and the only special-case is unsized types like
/// `str` and `[T]` which have special implementations.
pub trait Pointee: private::Sealed {
    #[doc(hidden)]
    type Pointer: Copy + PartialEq;
    #[doc(hidden)]
    fn debug(pointer: Self::Pointer, f: &mut fmt::Formatter) -> fmt::Result;
}

impl<T> Pointee for T {
    type Pointer = u32;
    fn debug(pointer: Self::Pointer, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "*guest {pointer:#x}")
    }
}

impl<T> Pointee for [T] {
    type Pointer = (u32, u32);
    fn debug(pointer: Self::Pointer, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "*guest {:#x}/{}", pointer.0, pointer.1)
    }
}

impl Pointee for str {
    type Pointer = (u32, u32);
    fn debug(pointer: Self::Pointer, f: &mut fmt::Formatter) -> fmt::Result {
        <[u8]>::debug(pointer, f)
    }
}

pub fn run_in_dummy_executor<F: std::future::Future>(future: F) -> Result<F::Output> {
    use std::pin::Pin;
    use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};

    let mut f = Pin::from(Box::new(future));
    let waker = dummy_waker();
    let mut cx = Context::from_waker(&waker);
    match f.as_mut().poll(&mut cx) {
        Poll::Ready(val) => return Ok(val),
        Poll::Pending =>
            bail!("Cannot wait on pending future: must enable wiggle \"async\" future and execute on an async Store"),
    }

    fn dummy_waker() -> Waker {
        return unsafe { Waker::from_raw(clone(5 as *const _)) };

        unsafe fn clone(ptr: *const ()) -> RawWaker {
            assert_eq!(ptr as usize, 5);
            const VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop);
            RawWaker::new(ptr, &VTABLE)
        }

        unsafe fn wake(ptr: *const ()) {
            assert_eq!(ptr as usize, 5);
        }

        unsafe fn wake_by_ref(ptr: *const ()) {
            assert_eq!(ptr as usize, 5);
        }

        unsafe fn drop(ptr: *const ()) {
            assert_eq!(ptr as usize, 5);
        }
    }
}