use std::marker::PhantomData;

use std::sync::{Arc, atomic};

use crate::guest::VmGuest;

use crate::{RwLock as SyncRwLock, RwLockWriteGuard as SyncRwLockWriteGuard};

use lyquid::mem::Wasm32;

use lyquor_api::store::LazyBytes;

use lyquor_primitives::StateCategory;

use lyquor_state::{Key, State, StateR, Value};

use tokio::sync::{Mutex, RwLock, RwLockWriteGuard};

use userspace_pagefault::{

    self, AccessType, AsyncPageFaultFuture, AsyncPageStore, PageChunks, PagedSegment, SharedMemory,

};

#[cfg(test)] use lyquor_test::test;

const LYTEPAGE_SIZE: usize = 1024; // the logical page size for LyteMemory

const BEGIN_GUARD_SIZE: usize = 2 << 30;

const END_GUARD_SIZE: usize = 2 << 30;

fn lytememory_setup() {

    static INITIALIZED: std::sync::Once = std::sync::Once::new();

    INITIALIZED.call_once(|| 
{19

        // this is a non-so-bad hack to wasmtime: creating a dummy wasmtime engine will trigger the

        // on-time initialization of the trap handler by wasmtime. Then we can override this trap

        // handler and only hand back when the signals are not triggered within any managed

        // LyteMemory.

        let _ = wasmtime::Engine::new(wasmtime::Config::new().macos_use_mach_ports(false));

    })

}

struct Bitmask(Box<[atomic::AtomicU64]>);

impl Bitmask {

    fn new(size: usize) -> Self {

        Self(

            
std::iter::repeat_with225
(|| atomic::AtomicU64::new(0))

                .take((size + 63) >> 6)

                .collect::<Vec<_>>()

                .into(),

        )

    }

    fn clear(&self) {

        for bits in self.0.iter() {

            bits.store(0, atomic::Ordering::Relaxed)

        }

    }

    #[inline(always)]

    fn mark(&self, idx: usize) {

        self.0[idx >> 6].fetch_or(1 << (idx & 63), atomic::Ordering::Relaxed);

    }

    #[inline(always)]

    fn is_marked(&self, idx: usize) -> bool {

        (self.0[idx >> 6].fetch_or(0, atomic::Ordering::Relaxed) >> (idx & 63)) & 1 == 1

    }

    const TABLE: [u64; 64] = [

        0x00, 0x3a, 0x01, 0x3b, 0x2f, 0x35, 0x02, 0x3c, 0x27, 0x30, 0x1b, 0x36, 0x21, 0x2a, 0x03, 0x3d, 0x33, 0x25,

        0x28, 0x31, 0x12, 0x1c, 0x14, 0x37, 0x1e, 0x22, 0x0b, 0x2b, 0x0e, 0x16, 0x04, 0x3e, 0x39, 0x2e, 0x34, 0x26,

        0x1a, 0x20, 0x29, 0x32, 0x24, 0x11, 0x13, 0x1d, 0x0a, 0x0d, 0x15, 0x38, 0x2d, 0x19, 0x1f, 0x23, 0x10, 0x09,

        0x0c, 0x2c, 0x18, 0x0f, 0x08, 0x17, 0x07, 0x06, 0x05, 0x3f,

    ];

    /// Turn a power of 2 (>0) to the exponent.

    #[inline(always)]

    fn fast_log(mut x: u64) -> u64 {

        x |= x >> 1;

        x |= x >> 2;

        x |= x >> 4;

        x |= x >> 8;

        x |= x >> 16;

        x |= x >> 32;

        Self::TABLE[((x.wrapping_mul(0x03f6eaf2cd271461)) >> 58) as usize]

    }

    #[inline]

    fn get_marked_and_clear(&self) -> impl Iterator<Item = usize> + '_ {

        
self.0120
.
iter120
().
zip120
((
0..120
).
step_by120
(64)).
flat_map120
(|(bits, i)| {

            let mut v = bits.swap(0, atomic::Ordering::Relaxed);

            
std::iter::from_fn479k
(move || {

                // this lowbit finding algorithm is suitable for a sparse bitmask, which is the

                // case typically for LyteMemory dirty pages

                let mut lowbit = v;

                if lowbit == 0 {

                    None

                } else {

                    lowbit = lowbit & lowbit.wrapping_neg();

                    v ^= lowbit; // clear the lowest bit

                    Some(i + Self::fast_log(lowbit) as usize)

                }

            })

        })

    }

}

