#![feature(ptr_metadata)]
#![feature(layout_for_ptr)]
#![feature(coerce_unsized)]

#[cfg(test)]
mod test;

pub mod prelude;

use core::panic;
#[cfg(feature = "only_store_meta")]
use std::{ptr::Pointee, mem::align_of_val_raw};
use std::{
    ptr::{NonNull, self, drop_in_place, metadata, Pointee},
    marker::PhantomData,
    alloc::{alloc, Layout},
    mem::{size_of, size_of_val, align_of_val, self, size_of_val_raw, align_of_val_raw},
    slice,
    fmt::Debug,
    ops::{CoerceUnsized, Index, IndexMut}
};

unsafe fn dealloc(ptr: *mut u8, layout: Layout) {
    std::alloc::dealloc(ptr, layout)
}

/// Copy `size` bytes of memory from `src` to `dst`.
/// 
/// # Safety
/// 
/// `src` must be valid for reads, `dst` must be valid for writes, etc, you get the idea.
// TODO: inline me! i didnt realize it was avaliable as `copy_from` until the code was mostly complete.
unsafe fn memcpy(src: *const u8, dst: *mut u8, size: usize) {
    dst.copy_from(src, size);
}

fn align_up<T: ?Sized>(ptr: *const T, align: usize) -> *const T {
    let (mut data, meta) = ptr.to_raw_parts();
    data = ((data as usize + align - 1) & !(align - 1)) as _;
    ptr::from_raw_parts(data, meta)
}

fn align_up_mut<T: ?Sized>(ptr: *mut T, align: usize) -> *mut T {
    align_up(ptr as _, align) as _
}

trait Strategy<T: ?Sized>: Sized {
    type Metadata: Copy;

    unsafe fn get_ptr(vec: &Vec<T, Self>, index: usize) -> *const T;

    unsafe fn next(vec: &Vec<T, Self>, prev: *const T, index: usize) -> *const T;

    fn create_meta_from_ptr(vec: &Vec<T, Self>, ptr: *const T) -> Self::Metadata;
}

struct StorePtr;

impl<T: ?Sized> Strategy<T> for StorePtr {
    type Metadata = *const T;

    unsafe fn get_ptr(vec: &Vec<T, Self>, index: usize) -> *const T {
        vec.get_ptr_to_meta(index + 1).cast::<*const T>().read()
    }

    unsafe fn next(vec: &Vec<T, Self>, _prev: *const T, index: usize) -> *const T {
        Self::get_ptr(vec, index)
    }

    fn create_meta_from_ptr(_vec: &Vec<T, Self>, ptr: *const T) -> Self::Metadata {
        ptr
    }
}

struct OnlyStoreMeta;

impl<T: ?Sized> Strategy<T> for OnlyStoreMeta {
    type Metadata = <T as Pointee>::Metadata;

    unsafe fn get_ptr(vec: &Vec<T, Self>, index: usize) -> *const T {
        let meta = vec.get_ptr_to_meta(1).cast::<<T as Pointee>::Metadata>().read();
        let fake = ptr::from_raw_parts::<T>(0 as _, meta);
        let ptr = align_up(vec.ptr.as_ptr(), align_of_val_raw(fake));
        let mut ptr = ptr::from_raw_parts(ptr.cast(), meta);

        for index in 1..=index {
            ptr = Self::next(vec, ptr, index);
        }

        ptr
    }

    unsafe fn next(vec: &Vec<T, Self>, prev: *const T, index: usize) -> *const T {
        let ptr = prev.cast::<u8>().wrapping_add(size_of_val_raw(prev));

        let meta = vec.get_ptr_to_meta(index + 1).cast::<<T as Pointee>::Metadata>().read();
        let fake = ptr::from_raw_parts::<T>(0 as _, meta);
        let ptr = align_up(ptr, align_of_val_raw(fake));
        ptr::from_raw_parts(ptr.cast(), meta)
    }

    fn create_meta_from_ptr(_vec: &Vec<T, Self>, ptr: *const T) -> Self::Metadata {
        metadata(ptr)
    }
}

/// A heap allocated, dynamically sized collection of `?Sized` elements.
/// 
/// See [`::alloc::vec::Vec`] (the standard library `Vec` type) for more information.
struct Vec<T: ?Sized, S: Strategy<T>> {
    ptr: NonNull<u8>,
    len: usize,
    capacity: usize,
    end_ptr: NonNull<u8>,
    _phantom: PhantomData<(*mut S, T)>
}

