//! Provide the virtual machine

use std::{
    collections::{HashMap, HashSet},
    ops::Deref,
};

use gc_arena::{Arena, Collect, Gc, Rootable};

use crate::{
    bytecode::{ByteCode, OpCode},
    value::{ConsCell, Primitive, SymbolId, SymbolTable},
};

type Memory = Rootable![Vec<StackFrame<'_>>];

/// Virtual machine
pub struct Vm {
    mem: Arena<Memory>,
    /// Main block for execution
    block: Block,
    // NOTE There should be a procedure type that
    // can also execute bytecode, and internally uses its own VM
    // TODO Handling VM top-level returns
}

/// Block of bytecode that supports multiple labels
pub struct Block {
    code: Box<[ByteCode]>,
    labels: HashMap<SymbolId, usize>,
    /// Symbol table for all bytecode labels in this block
    symbol_table: SymbolTable,
    /// breakpoints used to stop code execution in debug mode
    breakpoints: HashSet<usize>,
}

impl Block {
    /// Get the block offset of a specific label
    pub fn get_offset<S: AsRef<str>>(&self, label: S) -> Option<usize> {
        self.symbol_table
            .id(label)
            .and_then(|sym| self.labels.get(&sym).copied())
    }

    /// The symbol table of block labels
    pub fn symbols(&self) -> &SymbolTable {
        &self.symbol_table
    }

    /// Breakpoint access
    pub fn breakpoints(&self) -> &HashSet<usize> {
        &self.breakpoints
    }

    /// Mutable breakpoint access
    pub fn breakpoints_mut(&mut self) -> &mut HashSet<usize> {
        &mut self.breakpoints
    }

    fn get(&self, idx: usize) -> Option<ByteCode> {
        self.code.get(idx).copied()
    }
}

impl Deref for Block {
    type Target = [ByteCode];
    fn deref(&self) -> &Self::Target {
        &self.code
    }
}

/// Meant to represent a full stack frame
#[derive(Debug, Collect)]
#[collect(no_drop)]
pub struct StackFrame<'gc> {
    /// The stack of a stack frame
    pub stack: Vec<Gc<'gc, Primitive<'gc>>>,
    /// The SRL registers, mostly used to create self-referential
    /// datastructures
    pub srl_registers: [Option<Gc<'gc, ConsCell<'gc>>>; 255],
    /// The instruction pointer of a stack frame
    pub inst_ptr: usize,
}

impl Vm {
    /// Create a new VM starting at instruction 0
    pub fn new(block: Block) -> Self {
        Self::with_start(block, 0)
    }

    /// Create a new VM starting at a certain location
    pub fn with_start(block: Block, start: usize) -> Self {
        let mem = Arena::new(|_mc| {
            vec![StackFrame {
                stack: vec![],
                srl_registers: [None; 255],
                inst_ptr: start,
            }]
        });

        Self { mem, block }
    }

    /// Access the block of code this VM will execute
    pub fn block(&self) -> &Block {
        &self.block
    }

    /// Access the internal data of this VM
    pub fn internals(&self) -> &Arena<Memory> {
        &self.mem
    }

    /// Execute from the given stack frame, where all values that would be returned from the
    /// top-level of the script are passed to a "continuation"
    ///
    /// See [`Primitive::copy_to_arena`] for using VM data *outside* the VM.
    #[allow(clippy::result_unit_err)]
    pub fn execute<C>(&mut self, debug_mode: bool, continuation: C) -> Result<(), ()>
    where
        C: FnMut(Primitive),
    {
        // If `debug_mode`, respect breakpoints
        _ = debug_mode;
        _ = continuation;
        // TODO Move blocks and this loop to a "bytecode procedure"
        // which should just be a kind of procedure that is storable on the stack,
        // making room for unified procedure handling
        loop {
            let Some((inst, inst_ptr)) = self.mem.mutate(|_mc, root| {
                root.last()
                    .and_then(|frame| self.block.get(frame.inst_ptr).map(|i| (i, frame.inst_ptr)))
            }) else {
                // Invalid bytecode location. ERROR!!!
                // TODO Have an error type that can capture a VM
                // stack trace for debugging
                break Err(());
            };

            let mut next_ptr = inst_ptr + 1;
            match inst.op_code {
                OpCode::NoOp => {}
                OpCode::PushInt => todo!(),
                OpCode::PushSignedInt => todo!(),
                OpCode::PushConst => todo!(),
                OpCode::PushList => todo!(),
                OpCode::CreateList => todo!(),
                OpCode::Call => {
                    let a = inst.arguments.a() as usize;
                    let Some(jump_ptr) = self.mem.mutate_root(|_mc, root| {
                        root.last_mut().and_then(|frame| {
                            if a > frame.stack.len() {
                                None
                            } else {
                                let arguments: Vec<_> = frame
                                    .stack
                                    .drain(frame.stack.len().saturating_sub(a)..)
                                    .collect();
                                let proc = frame.stack.pop();
                                proc.map(|prim| {
                                    _ = prim;
                                    _ = arguments;
                                    // If the inner primitive is a raw procedure (a procedure written in Rust),
                                    // run it.
                                    // If the inner primtive is a "bytecode reference", jump to it's label by returning
                                    // Some(label location)
                                    // NOTE A continuation *replaces the information about a stack frame*,
                                    // then returns Some(continuation ptr)
                                    todo!()
                                })
                            }
                        })
                    }) else {
                        // jump to location not found.
                        break Err(());
                    };
                    next_ptr = jump_ptr;
                }
                OpCode::TailCall => todo!(),
                OpCode::Return => todo!(),
                OpCode::LoadLocal => todo!(),
                OpCode::StoreLocal => todo!(),
                OpCode::PushCurrentContinuation => todo!(),
                OpCode::JumpIfFalse => todo!(),
            }
            // Set the current frame's inst_ptr to next_ptr
            if !self.mem.mutate_root(|_mc, root| {
                if let Some(last) = root.last_mut() {
                    last.inst_ptr = next_ptr;
                    true
                } else {
                    false
                }
            }) {
                // The stack frame doesn't exist to continue executing...
                // For now, assume that code is done, and exit with an Ok!
                break Ok(());
            }
        }
    }
}