use crate::data::{Pointer, Value};
use std::collections::HashMap;
use std::iter::Iterator;

fn transmute(p: &u8) -> Result<&u64, String> {
    let u8_p = p as *const u8;
    let u64_p = u8_p as *const u64;
    if u64_p as usize % 8 != 0 {
        return Err(String::from("not aligned"));
    }
    unsafe {
        return Ok(&*u64_p);
    }
}

fn transmute_mut(p: &mut u8) -> Result<&mut u64, String> {
    let u8_p = p as *mut u8;
    let u64_p = u8_p as *mut u64;
    if u64_p as usize % 8 != 0 {
        return Err(String::from("not aligned"));
    }
    unsafe {
        return Ok(&mut *u64_p);
    }
}

fn rewrite_pointers(
    stack: &mut [Value],
    rewrites: &HashMap<Pointer, Pointer>,
) -> Result<(), String> {
    for val in stack.iter_mut() {
        let p = match val.to_pointer() {
            Ok(p) => p,
            Err(_) => {
                continue;
            }
        };
        let Pointer(u) = p;
        *val = Value::from_pointer(
            *rewrites
                .get(&p)
                .ok_or(format!("no rewrite found for {:x}", u))?,
        );
    }
    return Ok(());
}

pub struct Heap {
    heap: Vec<u8>,
    free_pointer: usize,
    spare_heap: Vec<u8>,
}

impl Heap {
    pub fn new() -> Self {
        Heap {
            heap: vec![0; 1024], // 1 kB
            free_pointer: 0,
            spare_heap: Vec::new(),
        }
    }

    fn alloc_size(&mut self, p: Pointer) -> Result<usize, String> {
        let Pointer(u) = p;
        // Alloc size is stored just below the pointer.
        let u8_p = &self.heap[u - 8];
        let u64_p = transmute(u8_p)?;
        return Ok(usize::try_from(*u64_p).unwrap() >> 1);
    }

    fn is_bytevector(&mut self, p: Pointer) -> Result<bool, String> {
        let Pointer(u) = p;
        let u8_p = &self.heap[u - 8];
        let u64_p = transmute(u8_p)?;
        // Low bit 1 means bytevector.
        return Ok(*u64_p & 1 != 0);
    }

    fn gc_process_value(
        &mut self,
        val: Value,
        spare_heap_ptr: &mut usize,
        rewrites: &mut HashMap<Pointer, Pointer>,
    ) -> Result<(), String> {
        let p = match val.to_pointer() {
            Ok(p) => p,
            Err(_) => {
                return Ok(());
            }
        };
        if rewrites.contains_key(&p) {
            // Already copied this one.
            return Ok(());
        }
        let Pointer(u) = p;
        let object_size = self.alloc_size(p)?;
        // Copy object and size.
        self.spare_heap[*spare_heap_ptr..*spare_heap_ptr + object_size + 8]
            .copy_from_slice(&self.heap[u - 8..u + object_size]);
        rewrites.insert(p, Pointer(*spare_heap_ptr + 8));
        *spare_heap_ptr += object_size + 8;
        if self.is_bytevector(p)? {
            // Don't process bytevectors recursively. We're all done.
            return Ok(());
        }
        for i in (u..u + object_size).step_by(8) {
            self.gc_process_value(self.peek(Pointer(i))?, spare_heap_ptr, rewrites)?;
        }
        return Ok(());
    }

    fn walk_gc_roots(
        &mut self,
        roots: &[Value],
        spare_heap_ptr: &mut usize,
        rewrites: &mut HashMap<Pointer, Pointer>,
    ) -> Result<(), String> {
        for &val in roots {
            self.gc_process_value(val, spare_heap_ptr, rewrites)?;
        }
        return Ok(());
    }