#[test]

fn test_bitmask() {

    let all_bits: Vec<_> = (0..64).collect();

    let cases = [

        (64, &all_bits[..]),

        (64, &[0, 2, 4, 7, 11, 33, 63][..]),

        (1234, &[0, 32, 63, 64, 100, 256, 1024, 1233][..]),

    ];

    for (size, pos) in cases {

        let bm = Bitmask::new(size);

        for i in pos.iter() {

            bm.mark(*i);

        }

        for (i, j) in bm.get_marked_and_clear().zip(pos.iter()) {

            assert_eq!(i, *j);

        }

    }

}

fn litepage_state_key(litepage_id: u32) -> Key {

    static PREFIX: std::sync::OnceLock<Key> = std::sync::OnceLock::new();

    let key = PREFIX.get_or_init(|| lyquid::LYTEMEM_PAGE_PREFIX.
into18
());

    // using big-endian here so the state debug will be easier to read

    let id: LazyBytes = litepage_id.to_be_bytes().into();

    id.prepend(key)

}

struct Segment {

    /// Dirty page bitmask. One position (bit) corresponds to one OS page.

    dirty: Bitmask,

    /// Loaded page bitmask. One position (bit) corresponds to one OS page.

    loaded: Bitmask,

    /// Size of an OS page on the host.

    page_size: usize,

    /// Offset of this segment from the beginning of the LytePage area.

    litepage_base: usize,

    /// Offest of this segment in the memory (PagedSegment).

    base: usize,

    /// Length of the segment.

    length: usize,

}

impl Segment {

    fn new(base: usize, length: usize, litepage_base: usize, page_size: usize) -> Self {

        // align to the closest page

        let bitmask_size = length.div_ceil(page_size);

        Self {

            dirty: Bitmask::new(bitmask_size),

            loaded: Bitmask::new(bitmask_size),

            page_size,

            litepage_base,

            base,

            length,

        }

    }

    #[inline(always)]

    fn addr_to_idx(&self, mem_addr: usize) -> usize {

        (mem_addr - self.base) / self.page_size

    }

    #[inline(always)]

    fn mark_dirty(&self, mem_addr: usize) {

        self.dirty.mark(self.addr_to_idx(mem_addr))

    }

    #[inline(always)]

    fn is_loaded(&self, mem_addr: usize) -> bool {

        self.loaded.is_marked(self.addr_to_idx(mem_addr))

    }

    #[inline(always)]

    fn mark_loaded(&self, mem_addr: usize) {

        self.loaded.mark(self.addr_to_idx(mem_addr))

    }

    /// Write the dirty pages to the key-value state and clear the dirty bitmask.

    async fn writeback_changes<S: State>(&self, state: &SyncRwLock<S>, mem: &[u8]) {

        let pending = {

            let state = state.read();

            let mut pending = Vec::new();

            for ospage_id in 
self.dirty115
.
get_marked_and_clear115
() {

                let segment_off = ospage_id * self.page_size;

                let litepage_start = (self.litepage_base + segment_off) / LYTEPAGE_SIZE;

                let mem_addr = self.base + segment_off;

                let ospage = &mem[mem_addr..mem_addr + self.page_size];

                for (i, litepage) in 
ospage230
.
chunks230
(LYTEPAGE_SIZE).
enumerate230
() {

                    let key = litepage_state_key((litepage_start + i) as u32);

                    let litepage: Value = litepage.to_vec().into();

                    pending.push((key.clone(), litepage, state.get(key)));

                }

            }

            pending

        };

        let mut changes = Vec::new();

        for (key, litepage, old) in 
pending115
 {

            let changed = match old.await {

                Some(old) => litepage.as_ref() != old.as_ref(), // if the page was previous persisted, compare

                None => !
litepage.as_ref().iter()849
.
all849
(|&b| b == 0), // otherwise check if it's empty

            };

            if changed {

                changes.push((key, litepage));

            }

        }

        let mut state = state.write();

        for (key, litepage) in 
changes115
 {

            state.set(key, Some(litepage));

        }

    }

    /// Revert the dirty pages to their original data and clear the dirty bitmask.

