pub mod consts;
pub mod error;
pub mod frame;
pub mod inst;
pub mod signal;

use crate::{
    obj::{builtin::BUILTIN_OBJS, reserved::*, prelude::*},
    vm::{consts::ConstPool, frame::*, inst::*, signal::*},
};

#[derive(Debug)]
pub struct Vm<'c> {
    stack: Vec<ObjRef>,
    frames: Vec<Frame>,
    pc: usize,
    condition: bool,
    const_pool: &'c ConstPool,
}

impl<'c> Vm<'c> {
    pub fn new(const_pool: &'c ConstPool,) -> Self {
        Self {
            stack: Default::default(),
            frames: vec![],
            pc: 0,
            condition: false,
            const_pool,
        }
    }

    /// Gets the current stack frame, if any.
    pub fn frame(&self) -> Option<&Frame> {
        self.frames.last()
    }

    /// Gets the current stack frame mutably, if any.
    pub fn frame_mut(&mut self) -> Option<&mut Frame> {
        self.frames.last_mut()
    }

    /// Gets the list of stack frames.
    pub fn frames(&self) -> &Vec<Frame> {
        &self.frames
    }

    /// Mutably gets the list of stack frames.
    pub(crate) fn frames_mut(&mut self) -> &mut Vec<Frame> {
        &mut self.frames
    }

    /// Gets the stack.
    pub fn stack(&self) -> &Vec<ObjRef> {
        &self.stack
    }

    /// Gets the stack, mutably.
    pub fn stack_mut(&mut self) -> &mut Vec<ObjRef> {
        &mut self.stack
    }

    /// Pushes a value to the stack.
    pub fn push(&mut self, value: ObjRef) {
        self.stack_mut().push(value);
    }

    /// Pops a value from the stack.
    pub fn pop(&mut self) -> Option<ObjRef> {
        self.stack_mut().pop()
    }

    /// Gets the current program counter address.
    pub fn pc(&self) -> usize {
        self.pc
    }

    /// Set the next program counter value.
    ///
    /// This may cause the running program to crash. Handle with care.
    pub fn set_pc(&mut self, pc: usize) {
        self.pc = pc;
    }

    /// Gets whether the condition flag has been set or not.
    pub fn condition(&self) -> bool {
        self.condition
    }

    /// Sets the condition flag to the specified value.
    pub fn set_condition(&mut self, condition: bool) {
        self.condition = condition;
    }

    /// Calls a function.
    pub fn call(&mut self, fun: ObjRef, args: Vec<ObjRef>) -> ObjRef {
        self.handle_signal(Signal::Call(fun, args));
        self.resume_until_return();
        self.pop().expect("return value")
    }

    /// Resumes execution of the current program.
    pub fn resume(&mut self) {
        while !self.frames().is_empty() {
            self.resume_until_return();
        }
    }

    /// Resume execution until the currently executing function returns.
    pub fn resume_until_return(&mut self) {
        // resume execution until control is returned to the originally calling function
        let start_frames = self.frames().len();
        if start_frames == 0 {
            return;
        }

        while self.frames().len() >= start_frames {
            let frame = self.frame().unwrap();
            let signal = match frame {
                Frame::Native(fun) => {
                    let callee = fun.callee().clone();
                    let args = fun.args().clone();
                    (*fun.fun_ptr())(callee, self, args)
                }
                Frame::User(_) => {
                    self.resume_user_fun()
                }
            };
            self.handle_signal(signal);
        }
    }

    fn resume_user_fun(&mut self) -> Signal {
        while let Some(Frame::User(_)) = self.frame() {
            if let Some(signal) = self.tick() {
                return signal;
            }
        }
        Signal::Return
    }

    /// Handles a signal targeting the VM.
    ///
    /// The signal may originate from the VM itself, or from an external location. Signals
    /// generally interrupt control flow of a program.
    pub fn handle_signal(&mut self, signal: Signal) {
        match signal {
            Signal::Call(callee, args) => {
                read_obj!(let callee_obj = callee);
                let frame = callee_obj.as_fun()
                    .expect("callable function")
                    .create_frame(callee.clone(), self, args);
                self.frames_mut().push(frame);
                // Jump to the first address of the new function call
                self.set_pc(0);
            }
            Signal::Return => {
                // 1. pop return value
                let ret_val = self.pop().expect("return value");

                // 2. pop stack frame
                let frame = self.frames_mut().pop().expect("stack frame");

                // 3. reset PC
                if let Frame::User(frame) = frame {
                    self.set_pc(frame.last_pc());
                }

                // 4. push return value
                self.push(ret_val);
            }
        }
    }

    /// Gets the current instruction.
    #[inline(always)]
    fn load_inst(&self) -> Inst {
        let frame = if let Frame::User(frame) = self.frame().expect("frame") {
            frame
        } else {
            panic!("expected user function frame")
        };
        frame.code()[self.pc()]
    }

