use crate::{
    expressions::{Expr, Value},
    lexing::Lexer,
    parsing::Parser,
    syntax::Punct,
};
use std::{
    collections::{BTreeMap, BTreeSet},
    fmt,
    ops::RangeInclusive,
};

#[cfg(test)]
mod tests;

pub mod compilation;

#[derive(Clone, PartialEq, Debug)]
pub(crate) enum RlValue {
    Unit,
    Addr(usize),
    Num(i32),
    Str(String),
    // A function is multiple instructions at a certain address
    // that `usize` is the "instruction pointer" of the first instruction of the function
    Function(usize, u8),
    Range(RangeInclusive<i32>),
}

impl fmt::Display for RlValue {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        use self::RlValue::*;
        match self {
            Unit => write!(f, "()"),
            Num(x) => write!(f, "{}", x),
            Str(ref s) => write!(f, "'{}'", s),
            _ => write!(f, ".."),
        }
    }
}

#[derive(Clone, PartialEq, Debug)]
pub(crate) enum Instruction {
    // Jumps to the address at the TOS - 1
    Ret,
    // Jumps to the address at the TOS
    Jump,
    Push(RlValue),
    PushLocal(String),
    Declare(String),
    Add,
    Sub,
    Mul,
    Call(String),
    CallNative(String),
    MakeRange,
}

#[derive(Default)]
struct Scope<'l> {
    parent: Option<&'l Scope<'l>>,
    bindings: BTreeMap<&'l str, RlValue>,
}

impl<'t> Scope<'t> {
    pub fn new(parent: &'t Scope) -> Self {
        Scope::init(Some(parent))
    }

    fn init(parent: Option<&'t Scope>) -> Self {
        Scope {
            parent,
            bindings: BTreeMap::new(),
        }
    }

    pub fn empty() -> Self {
        Default::default()
    }

    pub fn lookup(&self, ident: &str) -> Option<&RlValue> {
        self.bindings.get(ident)
    }
}

#[derive(Default)]
pub(crate) struct VM<'l> {
    instr_ptr: usize,
    current_scope: Scope<'l>,
    stack: Vec<RlValue>,
}

macro_rules! num_instr {
    ($self:expr, $op:expr) => {{
        let y = $self.stack.pop().unwrap();
        let x = $self.stack.pop().unwrap();

        match (x, y) {
            (RlValue::Num(x), RlValue::Num(y)) => $self.stack.push(RlValue::Num($op(x, y))),
            _ => unimplemented!(),
        }
    }};
}

impl<'l> VM<'l> {
    pub fn new() -> Self {
        Default::default()
    }

    pub fn lookup_num(&self, ident: &str) -> Option<i32> {
        match self.current_scope.bindings.get(ident) {
            Some(RlValue::Num(x)) => Some(*x),
            _ => None,
        }
    }

    pub fn lookup_str(&self, ident: &str) -> Option<String> {
        match self.current_scope.bindings.get(ident) {
            Some(RlValue::Str(ref s)) => Some(String::from(s)),
            _ => None,
        }
    }

    pub fn lookup_list(&self, ident: &str) -> Option<Vec<i32>> {
        match self.current_scope.bindings.get(ident) {
            Some(RlValue::Range(range)) => Some(range.clone().collect()),
            _ => None,
        }
    }

    fn lookup(&self, ident: &str) -> Option<&RlValue> {
        match self.current_scope.lookup(ident) {
            v @ Some(_) => v,
            _ => match self.current_scope.parent {
                Some(p) => p.lookup(ident),
                _ => None,
            },
        }
    }

    pub fn execute(&mut self, instrs: &'l [Instruction]) {
        while self.instr_ptr < instrs.len() {
            let instr = &instrs[self.instr_ptr];
            self.instr_ptr += 1;

            match instr {
                Instruction::Ret => {
                    // Functions push their result to the TOS,
                    // so in order to return, we need the address at TOS - 1
                    match self.stack.remove(self.stack.len() - 2) {
                        RlValue::Addr(addr) => {
                            self.instr_ptr = addr;
                        }
                        _ => unimplemented!(),
                    }
                }
                Instruction::Jump => match self.stack.pop() {
                    Some(RlValue::Addr(addr)) => self.instr_ptr = addr,
                    _ => unimplemented!(),
                },
                Instruction::Push(v) => self.stack.push(v.clone()),
                Instruction::PushLocal(ref ident) => {
                    self.stack.push(self.lookup(ident).unwrap().clone());
                }
                Instruction::Add => num_instr!(self, |x, y| x + y),
                Instruction::Sub => num_instr!(self, |x, y| x - y),
                Instruction::Mul => num_instr!(self, |x, y| x * y),
                Instruction::Declare(ident) => {
                    self.current_scope
                        .bindings
                        .insert(ident, self.stack.pop().unwrap());
                }
                Instruction::MakeRange => {
                    let to = self.stack.pop().unwrap();
                    let from = self.stack.pop().unwrap();
                    match (from, to) {
                        (RlValue::Num(from), RlValue::Num(to)) => {
                            self.stack.push(RlValue::Range(from..=to))
                        }
                        _ => unimplemented!(),
                    };
                }
                Instruction::Call(ref ident) => {
                    match self.lookup(ident) {
                        Some(RlValue::Function(addr, args)) => {
                            let n = *args;
                            let ptr = self.instr_ptr;

                            self.instr_ptr = *addr;

                            // need to return to the next instruction after the function call
                            // skip the top N args, which are the called function's args
                            self.stack
                                .insert(self.stack.len() - (n as usize), RlValue::Addr(ptr));
                        }
                        _ => {
                            unimplemented!("Function {} does not exist in the current scope", ident)
                        }
                    }
                }
                Instruction::CallNative(ref ident) => match ident.as_ref() {
                    "print" => {
                        let val = self.stack.pop().unwrap();
                        println!("{}", val);
                    }
                    _ => (),
                },
                _ => unimplemented!("{:?}", &instr),
            }
        }
    }
}