    async fn drop_changes<S: State>(&self, state: &SyncRwLock<S>) -> Vec<(usize, Option<Value>)> {

        let pending = {

            let state = state.read();

            let mut pending = Vec::new();

            for ospage_id in self.dirty.get_marked_and_clear() {

                let segment_off = ospage_id * self.page_size;

                let litepage_start = (self.litepage_base + segment_off) / LYTEPAGE_SIZE;

                let mem_addr = self.base + segment_off;

                for i in 
0..self.page_size / LYTEPAGE_SIZE2
 {

                    let key = litepage_state_key((litepage_start + i) as u32);

                    pending.push((mem_addr + i * LYTEPAGE_SIZE, state.get(key)));

                }

            }

            pending

        };

        let mut restores = Vec::with_capacity(pending.len());

        for (mem_addr, old) in 
pending2
 {

            restores.push((mem_addr, old.await));

        }

        restores

    }

}

#[derive(Clone)]

struct InstanceSegment {

    shm: SharedMemory,

    seg: Arc<Segment>,

    parking: Arc<crate::sync::ParkingSpot>,

}

impl InstanceSegment {

    fn new(

        base: usize, size: usize, litepage_base: usize, page_size: usize,

    ) -> Result<Self, userspace_pagefault::Error> {

        Ok(Self {

            shm: SharedMemory::new(size)
?0
,

            seg: Arc::new(Segment::new(base, size, litepage_base, page_size)),

            parking: Arc::new(crate::sync::ParkingSpot::new()),

        })

    }

}

#[derive(Clone)]

struct Context {

    state: crate::instance::LyquidState,

    network: Arc<Segment>,

    instance: InstanceSegment,

    litepage: std::ops::Range<usize>,

}

impl AsyncPageStore for Context {

    fn page_fault_async(&mut self, offset: usize, length: usize, access: AccessType) -> AsyncPageFaultFuture<'_> {

        Box::pin(async move {

            if !self.litepage.contains(&offset) {

                // the accessed page is outside LytePage (network/instance) address range. Returns None

                // because no loading is needed.

                return None;

            }

            /*

            println!(

                "{:?} pagefault {:?} @ {:x}",

                self as *mut Self,

                access,

                offset - BEGIN_GUARD_SIZE

            );

            */

            // the loaded OS page may span across multiple LytePages.

            let litepage_start = (offset - self.litepage.start) / LYTEPAGE_SIZE;

            let (seg, reads) = {

                // check which segment this address accesses.

                let (state, seg) = match offset {

                    
off223
 if off < self.instance.seg.base => 
(self.state.network.read(), &self.network)223
,

                    _ => (self.state.instance.read(), &self.instance.seg),

                };

                if let AccessType::Write = access {

                    // mark the page as dirty

                    seg.mark_dirty(offset)

                }

                if seg.is_loaded(offset) {

                    // the page was already loaded

                    return None;

                }

                //println!("loading @ {:x}", offset - BEGIN_GUARD_SIZE);

                let reads = (0..length / LYTEPAGE_SIZE)

                    .
map190
(|i| state.get(litepage_state_key((litepage_start + i) as u32)))

                    .collect::<Vec<_>>();

                (seg, reads)

            };

            // collect the contents of all LytePages read from state store.

            let mut chunks = Vec::with_capacity(reads.len());

            for read in 
reads190
 {

                chunks.push(Box::new(read.await.unwrap_or_else(|| 
{724

                    std::iter::repeat_with(|| 0)

                        .take(LYTEPAGE_SIZE)

                        .collect::<Vec<u8>>()

                        .into()

                
}724
)) as Box<dyn AsRef<[u8]>>);

            }

            seg.mark_loaded(offset);

            Some(Box::new(chunks.into_iter()) as PageChunks<'_>)

        })

    }

}

/// Factory and shared backing resources for Lyquid virtual memory objects.

pub struct LyteMemoryManager {

    instance: InstanceSegment,

    instance_mapped: userspace_pagefault::Segment,

    network_range: std::ops::Range<usize>,

    litepage: std::ops::Range<usize>,

    page_size: usize,

    total: usize,

}

impl LyteMemoryManager {

    /// Create the shared instance segment and reserve the LyteMemory virtual address range.