    fn collect_garbage(
        &mut self,
        size_hint: usize,
        stack: &mut [Value],
        locals: &mut [Value],
    ) -> Result<(), String> {
        const MAX_HEAP_SIZE: usize = 4 * 1024 * 1024; // 4 GB
                                                      // Always at least double the heap size (keeping in mind the max heap size).
        let mut size_hint = size_hint;
        if size_hint < self.heap.len() {
            size_hint = self.heap.len();
        }
        let mut new_heap_size = self.heap.len() + size_hint;
        if new_heap_size > MAX_HEAP_SIZE / 2 {
            new_heap_size = MAX_HEAP_SIZE / 2;
        }
        self.spare_heap.resize(new_heap_size, 0);
        let mut spare_heap_ptr = 0;
        let mut rewrites = HashMap::new();
        self.walk_gc_roots(stack, &mut spare_heap_ptr, &mut rewrites)?;
        self.walk_gc_roots(locals, &mut spare_heap_ptr, &mut rewrites)?;
        // Walk the stacks and rewrite.
        rewrite_pointers(stack, &rewrites)?;
        rewrite_pointers(locals, &rewrites)?;
        // Activate the new heap!
        std::mem::swap(&mut self.heap, &mut self.spare_heap);
        self.free_pointer = spare_heap_ptr;
        // Walk objects in the heap and rewrite pointers. First object is at address 8.
        let mut i = 8;
        while i < self.free_pointer {
            let p = Pointer(i);
            if self.is_bytevector(p)? {
                i += self.alloc_size(p)?;
                continue;
            }
            for j in (i..i + self.alloc_size(p)?).step_by(8) {
                let q = Pointer(j);
                let val = self.peek(q)?;
                let vp = match val.to_pointer() {
                    Ok(x) => x,
                    Err(_) => {
                        continue;
                    }
                };
                let Pointer(u) = vp;
                self.poke(
                    Value::from_pointer(
                        *rewrites
                            .get(&vp)
                            .ok_or(format!("no rewrite found for {:x}", u))?,
                    ),
                    q,
                )?;
            }
        }
        // Done??
        return Ok(());
    }

    fn alloc_b(
        &mut self,
        n: usize,
        stack: &mut [Value],
        locals: &mut [Value],
        bytevector_p: bool,
    ) -> Result<Pointer, String> {
        if self.heap.len() - self.free_pointer < n {
            self.collect_garbage(n, stack, locals)?;
            if self.heap.len() - self.free_pointer < n {
                return Err(String::from("out of space"));
            }
        }
        let len_p = transmute_mut(&mut self.heap[self.free_pointer])?;
        *len_p = u64::try_from(n).unwrap() << 1;
        if bytevector_p {
            *len_p |= 1;
        }
        self.free_pointer += 8;
        let p = Pointer(self.free_pointer);
        self.heap[self.free_pointer..self.free_pointer + n].fill(0);
        self.free_pointer += n;
        return Ok(p);
    }

    pub fn alloc(
        &mut self,
        n: usize,
        stack: &mut [Value],
        locals: &mut [Value],
    ) -> Result<Pointer, String> {
        return self.alloc_b(n, stack, locals, false);
    }

    pub fn alloc_bytevector(
        &mut self,
        n: usize,
        stack: &mut [Value],
        locals: &mut [Value],
    ) -> Result<Pointer, String> {
        return self.alloc_b(n, stack, locals, true);
    }

    pub fn peek(&self, p: Pointer) -> Result<Value, String> {
        let Pointer(x) = p;
        let u8_p = &self.heap[x];
        let u64_p = transmute(u8_p)?;
        return Ok(Value(*u64_p));
    }

    pub fn poke(&mut self, v: Value, p: Pointer) -> Result<(), String> {
        let Pointer(x) = p;
        let Value(u) = v;
        let u8_p = &mut self.heap[x];
        let u64_p = transmute_mut(u8_p)?;
        *u64_p = u;
        return Ok(());
    }

    pub fn peek_byte(&self, p: Pointer) -> u8 {
        let Pointer(x) = p;
        return self.heap[x];
    }

    pub fn poke_byte(&mut self, u: u8, p: Pointer) {
        let Pointer(x) = p;
        self.heap[x] = u;
    }
}