// keeps this file cleaner
mod impls;
mod iter;
pub use iter::*;

/// The data stored as metadata at the end of the allocation.
#[cfg(feature = "only_store_meta")]
type Meta<T> = <T as Pointee>::Metadata;
#[cfg(not(feature = "only_store_meta"))]
type Meta<T> = *const T;

#[cfg(feature = "only_store_meta")]
unsafe fn size_of_val_meta<T: ?Sized>(meta: Meta<T>) -> usize {
    size_of_val_raw(ptr::from_raw_parts::<T>(0 as *const (), meta))
}

#[cfg(feature = "only_store_meta")]
unsafe fn align_of_val_meta<T: ?Sized>(meta: Meta<T>) -> usize {
    align_of_val_raw(ptr::from_raw_parts::<T>(0 as *const (), meta))
}

impl<T: ?Sized, S: Strategy<T>> Vec<T, S> {
    /// Creates a new, empty `Vec`.
    pub fn new() -> Self {
        let ptr = NonNull::dangling();
        Self {
            ptr,
            len: 0,
            capacity: 0,
            end_ptr: ptr,
            _phantom: PhantomData
        }
    }

    /// Appends an element to the end of the `Vec`.
    /// 
    /// Only avaliable if `T: Sized`.
    pub fn push(&mut self, v: T) where T: Sized {
        unsafe { self.push_raw(&v) }
        mem::forget(v);
    }

    /// Appends an (possibly unsized) boxed element to the end of the `Vec`.
    pub fn push_box(&mut self, v: Box<T>) {
        let ptr = Box::into_raw(v);
        unsafe {
            self.push_raw(ptr);
            dealloc(ptr.cast(), Layout::for_value_raw(ptr));
        }
    }

