pub mod consts;
pub mod error; // TODO : not needed?
pub mod frame;
pub mod inst;
pub(crate) mod signal;

use crate::{
    obj::{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
    }

    /// 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;
    }

    /// Resumes execution of the current program.
    pub fn resume(&mut self) {
        while let Some(frame) = self.frame().cloned() {
            let signal = match frame.kind() {
                FrameKind::Native(fun, args) => {
                    read_obj!(let fun_obj = fun);
                    (fun_obj.fun())(fun.clone(), self, args.clone())
                }
                FrameKind::User { .. } => {
                    // run the user function until it returns control to the VM
                    loop {
                        if let Some(signal) = self.tick() {
                            break signal;
                        }
                    }
                }
            };
            self.handle_signal(signal);
        }
    }

    /// 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(caller, args) => {
                self.begin_call(caller, args);
            }
            Signal::Return => {
                let retval = self.stack.pop()
                    .expect("return value");
                match self.frames.pop().unwrap().kind() {
                    // no-op; return value should be the TOS
                    FrameKind::Native(_, _) => {  },
                    FrameKind::User { last_pc, stack_base, fun: _, } => {
                        // pop return value, reset PC, clean stack, and push return value
                        self.set_pc(*last_pc);
                        self.stack.truncate(*stack_base);
                    }
                }
                self.stack.push(retval);
            }
        }
    }

    /// Sets up the stack for a function call.
    fn begin_call(&mut self, caller: ObjRef, args: Vec<ObjRef>) {
        // create stack frame
        let stack_frame = if let Some(user_fun) = std::any::Any::downcast_ref::<UserFunRef>(&caller) {
            let names: FrameLocals = {
                read_obj!(let fun_ref = user_fun);
                fun_ref.locals()
                    .iter()
                    .enumerate()
                    .zip(args.into_iter())
                    .map(|((index, _), arg)| (index, arg.clone()))
                    .collect()
            };
            // TODO : check function arity vs argument count

            Frame::new(names, FrameKind::User {
                last_pc: self.pc(),
                stack_base: self.stack().len(),
                fun: user_fun.clone(),
            })
        } else if let Some(native_fun) = std::any::Any::downcast_ref::<NativeFunRef>(&caller) {
            Frame::new(Default::default(), FrameKind::Native(native_fun.clone(), args))
        } else {
            // TODO : throw an error when error handling is figured out
            panic!("can't call object {:?}", caller);
        };

        // push a new stack frame
        self.frames.push(stack_frame);
    }

    /// Gets the current instruction.
    #[inline(always)]
    fn load_inst(&self) -> Inst {
        let user_fun = if let FrameKind::User { fun, .. } = self.frame().expect("frame").kind() {
            fun
        } else {
            panic!("expected user function frame")
        };
        read_obj!(let user_fun = user_fun);
        user_fun.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 obj_ref = self.frames()
                    .iter()
                    .rev()
                    .filter_map(|frame| frame.locals().get(&local.index()))
                    .cloned()
                    .next();
                if let Some(obj_ref) = obj_ref {
                    self.push(obj_ref);
                } else {
                    todo!("TODO: implement \"local value not found\" lookup")
                }
            }
            Inst::LoadGlobal(global) => {
                let obj_ref = self.frames()
                    .first()
                    .and_then(|frame| frame.locals().get(&global.index()))
                    .cloned();
                if let Some(obj_ref) = obj_ref {
                    self.push(obj_ref);
                } else {
                    todo!("TODO: implement \"local value not found\" lookup")
                }
            }
            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 set_local(&mut self, name: Name, value: ObjRef) {
        let frame = self.frame_mut().unwrap();
        let locals = frame.locals_mut();
        locals.insert(name.index(), value);
    }

    fn set_global(&mut self, name: Name, value: ObjRef) {
        let frame = self.frames.first_mut().unwrap();
        let locals = frame.locals_mut();
        locals.insert(name.index(), value);
    }
}