    pub fn new() -> Result<Self, userspace_pagefault::Error> {

        lytememory_setup();

        let wasm_usable = lyquid::LYTEMEM_SIZE_IN_MB << 20;

        let network = lyquid::NETWORK_MEMSIZE_IN_MB << 20;

        let instance = lyquid::INSTANCE_MEMSIZE_IN_MB << 20;

        // WASM 32-bit starting address for LytePage area

        let start = BEGIN_GUARD_SIZE + lyquid::LYTEMEM_BASE;

        // WASM 32-bit ending address for LytePage area

        let end = start + instance + network;

        // total virtual size, including guards

        let total = BEGIN_GUARD_SIZE + wasm_usable + END_GUARD_SIZE;

        assert!(end <= total);

        let page_size = userspace_pagefault::get_page_size()
?0
;

        let network_range = start..start + network;

        let instance_range = start + network..end;

        // all LyteMemories generated from this LyteMemoryManager will share the same

        // InstanceSegment

        let instance = InstanceSegment::new(

            instance_range.start,

            instance_range.len(),

            network_range.len(),

            page_size,

        )?;

        let instance_mapped =

            userspace_pagefault::Segment::new(None, total, page_size, userspace_pagefault::ProtFlags::PROT_NONE)
?0
;

        instance_mapped.make_shared(

            instance.seg.base,

            &instance.shm,

            userspace_pagefault::ProtFlags::PROT_READ | userspace_pagefault::ProtFlags::PROT_WRITE,

        )?;

        Ok(Self {

            instance,

            instance_mapped,

            network_range,

            litepage: start..end,

            page_size,

            total,

        })

    }

    /// Flush the dirty pages in the instance segement to the key-value state and clear the dirty

    /// bitmask.

    pub(crate) async fn flush_instance_state<'a, T, S: State>(

        &self, lock: &'a RwLock<T>, state: &'a SyncRwLock<S>,

    ) -> (RwLockWriteGuard<'a, T>, SyncRwLockWriteGuard<'a, S>) {

        let guard = lock.write().await;

        self.instance

            .seg

            .writeback_changes(state, self.instance_mapped.as_slice())

            .await;

        let state = state.write();

        (guard, state)

    }

    /// Create a new LyteMemory object. The object can be clone-shared but the underlying resource

    /// is the same.

    pub fn new_lytememory<G: VmGuest>(

        &self, state: crate::instance::LyquidState, category: StateCategory,

    ) -> Result<LyteMemory<G>, userspace_pagefault::Error> {

        let ctx = Context {

            state,

            network: Arc::new(Segment::new(

                self.network_range.start,

                self.network_range.len(),

                0,

                self.page_size,

            )),

            instance: self.instance.clone(),

            litepage: self.litepage.clone(),

        };

        let mem = Arc::new(PagedSegment::new_async(self.total, ctx.clone(), None)
?0
);

        mem.make_shared(self.instance.seg.base, &self.instance.shm)
?0
;

        Ok(LyteMemory {

            mem,

            ctx,

            init: Arc::new(Mutex::new(false)),

            category,

            guest: PhantomData,

        })

    }

}

/// A sharable object that represents a userspace-paged virtual memory used by LVM.

#[derive(Clone)]

pub struct LyteMemory<G: VmGuest = Wasm32> {

    mem: Arc<PagedSegment<'static>>,

    ctx: Context,

    init: Arc<Mutex<bool>>,

    category: StateCategory,

    guest: PhantomData<G>,

}

/// Guard returned to the caller that wins one-time LyteMemory initialization.

#[allow(unused)]

pub struct LyteMemoryInitGuard<'a>(tokio::sync::MutexGuard<'a, bool>);

impl<G: VmGuest> LyteMemory<G> {

    /// Reset the dirty-page detection of the instance segment. This is required after a flush

    /// write-back of the instance segment, otherwise future changes will not be detected.

    pub fn reset_instance_write_detection(&self) {

        self.mem

            .reset_write_detection(self.ctx.instance.seg.base, self.ctx.instance.seg.length)

            .expect("Failed to reset write detection."); // FIXME: handle this error

    }

    /// Flush the dirty pages of the network segment to the key-value state carried by LyteMemory.

    /// The dirty-page detection is also reset properly after flushing the pages.

    pub async fn flush_network_state(&self) {

        self.ctx

            .network

            .writeback_changes(self.ctx.state.network.as_ref(), self.mem.as_slice())

            .await;

        self.mem

            .reset_write_detection(self.ctx.network.base, self.ctx.network.length)

            .expect("Failed to reset write detection."); // FIXME: handle this error

        // the lock is released at the end

    }

    /// Revert all dirty pages of the network segment.

    /// The dirty-page detection is also reset properly after flushing the pages.