    /// Appends a sized element of type `U` to the end of the `Vec`, given that it can be `CoerceUnsized` to a `T`.
    pub fn push_unsize<U>(&mut self, v: U) where for<'a> &'a U: CoerceUnsized<&'a T> {
        let v_unsized: &T = &v;
        unsafe { self.push_raw(v_unsized) };
        mem::forget(v);
    }

    unsafe fn push_raw(&mut self, v: *const T) {
        let size = size_of_val(&*v);

        if !self.will_fit(&*v) {
            // oh no! allocation too small!

            // make sure we have enough space for a new element, but also space for future elements
            // this bit is tricky, we must make sure we have enough space for padding too, so its probably UB somehow
            // FIXME: ^^^
            let new_alloc_size = self.capacity * 2 + size * 2 + size_of::<S::Metadata>();
            self.realloc(new_alloc_size);
        }

        self.push_raw_unchecked(v);
    }

    /// Given an element, returns a pointer to where it would be written if it was pushed, assuming no reallocation is needed.
    /// 
    /// The pointer will be aligned, but writing to it may overwrite data belonging to the Vec.
    /// To check for this, call `will_fit`.
    /// In addition, the metedata for the element must be set using `set_meta_from_ptr`.
    fn get_next_elem_ptr(&self, v: &T) -> *mut u8 {
        align_up_mut(self.end_ptr.as_ptr(), align_of_val(v))
    }

    /// Checks if a given element will fill in the `Vec` without reallocations.
    pub fn will_fit(&self, v: &T) -> bool {
        let remaining_space = self.get_ptr_to_meta(self.len) as usize - self.end_ptr.as_ptr() as usize;
        let needed_space = size_of_val(v) + size_of::<S::Metadata>();
        remaining_space >= needed_space
    }

    unsafe fn push_raw_unchecked(&mut self, v: *const T) {
        let size = size_of_val(&*v);
        let dest = self.get_next_elem_ptr(&*v); // this is mentioned by the `// SAFETY:` in `as_slice_flatten`

        memcpy(v.cast(), dest, size);

        let new_ptr = ptr::from_raw_parts::<T>(dest.cast(), metadata(v));
        self.set_meta_from_ptr(self.len, new_ptr);

        self.end_ptr = NonNull::new_unchecked(dest.wrapping_add(size));
        self.len += 1;
    }

    unsafe fn realloc(&mut self, size: usize) {
        let layout = Layout::from_size_align_unchecked(size, 8).pad_to_align();
        if self.capacity == 0 {
            // will panic if OOM
            self.ptr = NonNull::new(alloc(layout)).unwrap();

            self.end_ptr = self.ptr;
        } else {
            // cannot use realloc here

            let new_alloc = NonNull::new(alloc(layout)).unwrap();

            // data
            let mut ptr = new_alloc.as_ptr();
            for i in 0..self.len {
                let v = self.get_unchecked(i);

                let size = size_of_val(v);
                ptr = align_up_mut(ptr, align_of_val(v));
                memcpy(v as *const _ as _, ptr, size);
                let meta = self.get_ptr(i).to_raw_parts().1;
                self.set_meta_from_ptr(i, ptr::from_raw_parts(ptr.cast(), meta));
                ptr = ptr.wrapping_add(size);
            }
            self.end_ptr = NonNull::new_unchecked(ptr);

            // metadata
            let meta_src = self.get_ptr_to_meta(self.len);
            let meta_dst = {
                let current_alloc_end = self.ptr.as_ptr().wrapping_add(self.capacity);
                let new_alloc_end = new_alloc.as_ptr().wrapping_add(layout.size());
                let meta_len = current_alloc_end as usize - meta_src as usize;
                new_alloc_end.wrapping_sub(meta_len)
            };
            let meta_size = self.len * size_of::<Meta<T>>();
            memcpy(meta_src.cast(), meta_dst, meta_size);

            dealloc(self.ptr.as_ptr(), Layout::from_size_align_unchecked(self.capacity, 8));

            self.ptr = new_alloc;
        }

        self.capacity = layout.size();
    }

    /// for internal use
    /// 
    /// NOTE: 1-indexed, to allow getting a pointer to the end of the alloc easily
    fn get_ptr_to_meta(&self, index: usize) -> *mut S::Metadata {
        self.ptr.as_ptr()
        .wrapping_add(self.capacity)
        .cast::<S::Metadata>()
        .wrapping_sub(index)
    }

    /// for internal use
    unsafe fn get_meta(&self, index: usize) -> S::Metadata {
        *self.get_ptr_to_meta(index + 1)
    }

    unsafe fn set_meta_from_ptr(&self, index: usize, ptr: *const T) {
        self.get_ptr_to_meta(index + 1).write(S::create_meta_from_ptr(self, ptr));
    }

    unsafe fn get_ptr(&self, index: usize) -> *const T {
        S::get_ptr(self, index)
    }

    pub fn get(&self, index: usize) -> Option<&T> {
        if index < self.len {
            Some(unsafe { self.get_unchecked(index) })
        } else {
            None
        }
    }

    pub unsafe fn get_unchecked(&self, index: usize) -> &T {
        &*self.get_ptr(index)
    }

    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
        if index < self.len {
            Some(unsafe { self.get_unchecked_mut(index) })
        } else {
            None
        }
    }

    pub unsafe fn get_unchecked_mut(&mut self, index: usize) -> &mut T {
        &mut *(self.get_ptr(index) as *mut _)
    }

    /// Returns the length of the `Vec`, which is how many items it contains.
    pub fn len(&self) -> usize {
        self.len
    }

    /// Returns the capacity, which is the size of the allocation in bytes.
    pub fn capacity(&self) -> usize {
        self.capacity
    }

    /// Returns a pointer to the allocation of the `Vec`.
    pub fn as_ptr(&self) -> *const u8 {
        self.ptr.as_ptr()
    }

    /// Returns a mutable pointer to the allocation of the `Vec`.
    pub fn as_mut_ptr(&mut self) -> *mut u8 {
        self.ptr.as_ptr()
    }

    /// Iterates over the `Vec` by-ref.
    pub fn iter(&self) -> Iter<T> {
        Iter::new(self)
    }

    /// Iterates over the `Vec` by-mut.
    pub fn iter_mut(&mut self) -> IterMut<T> {
        IterMut::new(self)
    }

    /// Converts a `Vec<T: Sized>` into a `Vec<U: ?Sized>`, given that `T` can be `CoerceUnsized` into `U`.
    pub fn unsize<U: ?Sized>(self) -> Vec<U> where for<'a> &'a T: CoerceUnsized<&'a U>, S: Strategy<U> {
        let new_vec = Vec::<U> {
            ptr: self.ptr,
            len: self.len,
            capacity: self.capacity,
            end_ptr: self.end_ptr,
            _phantom: PhantomData,
        };

        if size_of::<<S as Strategy<U>>::Metadata>() > size_of::<<S as Strategy<T>>::Metadata>() {
            // new meta larger than old meta, must go from back to front

            for i in (0..self.len).rev() {
                let current = unsafe { &*self.get_ptr(i) };
                unsafe { new_vec.set_meta_from_ptr(i, current as &U) }
            }
        } else {
            // new meta smaller or same size as old meta, must go from front to back

            for i in 0..self.len {
                let current = unsafe { &*self.get_ptr(i) };
                unsafe { new_vec.set_meta_from_ptr(i, current as &U) }
            }
        }

        mem::forget(self);
        new_vec
    }
}

