use crate::compile::Compile;
use crate::object::*;
use crate::syn::{ast::SpExpr, words::*};
use crate::vm::{error::*, inst::*};
use std::cell::RefCell;
use std::collections::BTreeMap;
use std::rc::Rc;

#[derive(Debug)]
pub struct Frame {
    locals: BTreeMap<Word, Option<Value>>,
    code: Vec<Inst>, // TODO - deduplicate this with some kind of shared pointer
    pc: usize,
}

impl Frame {
    pub fn inst(&self) -> Option<&Inst> {
        self.code.get(self.pc)
    }
}

/// The current state of a VM.
#[derive(Debug)]
pub struct Machine {
    stack: Vec<Value>,
    max_stack_size: Option<usize>,
    arena: Rc<RefCell<Arena>>,
    quote_table: QuoteTable,
    scope_stack: ScopeStack,
    call_stack: Vec<Frame>,
}

impl Machine {
    pub fn new(max_stack_size: Option<usize>, arena: Arena) -> Self {
        Machine {
            stack: Default::default(),
            max_stack_size,
            arena: Rc::new(RefCell::new(arena)),
            quote_table: Default::default(),
            scope_stack: Default::default(),
            call_stack: Default::default(),
        }
    }

    // /////////////////////////////////////////////////////////////////////////
    // Properties
    // /////////////////////////////////////////////////////////////////////////

    pub fn max_stack_size(&self) -> Option<usize> {
        self.max_stack_size
    }

    pub fn stack(&self) -> &Vec<Value> {
        &self.stack
    }

    pub fn stack_mut(&mut self) -> &mut Vec<Value> {
        &mut self.stack
    }

    pub fn stack_push(&mut self, value: Value) -> Result<()> {
        if let Some(max) = self.max_stack_size() {
            if self.stack().len() >= max {
                return Err(RuntimeError::StackOverflow);
            }
        }
        self.stack_mut().push(value);
        Ok(())
    }

    pub fn stack_pop(&mut self) -> Result<Value> {
        self.stack_mut()
            .pop()
            .ok_or_else(|| RuntimeError::StackUnderflow)
    }

    pub fn lookup_word(&self, word: &Word) -> Option<&Value> {
        self.call_stack
            .iter()
            .rev()
            .filter_map(|frame| frame.locals.get(word))
            .next()?
            .as_ref()
    }

    pub fn store_local(&mut self, word: Word, value: Value) {
        let last = self.call_stack.last_mut().expect("no call stack");
        last.locals.insert(word, Some(value));
    }

    pub fn frame(&self) -> Option<&Frame> {
        self.call_stack.last()
    }

    pub fn frame_mut(&mut self) -> Option<&mut Frame> {
        self.call_stack.last_mut()
    }

    fn pc(&self) -> usize {
        self.frame().unwrap().pc
    }

    // /////////////////////////////////////////////////////////////////////////
    // Eval
    // /////////////////////////////////////////////////////////////////////////

    pub fn eval(&mut self, exprs: &Vec<SpExpr>) -> Result<()> {
        self.scope_stack.push_scope();
        let mut compile = Compile::new(&mut self.scope_stack, &mut self.quote_table);
        let code = compile.compile(exprs);
        let locals = self.scope_stack.pop_scope().unwrap();
        self.call_stack.push(Frame {
            locals: locals.keys().map(|w| (*w, None)).collect(),
            code,
            pc: 0,
        });

        // Run instructions
        loop {
            // nothing left to execute
            if self.call_stack.is_empty() {
                break;
            }
            let frame = self.frame().unwrap();
            let inst = frame.inst();
            if let Some(inst) = inst {
                let inst = inst.clone();
                self.eval_inst(inst)?;
            } else {
                self.call_stack.pop();
            }
        }

        Ok(())
    }

    fn eval_inst(&mut self, inst: Inst) -> Result<()> {
        let mut new_frame = None;
        let next_pc = self.pc() + 1;
        match inst {
            Inst::PushValue(value) => self.stack_push(value)?,
            Inst::Load(word) => {
                let value = self
                    .lookup_word(&word)
                    .ok_or_else(|| {
                        RuntimeError::UnsetWord(self.scope_stack.get(&word).to_string())
                    })?
                    .clone();
                self.stack_push(value)?;
            }
            Inst::Store(word) => {
                let value = self.stack_pop()?;
                self.store_local(word, value);
            }
            Inst::Call => {
                let value = self.stack_pop()?;
                let quote = if let Value::Quote(quote) = value {
                    quote
                } else {
                    return Err(RuntimeError::CannotCall(value.name().to_string()));
                };
                let (_span, locals, _expr, code) = self.quote_table.get(quote);

                // create a new stack frame
                new_frame = Some(Frame {
                    locals: locals.keys().map(|w| (*w, None)).collect(),
                    code: code.clone(),
                    pc: 0,
                });
            }
            Inst::Print => {
                let value = self.stack_pop()?;
                println!("{}", value);
            }
        };

        self.frame_mut().unwrap().pc = next_pc;

        // Push the new stack frame *after* updating the PC
        if let Some(frame) = new_frame {
            self.call_stack.push(frame);
        }

        Ok(())
    }
}

#[derive(Debug, Default)]
pub struct MachineBuilder {
    max_stack_size: Option<usize>,
    max_arena_objects: Option<usize>,
}

impl MachineBuilder {
    pub fn max_stack_size(mut self, max_stack_size: Option<usize>) -> Self {
        self.max_stack_size = max_stack_size;
        self
    }

    pub fn max_arena_objects(mut self, max_arena_objects: Option<usize>) -> Self {
        self.max_arena_objects = max_arena_objects;
        self
    }

    pub fn finish(self) -> Machine {
        Machine::new(self.max_stack_size, Arena::new(self.max_arena_objects))
    }
}