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

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(());
}

fn open_ptr(p: Pointer) -> (usize, usize) {
    let Pointer(u) = p;
    return (u / 8, u % 8);
}

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

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

    fn alloc_size(&mut self, p: Pointer) -> Result<usize, String> {
        let (u, _) = open_ptr(p);
        if u == 0 {
            return Err("cannot get the size of a nil pointer");
        }
        // Alloc size is stored just below the pointer.
        let size = self.heap[u - 1];
        return Ok(usize::try_from(size).unwrap() >> 1);
    }

    fn is_bytevector(&mut self, p: Pointer) -> Result<bool, String> {
        let (u, _) = open_ptr(p);
        let size = self.heap[u - 1];
        // Low bit 1 means bytevector.
        return Ok(size & 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> {
        let n_cells = (n + 7) / 8;
        if self.heap.len() - self.free_pointer < n_cells {
            self.collect_garbage(n, stack, locals)?;
        }
        let len_p = &mut self.heap[self.free_pointer]?;
        *len_p = u64::try_from(n).unwrap() << 1;
        if bytevector_p {
            *len_p |= 1;
        }
        self.free_pointer += 1;
        let p = Pointer(self.free_pointer * 8);
        let n_cells = n / 8;
        self.heap[self.free_pointer..self.free_pointer + n].fill(0);
        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 (u, _) = open_ptr(p);
        if u >= self.heap.len() {
            return Err(format!("invalid pointer {:x}", p));
        }
        return Ok(Value(self.heap[u]));
    }

    pub fn poke(&mut self, v: Value, p: Pointer) -> Result<(), String> {
        let (u, _) = open_ptr(p);
        let Value(x) = v;
        if u >= self.heap.len() {
            return Err(format!("invalid pointer {:x}", p));
        self.heap[u] = x;
        return Ok(());
    }

    fn peek_byte(&self, p: Pointer) -> Result<u8, String> {
        let (word_cnt, word_offset) = open_ptr(p);
        if word_cnt >= self.heap.len() {
            return Err(format!("invalid pointer {:x}", p));
        }
        return (self.heap[word_cnt] >> word_offset * 8) as u8;
    }

    pub fn poke_byte(&mut self, u: u8, p: Pointer) -> Result<(), String> {
        let (word_cnt, word_offset) = open_ptr(p);
        if word_cnt >= self.heap.len() {
            return Err(format!("invalid pointer {:x}", p));
        }
        let surrounding_word = self.heap[word_cnt];
        let mask = !(0xff << word_offset * 8);
        self.heap[word_cnt] = surrounding_word & mask | u << word_offset * 8;
    }
}