impl<T> Vec<[T]> {
    pub fn as_slice_flatten(&self) -> &[T] {
        assert!(self.len > 0);

        // SAFETY: the slices should be contiguous by the logic of `push_raw_unchecked`
        unsafe {
            slice::from_raw_parts(self.get_ptr(0).to_raw_parts().0 as _, {
                let start = self.get_ptr(0).to_raw_parts().0 as usize;
                let end = self.end_ptr.as_ptr() as usize;
                (end - start) / size_of::<T>() // integer division!
            })
        }
    }

    pub fn as_mut_slice_flatten(&mut self) -> &mut [T] {
        assert!(self.len > 0);

        // SAFETY: the slices should be contiguous by the logic of `push_raw_unchecked`
        unsafe {
            slice::from_raw_parts_mut(self.get_ptr(0).to_raw_parts().0 as _, {
                let start = self.get_ptr(0).to_raw_parts().0 as usize;
                let end = self.end_ptr.as_ptr() as usize;
                (end - start) / size_of::<T>() // integer division!
            })
        }
    }
}

impl<T: ?Sized> Drop for Vec<T> {
    fn drop(&mut self) {
        unsafe {
            for i in 0..self.len {
                println!("dropping {}", i);
                drop_in_place(self.get_unchecked_mut(i));
            }
            dealloc(self.ptr.as_ptr(), Layout::from_size_align_unchecked(self.capacity, 8));
        }
    }
}

impl<T: ?Sized> Index<usize> for Vec<T> {
    type Output = T;

    #[track_caller]
    fn index(&self, index: usize) -> &Self::Output {
        match self.get(index) {
            Some(v) => v,
            None => panic!("index out of bounds: the len is {} but the index is {}", self.len, index),
        }
    }
}

impl<T: ?Sized> IndexMut<usize> for Vec<T> {
    #[track_caller]
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        let len = self.len;
        match self.get_mut(index) {
            Some(v) => v,
            None => panic!("index out of bounds: the len is {} but the index is {}", len, index),
        }
    }
}

/// Creates a [`Vec`].
/// 
/// # Examples
/// 
/// ```
/// # use dyn_vec::prelude::{vec, Vec};
/// # use std::fmt::Debug;
/// let vec1: Vec<dyn Debug> = vec![1, 2, 3].unsize();
/// let vec2: Vec<dyn Debug> = vec![box:
///     Box::new(1) as _,
///     Box::new(String::from("foo")) as _,
///     Box::new(true) as _
/// ];
/// let vec3: Vec<dyn Debug> = vec![unsized: 1, String::from("foo"), true];
/// ```
#[macro_export]
macro_rules! vec {
    () => {
        $crate::Vec::new();   
    };
    (box: $($elem:expr),+ $(,)?) => {{
        let mut vec = $crate::Vec::new();
        $(vec.push_box($elem);)+
        vec
    }};
    (unsized: $($elem:expr),+ $(,)?) => {{
        let mut vec = $crate::Vec::new();
        $(vec.push_unsize($elem);)+
        vec
    }};
    ($elem:expr; $n:expr) => {
        compile_error!("dyn_vec::vec![T; N] is currently not supported");
    };
    ($($elem:expr),+ $(,)?) => {{
        let mut vec = $crate::Vec::new();
        $(vec.push($elem);)+
        vec
    }};
}