    /// Run a single instruction - inlined version.
    ///
    /// Since this is inlined, it is probably a bad idea to allow users to use it everywhere. The
    /// exposed API function `tick()` calls this function, but does not inline itself (unless the
    /// optimizer thinks it's a good idea to do so).
    #[inline]
    fn tick(&mut self) -> Option<Signal> {
        let inst = self.load_inst();
        let mut next_pc = self.pc() + 1;
        let mut signal = None;
        match inst {
            Inst::PushSym(sym) => {
                let sym_ref = global_sym_ref(sym);
                self.push(sym_ref);
            }
            Inst::PushConst(hdl) => {
                let obj_ref = self.const_pool.get(hdl).clone();
                self.push(obj_ref);
            }
            Inst::LoadLocal(local) => {
                let value = self.get_local(local)
                    .expect("TODO: throw error for missing local");
                self.push(value);
            }
            Inst::LoadGlobal(global) => {
                let value = self.get_global(global)
                    .or_else(|| self.get_builtin(global))
                    .expect("TODO: throw error for missing global");
                self.push(value);
            }
            Inst::PopLocal(name) => {
                let tos = self.pop()
                    .expect("stack underflow");
                // pop into name
                if let Some(name) = name {
                    self.set_local(name, tos);
                }
                // else discard
            }
            Inst::PopGlobal(name) => {
                let tos = self.pop()
                    .expect("stack underflow");
                // pop into name
                if let Some(name) = name {
                    self.set_global(name, tos);
                }
                // else discard
            }
            Inst::GetAttr(sym) => {
                let obj_ref = self.pop().expect("getattr object");
                read_obj!(let obj = obj_ref);
                let attr = obj
                    .get_attr(sym)
                    .unwrap_or_else(|| global_sym_ref(NIL_NAME.sym));
                self.push(attr);
            }
            Inst::SetAttr(sym) => {
                let target = self.pop()
                    .expect("no target available for SetAttr");
                let source = self.pop()
                    .expect("no source available for SetAttr");
                write_obj!(let target = target);
                if let Some(attrs) = target.attrs_mut() {
                    attrs.insert(sym, source);
                } else {
                    todo!("TODO: throw an error for attributes that can't be set");
                }
            }
            Inst::Jump(addr) => {
                next_pc = addr;
            }
            Inst::JumpTrue(addr) => {
                if self.condition {
                    next_pc = addr;
                }
            }
            Inst::Call(argc) => {
                let stack_top = self.stack.len() - argc;
                let args = self.stack.split_off(stack_top);
                let tos = self.pop()
                    .expect("stack underflow");
                read_obj!(let tos = tos);
                let callee = tos.get_call()
                    .expect("TODO: throw an error for missing __call__ attr");
                signal = Some(Signal::Call(callee, args));
            }
            Inst::Index => todo!(),
            Inst::Return => { signal = Some(Signal::Return); }
            Inst::UnNeg => todo!(),
            Inst::UnPos => todo!(),
            Inst::BinPlus => {
                let rhs = self.pop().unwrap();
                let lhs = self.pop().unwrap();
                let fun = {
                    read_obj!(let lhs = lhs);
                    lhs.get_plus().expect("TODO: throw an error for missing __plus__ attr")
                };
                signal = Some(Signal::Call(fun, vec![lhs, rhs]));
            }
            Inst::BinMinus => {
                let rhs = self.pop().unwrap();
                let lhs = self.pop().unwrap();
                let fun = {
                    read_obj!(let lhs = lhs);
                    lhs.get_minus().expect("TODO: throw an error for missing __minus__ attr")
                };
                signal = Some(Signal::Call(fun, vec![lhs, rhs]));
            }
            Inst::BinMul => {
                let rhs = self.pop().unwrap();
                let lhs = self.pop().unwrap();
                let fun = {
                    read_obj!(let lhs = lhs);
                    lhs.get_mul().expect("TODO: throw an error for missing __mul__ attr")
                };
                signal = Some(Signal::Call(fun, vec![lhs, rhs]));
            }
            Inst::BinDiv => todo!(),
            Inst::BinEq => {
                let rhs = self.pop().unwrap();
                let lhs = self.pop().unwrap();
                let fun = {
                    read_obj!(let lhs = lhs);
                    lhs.get_eq().expect("TODO: throw an error for missing __eq__ attr")
                };
                signal = Some(Signal::Call(fun, vec![lhs, rhs]));
            }
            Inst::BinNeq => todo!(),
            Inst::BinLt => todo!(),
            Inst::BinLe => todo!(),
            Inst::BinGt => todo!(),
            Inst::BinGe => todo!(),
            Inst::BinAnd => todo!(),
            Inst::BinOr => todo!(),
        }
        self.set_pc(next_pc);

        signal
    }

    fn get_local(&self, name: Name) -> Option<ObjRef> {
        self.frame()?
            .user_frame()?
            .bindings()
            .get(&name.index())
            .cloned()
    }
    
    fn set_local(&mut self, name: Name, value: ObjRef) {
        let frame = self
            .frame_mut()
            .expect("user stack frame")
            .user_frame_mut()
            .unwrap();
        let bindings = frame.bindings_mut();
        bindings.insert(name.index(), value);
    }

    fn get_global(&self, name: Name) -> Option<ObjRef> {
        self.frames()
            .first()?
            .user_frame()?
            .bindings()
            .get(&name.index())
            .cloned()
    }

    fn set_global(&mut self, name: Name, value: ObjRef) {
        let frame = self.frames_mut()
            .first_mut()
            .expect("global stack frame")
            .user_frame_mut()
            .expect("user-defined function stack frame");
        let bindings = frame.bindings_mut();
        bindings.insert(name.index(), value);
    }

    fn get_builtin(&self, name: Name) -> Option<ObjRef> {
        read_obj!(let fun = self.frames()
            .first()?
            .user_frame()?
            .callee()
        );
        let fun: &UserFun = std::any::Any::downcast_ref(fun.as_any()).unwrap();
        let sym = fun.locals()[name.index()];
        BUILTIN_OBJS.get(&sym).cloned()
    }
}