    pub async fn revert_network_state(&self) {

        let restores = self.ctx.network.drop_changes(self.ctx.state.network.as_ref()).await;

        let (base, len) = self.mem.as_raw_parts();

        let mem = unsafe { std::slice::from_raw_parts_mut(base, len) };

        for (mem_addr, old) in 
restores2
 {

            let litepage = &mut mem[mem_addr..mem_addr + LYTEPAGE_SIZE];

            match old {

                Some(old) => litepage.copy_from_slice(&old),

                None => litepage.fill(0),

            }

        }

        self.mem

            .reset_write_detection(self.ctx.network.base, self.ctx.network.length)

            .expect("Failed to reset write detection."); // FIXME: handle this error

        // the lock is released at the end

    }

    /// Release all in-memory changes and allocated pages.

    #[allow(dead_code)]

    pub async fn release(&self) -> Result<(), userspace_pagefault::Error> {

        self.mem.release_all_pages()?;

        self.mem

            .make_shared(self.ctx.instance.seg.base, &self.ctx.instance.shm)?;

        self.ctx.network.dirty.clear();

        self.ctx.network.loaded.clear();

        *self.init.lock().await = false;

        Ok(())

    }

    /// Returns if the memory needs to be initialized.

    pub async fn need_init(&self) -> Option<LyteMemoryInitGuard<'_>> {

        let mut init = self.init.lock().await;

        match *init {

            true => None,

            false => {

                // only initialize once

                *init = true;

                Some(LyteMemoryInitGuard(init))

            }

        }

    }

    /// Manually set the initialized flag.

    pub async fn set_init(&self, flag: bool) {

        *self.init.lock().await = flag

    }

    /// Get the raw pointer to the host memory at a given WASM address.

    #[inline(always)]

    pub fn host_ptr_by_wasm_addr(&self, addr: G::Usize) -> *mut u8 {

        let (base, _) = self.mem.as_raw_parts();

        base.wrapping_add(G::usize_to_host(addr) + BEGIN_GUARD_SIZE)

    }

    /// Get a read-only slice of memory by WASM address.

    pub fn slice_by_wasm_addr(&self, offset: G::Usize, length: G::Usize) -> &[u8] {

        unsafe { std::slice::from_raw_parts(self.host_ptr_by_wasm_addr(offset), G::usize_to_host(length)) }

    }

    /// Get a mutable slice of the memory with WASM address.

    ///

    /// # Safety

    ///

    /// The caller must ensure no other live references alias the returned range for the duration

    /// of the returned slice.

    #[allow(clippy::mut_from_ref)]

    pub unsafe fn slice_by_wasm_addr_mut(&self, offset: G::Usize, length: G::Usize) -> &mut [u8] {

        unsafe { std::slice::from_raw_parts_mut(self.host_ptr_by_wasm_addr(offset), G::usize_to_host(length)) }

    }

    /// Get the function category of this LyteMemory's owner.

    pub fn category(&self) -> StateCategory {

        self.category

    }

    /// Return the wait/notify queue associated with the shared instance segment.

    pub(crate) fn instance_parking_spot(&self) -> Arc<crate::sync::ParkingSpot> {

        self.ctx.instance.parking.clone()

    }

}

struct WasmRuntimeMemory<G: VmGuest> {

    inner: LyteMemory<G>,

    size: usize,

}

unsafe impl<G: VmGuest> wasmtime::LinearMemory for WasmRuntimeMemory<G> {

    fn byte_size(&self) -> usize {

        self.size

    }

    fn byte_capacity(&self) -> usize {

        self.size

    }

    fn grow_to(&mut self, size: usize) -> wasmtime::Result<()> {

        if size > self.size {

            return Err(wasmtime::Error::msg("LyteMemory is not growable."));

        }

        Ok(())

    }

    fn as_ptr(&self) -> *mut u8 {

        let (ptr, _) = self.inner.mem.as_raw_parts();

        unsafe { ptr.add(BEGIN_GUARD_SIZE) }

    }

}

/// Memory factory that makes wasmtime memory for the given LyteMemory.

pub struct WasmMemoryCreator<G: VmGuest = Wasm32>(pub LyteMemory<G>);

unsafe impl<G: VmGuest> wasmtime::MemoryCreator for WasmMemoryCreator<G> {

    fn new_memory(

        &self, _mt: wasmtime::MemoryType, _min: usize, _max: Option<usize>, _reserved: Option<usize>, _guard: usize,

    ) -> Result<Box<dyn wasmtime::LinearMemory>, String> {

        Ok(Box::new(WasmRuntimeMemory {

            inner: self.0.clone(),

            size: lyquid::LYTEMEM_SIZE_IN_MB << 20,

        }))

    }

}