xbasic 0.3.2

use crate::basic_io::BasicIO;
use crate::chunk::Chunk;
use crate::error_handler::ErrorHandler;
use crate::expr::ExprValue;
use crate::native_function::NativeFunction;
use crate::opcodes::OpCode;
use num_traits::{FromPrimitive, Pow};

struct CallFrame {
	chunk_id: usize,
	offset: usize,
	ip: usize,
}

impl CallFrame {
	fn new(chunk_id: usize, ip: usize, offset: usize) -> Self {
		Self {
			chunk_id,
			offset,
			ip,
		}
	}
}

// TODO we need a way to clean up the stack when there's an error
// Probably we can just pop off the stack until we're only left with the # of variables?
pub(crate) struct VirtualMachine<T: 'static> {
	ip: usize,
	stack: Vec<ExprValue>,
	pub stdio: T,
	error_handler: ErrorHandler,
	chunks: Vec<Chunk>, // Holds core + function chunks
	cur_chunk: usize,
	call_stack: Vec<CallFrame>,
	offset: usize,
	buffer: String,
	native_functions: Vec<NativeFunction<T>>,
	compute_count: usize,
	compute_limit: usize,
}

impl<T> VirtualMachine<T>
where
	T: BasicIO,
{
	pub(crate) fn new(
		stdio: T,
		native_functions: Vec<NativeFunction<T>>,
		compute_limit: usize,
	) -> Self {
		Self {
			ip: 0,
			stack: Vec::new(),
			stdio,
			error_handler: ErrorHandler::new(),
			chunks: Vec::new(),
			cur_chunk: 0,
			call_stack: Vec::new(),
			offset: 0,
			buffer: String::new(),
			native_functions,
			compute_limit,
			compute_count: 0,
		}
	}

	pub(crate) fn run(
		&mut self,
		chunk: Chunk,
		function_chunks: Vec<Chunk>,
		error_handler: &mut ErrorHandler,
	) -> bool {
		// Allocate space for variables
		for _ in 0..chunk.num_variables {
			self.push(ExprValue::Integer(0))
		}

		self.ip = 0;
		self.cur_chunk = 0;
		self.chunks = vec![chunk];
		self.chunks.extend(function_chunks);
		self.compute_count = 0;

		while self.ip < self.instructions().len() {
			if self.error_handler.had_runtime_error {
				break;
			}
			self.process_instruction()
		}

		error_handler.merge(&self.error_handler);
		!self.error_handler.had_runtime_error
	}

	pub(crate) fn clear_errors(&mut self) {
		self.error_handler = ErrorHandler::new();
	}

	fn process_instruction(&mut self) {
		let opcode = FromPrimitive::from_u8(self.instructions()[self.ip]).unwrap();
		self.ip += 1;

		// Check that we're not at our compute limit
		if self.compute_limit != 0 {
			self.compute_count += 1;
			if self.compute_count >= self.compute_limit {
				self.error_handler
					.runtime_error(self.line(), "Compute limit exceeded.");
				return;
			}
		}

		match opcode {
			OpCode::Pop => {
				self.pop();
			}
			OpCode::Print => {
				let num_values = self.read_byte();
				for _ in 0..num_values {
					let value = self.stack.pop().unwrap();
					self.write_exprvalue(value);
				}
				self.write_newline();
			}
			OpCode::Input => {
				let var_id = self.read_byte() as usize;
				self.stack[var_id] = ExprValue::String(self.stdio.read_line());
			}
			OpCode::Literal8 => {
				let byte = self.read_byte();
				// TODO cloning is very slow - especially if the literal is a string!
				self.stack.push(self.literal(byte as usize).clone())
			}
			OpCode::Literal16 => {
				let word = self.read_word();
				self.stack.push(self.literal(word as usize).clone());
			}
			OpCode::Add => {
				let v2 = self.pop();
				let v1 = self.pop();

				if v1.is_string() || v2.is_string() {
					self.push(ExprValue::String(v1.into_string() + &v2.into_string()));
				} else if v1.is_integer() && v2.is_integer() {
					self.push(ExprValue::Integer(v1.into_integer() + v2.into_integer()));
				} else {
					self.push(ExprValue::Decimal(v1.into_decimal() + v2.into_decimal()));
				}
			}
			OpCode::Sub => {
				let v2 = self.pop();
				let v1 = self.pop();

				if v1.is_string() || v2.is_string() {
					self.error_handler
						.runtime_error(self.line(), "Cannot subtract strings.");
				} else if v1.is_integer() && v2.is_integer() {
					self.push(ExprValue::Integer(v1.into_integer() - v2.into_integer()));
				} else {
					self.push(ExprValue::Decimal(v1.into_decimal() - v2.into_decimal()));
				}
			}
			OpCode::Mul => {
				let v2 = self.pop();
				let v1 = self.pop();

				if v1.is_string() || v2.is_string() {
					self.error_handler
						.runtime_error(self.line(), "Cannot multiply strings.");
				} else if v1.is_integer() && v2.is_integer() {
					self.push(ExprValue::Integer(v1.into_integer() * v2.into_integer()));
				} else {
					self.push(ExprValue::Decimal(v1.into_decimal() * v2.into_decimal()));
				}
			}
			OpCode::Div => {
				let v2 = self.pop();
				let v1 = self.pop();

				if v1.is_string() || v2.is_string() {
					self.error_handler
						.runtime_error(self.line(), "Cannot divide strings.");
				} else {
					self.push(ExprValue::Decimal(v1.into_decimal() / v2.into_decimal()));
				}
			}
			OpCode::Pow => {
				let v2 = self.pop();
				let v1 = self.pop();
				if v1.is_string() || v2.is_string() {
					self.error_handler
						.runtime_error(self.line(), "Cannot pow strings.");
				} else if v1.is_decimal() || v2.is_decimal() {
					let v1 = v1.into_decimal();
					let v2 = v2.into_decimal();
					self.push(ExprValue::Decimal(v1.pow(v2)));
				} else {
					let v1 = v1.into_integer();
					let v2 = v2.into_integer();
					if v2 > u32::MAX as i64 || v2 < 0 {
						self.push(ExprValue::Decimal((v1 as f64).pow(v2 as f64)));
					} else {
						self.push(ExprValue::Integer(v1.pow(v2 as u32)));
					}
				}
			}
			OpCode::Mod => {
				let v2 = self.pop();
				let v1 = self.pop();
				if v1.is_string() || v2.is_string() {
					self.error_handler
						.runtime_error(self.line(), "Cannot mod strings.");
				} else if v1.is_decimal() || v2.is_decimal() {
					let v1 = v1.into_decimal();
					let v2 = v2.into_decimal();
					self.push(ExprValue::Decimal(v1 % v2));
				} else {
					let v1 = v1.into_integer();
					let v2 = v2.into_integer();
					self.push(ExprValue::Integer(v1 % v2));
				}
			}
			OpCode::And => {
				let v2 = self.pop().into_boolean();
				let v1 = self.pop().into_boolean();
				if let Some(v1) = v1 {
					if let Some(v2) = v2 {
						self.push(ExprValue::Boolean(v1 && v2));
						return;
					}
				}
				self.error_handler
					.runtime_error(self.line(), "Cannot cast string to boolean.");
			}
			OpCode::Or => {
				let v2 = self.pop().into_boolean();
				let v1 = self.pop().into_boolean();
				if let Some(v1) = v1 {
					if let Some(v2) = v2 {
						self.push(ExprValue::Boolean(v1 || v2));
						return;
					}
				}
				self.error_handler
					.runtime_error(self.line(), "Cannot cast string to boolean.");
			}
			OpCode::Not => {
				let v = self.pop().into_boolean();
				if let Some(v) = v {
					self.push(ExprValue::Boolean(!v));
					return;
				}
				self.error_handler
					.runtime_error(self.line(), "Cannot cast string to boolean.");
			}
			OpCode::Minus => {
				let v = self.pop();
				if v.is_string() {
					self.error_handler
						.runtime_error(self.line(), "Cannot cast string to number.")
				} else if v.is_integer() {
					self.push(ExprValue::Integer(-v.into_integer()))
				} else {
					self.push(ExprValue::Decimal(-v.into_decimal()))
				}
			}
			OpCode::Var => {
				let var_id = self.read_byte() as usize;
				self.push(self.stack[self.offset + var_id].clone());
			}
			OpCode::SetVar => {
				let var_id = self.read_byte() as usize;
				self.stack[self.offset + var_id] = self.pop();
			}
			OpCode::Equal => {
				let v1 = self.pop();
				let v2 = self.pop();
				if v1.is_string() && v2.is_string() {
					self.push(ExprValue::Boolean(v1.into_string() == v2.into_string()));
				} else if v1.is_string() || v2.is_string() {
					self.error_handler
						.runtime_error(self.line(), "Cannot compare string with non-string");
				} else if v1.is_integer() && v2.is_integer() {
					self.push(ExprValue::Boolean(v1.into_integer() == v2.into_integer()));
				} else {
					self.push(ExprValue::Boolean(
						(v1.into_decimal() - v2.into_decimal()).abs() < 0.000001,
					));
				}
			}
			OpCode::Less => {
				let v2 = self.pop();
				let v1 = self.pop();
				if v1.is_string() || v2.is_string() {
					// Yes, I know this is allowed in qbasic
					self.error_handler
						.runtime_error(self.line(), "Cannot apply equalities to strings");
				} else if v1.is_integer() && v2.is_integer() {
					self.push(ExprValue::Boolean(v1.into_integer() < v2.into_integer()));
				} else {
					self.push(ExprValue::Boolean(v1.into_decimal() < v2.into_decimal()));
				}
			}
			OpCode::LessEqual => {
				let v2 = self.pop();
				let v1 = self.pop();
				if v1.is_string() || v2.is_string() {
					self.error_handler
						.error(self.line(), "Cannot apply equalities to strings");
				} else if v1.is_integer() && v2.is_integer() {
					self.push(ExprValue::Boolean(v1.into_integer() <= v2.into_integer()));
				} else {
					self.push(ExprValue::Boolean(v1.into_decimal() <= v2.into_decimal()));
				}
			}
			OpCode::Greater => {
				let v2 = self.pop();
				let v1 = self.pop();
				if v1.is_string() || v2.is_string() {
					self.error_handler
						.error(self.line(), "Cannot apply equalities to strings");
				} else if v1.is_integer() && v2.is_integer() {
					self.push(ExprValue::Boolean(v1.into_integer() > v2.into_integer()));
				} else {
					self.push(ExprValue::Boolean(v1.into_decimal() > v2.into_decimal()));
				}
			}
			OpCode::GreaterEqual => {
				let v2 = self.pop();
				let v1 = self.pop();
				if v1.is_string() || v2.is_string() {
					self.error_handler
						.error(self.line(), "Cannot apply equalities to strings");
				} else if v1.is_integer() && v2.is_integer() {
					self.push(ExprValue::Boolean(v1.into_integer() >= v2.into_integer()));
				} else {
					self.push(ExprValue::Boolean(v1.into_decimal() >= v2.into_decimal()));
				}
			}
			OpCode::Jne => {
				let condition = self.pop();
				let jmp_amount = self.read_byte() as usize;
				match condition.into_boolean() {
					Some(x) => {
						if !x {
							self.ip += jmp_amount
						}
					}
					None => self
						.error_handler
						.error(self.line(), "Cannot evaluate string as boolean."),
				}
			}
			OpCode::Jmp => {
				self.ip += self.read_byte() as usize;
			}
			OpCode::JmpBack => {
				self.ip -= self.read_byte() as usize;
			}
			OpCode::Inc => {
				let var_id = self.read_byte();
				match self.stack[var_id as usize].increment() {
					Ok(_) => {}
					Err(x) => self.error_handler.error(self.line(), x),
				}
			}
			OpCode::Call => {
				// Check for stack overflow
				// TODO is this too high?
				if self.call_stack.len() > 50000 {
					self.error_handler
						.runtime_error(self.line(), "Stack overflow.");
					return;
				}

				let old_chunk = self.cur_chunk;
				self.cur_chunk = self.read_byte() as usize + 1; // Switch to function chunk
				self.call_stack
					.push(CallFrame::new(old_chunk, self.ip, self.offset));
				self.offset = self.stack.len() - self.arity();
				self.ip = 0; // Reset IP

				// Allocate space on stack for variables
				for _ in 0..self.chunks[self.cur_chunk].num_variables {
					self.push(ExprValue::Integer(0));
				}
			}
			OpCode::Ret => {
				if self.call_stack.is_empty() {
					// We're in the core chunk. Exit gracefully
					return;
				}
				let return_value = self.stack.pop().unwrap();
				// Pop off function variables
				while self.stack.len() > self.offset {
					self.stack.pop().unwrap();
				}
				// Push actual return value
				self.push(return_value);

				// Restore cur chunk
				let new_frame = self.call_stack.pop().unwrap();
				self.ip = new_frame.ip;
				self.cur_chunk = new_frame.chunk_id;
				self.offset = new_frame.offset;
			}
			OpCode::Nil => {
				// TODO a nil type is probably a good idea
				self.push(ExprValue::Integer(0));
			}
			OpCode::Native => {
				let function_id = self.read_byte() as usize;
				let mut args: Vec<ExprValue> = Vec::new();
				for _ in 0..self.native_functions[function_id].arity {
					args.push(self.pop());
				}
				args.reverse();
				let function = &self.native_functions[function_id];
				let result = (function.function)(args, &mut self.stdio);
				self.push(result);
			}
		}
	}

	fn read_byte(&mut self) -> u8 {
		let byte = self.instructions()[self.ip];
		self.ip += 1;
		byte
	}

	fn read_word(&mut self) -> u16 {
		let upper = self.read_byte() as u16;
		let lower = self.read_byte();

		(upper << 8) | (lower as u16)
	}

	fn push(&mut self, value: ExprValue) {
		self.stack.push(value);
	}

	pub(crate) fn pop(&mut self) -> ExprValue {
		self.stack.pop().unwrap()
	}

	fn write_exprvalue(&mut self, value: ExprValue) {
		if !self.buffer.is_empty() {
			self.buffer += " "
		}
		self.buffer += &value.into_string();
	}

	fn write_newline(&mut self) {
		self.stdio.write_line(self.buffer.clone());
		self.buffer = String::new();
	}

	fn line(&self) -> usize {
		self.chunks[self.cur_chunk].lines[self.ip - 1]
	}

	fn instructions(&self) -> &Vec<u8> {
		&self.chunks[self.cur_chunk].instructions
	}

	fn literal(&self, i: usize) -> &ExprValue {
		&self.chunks[self.cur_chunk].literals[i]
	}

	fn arity(&self) -> usize {
		self.chunks[self.cur_chunk].arity
	}
}