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

fn rewrite_pointers(
    stack: &mut [Value],
    rewrites: &HashMap<Pointer, Pointer>,
) -> Result<(), String> {
    for val in stack.iter_mut() {
        let p = if val.is_pointer() {
            val.to_pointer().unwrap()
        } else {
            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>,
    spare_heap: Vec<u64>,
    last_reachable_cells: usize,
}

impl Heap {
    pub fn new() -> Self {
        Heap {
            heap: Vec::with_capacity(512),
            spare_heap: Vec::with_capacity(512),
            last_reachable_cells: 0,
        }
    }

    fn alloc_size(&self, p: Pointer) -> Result<usize, String> {
        let (u, _) = open_ptr(p);
        if u == 0 {
            return Err(String::from("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 walk_gc_roots(
        &mut self,
        roots: &[Value],
        rewrites: &mut HashMap<Pointer, Pointer>,
    ) -> Result<(), String> {
        let mut queue: VecDeque<Value> = VecDeque::with_capacity(roots.len());
        queue.extend(roots);
        while let Some(val) = queue.pop_front() {
            let p = if val.is_pointer() {
                val.to_pointer().unwrap()
            } else {
                continue;
            };
            if rewrites.contains_key(&p) {
                // Already copied this one.
                continue;
            }
            let (u, _) = open_ptr(p);
            let object_size = self.alloc_size(p)?;
            let object_size_words = object_size / 8;
            // Copy object and size.
            rewrites.insert(p, Pointer(self.spare_heap.len() * 8 + 8));
            self.spare_heap
                .extend(&self.heap[u - 1..u + object_size_words]);
            if self.is_bytevector(p)? {
                // Don't process bytevectors recursively. We're all done.
                continue;
            }
            for i in u..u + object_size_words {
                queue.push_back(Value(self.heap[i]));
            }
        }
        return Ok(());
    }

    fn collect_garbage(&mut self, locals: &mut [Value]) -> Result<(), String> {
        self.spare_heap.truncate(0);
        let mut rewrites = HashMap::new();
        self.walk_gc_roots(locals, &mut rewrites)?;
        // Rewrite values.
        rewrite_pointers(locals, &rewrites)?;
        // Activate the new heap!
        std::mem::swap(&mut self.heap, &mut self.spare_heap);
        // Walk objects in the heap and rewrite pointers. First object is at address 8.
        let mut i = 8;
        while i < self.heap.len() * 8 {
            let p = Pointer(i);
            if self.is_bytevector(p)? {
                i += self.alloc_size(p)?;
                continue;
            }
            let sz = self.alloc_size(p)?;
            for j in (0..sz).step_by(8) {
                let val = self.peek(p, j)?;
                let vp = if val.is_pointer() {
                    val.to_pointer().unwrap()
                } else {
                    continue;
                };
                let Pointer(u) = vp;
                self.poke(
                    Value::from_pointer(
                        *rewrites
                            .get(&vp)
                            .ok_or(format!("no rewrite found for {:x}", u))?,
                    ),
                    p,
                    j,
                )?;
            }
            if sz == 0 {
                return Err(String::from("object with size 0 found on the heap"));
            }
            i += sz + 8;
        }
        self.last_reachable_cells = self.heap.len();
        // Done??
        return Ok(());
    }

    fn alloc_b(
        &mut self,
        n: usize,
        locals: &mut [Value],
        bytevector_p: bool,
    ) -> Result<Pointer, String> {
        if self.heap.len() > 2 * self.last_reachable_cells {
            self.collect_garbage(locals)?;
        }
        let n_cells = (n + 7) / 8;
        let len_idx = self.heap.len();
        self.heap.resize(self.heap.len() + n_cells + 1, 0);
        let len_p = &mut self.heap[len_idx];
        *len_p = u64::try_from(n).unwrap() << 1;
        if bytevector_p {
            *len_p |= 1;
        }
        let obj_start = len_idx + 1;
        let p = Pointer(obj_start * 8);
        self.heap[obj_start..obj_start + n_cells].fill(0);
        return Ok(p);
    }

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

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

    pub fn peek(&self, p: Pointer, offset: usize) -> Result<Value, String> {
        let offset_bytes = offset * 8;
        let obj_sz = self.alloc_size(p)?;
        if offset_bytes >= obj_sz {
            return Err(format!(
                "invalid offset {offset_bytes} in an object of size {obj_sz}"
            ));
        }
        let (u, _) = open_ptr(p.offset(offset_bytes));
        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, offset: usize) -> Result<(), String> {
        let obj_sz = self.alloc_size(p)?;
        if offset >= obj_sz {
            return Err(format!(
                "invalid offset {offset} in an object of size {obj_sz}"
            ));
        }
        let (u, _) = open_ptr(p.offset(offset));
        let Value(x) = v;
        if u >= self.heap.len() {
            return Err(format!("invalid pointer {:x}", p));
        }
        self.heap[u] = x;
        return Ok(());
    }

    pub 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));
        }
        Ok((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 | u64::from(u) << word_offset * 8;
        Ok(())
    }
}