use super::*;

impl VirtualMachine {
    pub(super) fn handle_safepoint(&mut self) {
        if self.gc_alloc_threshold > 0 {
            let live_now = self.heap.len();
            let since_last = live_now.saturating_sub(self.last_gc_live_count);
            if since_last >= self.gc_alloc_threshold {
                struct CollectRoots(Vec<prometeu_bytecode::HeapRef>);
                impl crate::roots::RootVisitor for CollectRoots {
                    fn visit_heap_ref(&mut self, r: prometeu_bytecode::HeapRef) {
                        self.0.push(r);
                    }
                }
                let mut collector = CollectRoots(Vec::new());
                self.visit_roots(&mut collector);
                if let Some(cur) = self.current_coro {
                    collector.0.push(cur);
                }
                let mut coro_roots = self.heap.suspended_coroutine_handles();
                collector.0.append(&mut coro_roots);
                self.heap.mark_from_roots(collector.0);
                self.heap.sweep();
                self.last_gc_live_count = self.heap.len();
            }
        }

        self.current_tick = self.current_tick.wrapping_add(1);
        self.scheduler.wake_ready(self.current_tick);

        let mut switched_out = false;
        if let Some(cur) = self.current_coro {
            if let Some(wake) = self.sleep_requested_until.take() {
                if let Some(co) = self.heap.coroutine_data_mut(cur) {
                    co.pc = self.pc;
                    co.stack = std::mem::take(&mut self.operand_stack);
                    co.frames = std::mem::take(&mut self.call_stack);
                    co.state = CoroutineState::Sleeping;
                    co.wake_tick = wake;
                }
                self.scheduler.sleep_until(cur, wake);
                self.current_coro = None;
                self.scheduler.clear_current();
                switched_out = true;
            } else if self.yield_requested {
                if let Some(co) = self.heap.coroutine_data_mut(cur) {
                    co.pc = self.pc;
                    co.stack = std::mem::take(&mut self.operand_stack);
                    co.frames = std::mem::take(&mut self.call_stack);
                    co.state = CoroutineState::Ready;
                }
                self.scheduler.enqueue_ready(cur);
                self.current_coro = None;
                self.scheduler.clear_current();
                switched_out = true;
            } else if self.halted || self.pc >= self.program.rom.len() {
                if let Some(co) = self.heap.coroutine_data_mut(cur) {
                    co.pc = self.pc;
                    co.stack = std::mem::take(&mut self.operand_stack);
                    co.frames = std::mem::take(&mut self.call_stack);
                    co.state = CoroutineState::Finished;
                }
                self.current_coro = None;
                self.scheduler.clear_current();
                switched_out = true;
            }
        }

        if self.current_coro.is_none() {
            if let Some(next) = self.scheduler.dequeue_next() {
                if let Some(co) = self.heap.coroutine_data_mut(next) {
                    self.pc = co.pc;
                    self.operand_stack = std::mem::take(&mut co.stack);
                    self.call_stack = std::mem::take(&mut co.frames);
                    co.state = CoroutineState::Running;
                }
                self.current_coro = Some(next);
                self.scheduler.set_current(self.current_coro);
            } else if switched_out && !self.scheduler.has_sleeping() {
                self.halted = true;
            }
        } else {
            self.scheduler.set_current(self.current_coro);
        }

        self.yield_requested = false;
    }

    pub fn visit_roots<V: RootVisitor + ?Sized>(&self, visitor: &mut V) {
        for value in &self.operand_stack {
            visit_value_for_roots(value, visitor);
        }

        for frame in &self.call_stack {
            if let Some(func_meta) = self.program.functions.get(frame.func_idx) {
                let start = frame.stack_base;
                let frame_slots =
                    (func_meta.param_slots as usize) + (func_meta.local_slots as usize);
                let mut end = start.saturating_add(frame_slots);
                if end > self.operand_stack.len() {
                    end = self.operand_stack.len();
                }
                for index in start..end {
                    if let Some(value) = self.operand_stack.get(index) {
                        visit_value_for_roots(value, visitor);
                    }
                }
            }
        }

        for global in &self.globals {
            visit_value_for_roots(global, visitor);
        }
    }
}