use std::collections::BTreeMap;

enum Instruction {
	VarArgPrep (i32),
	
	// Get Table, Upvalue
	GetTabUp (u8, u8, u8),
	
	// Get Immediate?
	GetI (u8, u8, u8),
	
	// Equals Constant?
	EqK (u8, u8, u8),
	
	// Jump
	Jmp (i32),
	
	// Load Integer?
	LoadI (u8, i32),
	
	// Load Constant
	LoadK (u8, i32),
	
	
	Call (u8, u8, u8),
	Closure (u8, i32),
	
	// (A, B, _C) Return B - 1 registers starting with  A
	Return (u8, u8, u8),
}

#[derive (Clone, Debug, PartialEq)]
enum Value {
	Nil,
	False,
	True,
	Float (f64),
	Integer (i64),
	String (String),
	
	// These are all bogus, I haven't figured out how to implement
	// tables and function pointers yet
	
	BogusArg (Vec <String>),
	BogusEnv (BTreeMap <String, Value>),
	BogusPrint,
}

impl Default for Value {
	fn default () -> Self {
		Self::Nil
	}
}

impl From <String> for Value {
	fn from (x: String) -> Self {
		Self::String (x)
	}
}

impl From <&str> for Value {
	fn from (x: &str) -> Self {
		Self::from (String::from (x))
	}
}

impl From <i32> for Value {
	fn from (x: i32) -> Self {
		Self::Integer (x as i64)
	}
}

struct Chunk {
	instructions: Vec <Instruction>,
	constants: Vec <Value>,
}

struct VirtualMachine {
	registers: Vec <Value>,
	
	// i32 makes it a little easier to implement jumps
	program_counter: i32,
}

impl Default for VirtualMachine {
	fn default () -> Self {
		Self {
			registers: vec! [Value::Nil; 256],
			program_counter: 0,
		}
	}
}

impl VirtualMachine {
	fn execute_chunk (&mut self, chunk: &Chunk, upvalues: &Vec <Value>) 
	-> Vec <Value> {
		let max_iters = 2000;
		
		for _ in 0..max_iters {
			let instruction = chunk.instructions.get (usize::try_from (self.program_counter).unwrap ()).unwrap ();
			
			let r = &mut self.registers;
			let k = &chunk.constants;
			
			match instruction {
				Instruction::Call (a, b, c) => {
					// Take arguments from registers [a + 1, a + b)
					// Call the function in register [a]
					// Return values in registers [a, a + c - 1)
					// 
					// That is, call a with b - 1 arguments and expect c returns
					// 
					// e.g. CALL 0 2 1 mean "Call 0 with 1 argument, return 1 value", like for printing a constant
					
					// TODO: Only implement printing constants for now
					
					let a = usize::try_from (*a).unwrap ();
					
					assert_eq! (*b, 2);
					assert_eq! (*c, 1);
					
					println! ("{:?}", r [a + 1]);
				},
				Instruction::EqK (a, b, c_k) => {
					let a = usize::try_from (*a).unwrap ();
					let b = usize::try_from (*b).unwrap ();
					
					let equal = r [a] == k [b];
					
					match (equal, c_k) {
						(true, 0) => self.program_counter += 1,
						(false, 1) => self.program_counter += 1,
						_ => (),
					}
				},
				Instruction::GetTabUp (a, b, c) => {
					let a = usize::try_from (*a).unwrap ();
					let b = usize::try_from (*b).unwrap ();
					let c = usize::try_from (*c).unwrap ();
					
					let env = match upvalues.get (b).unwrap () {
						Value::BogusEnv (x) => x,
						_ => panic! ("Only allowed upvalue is BogusEnv"),
					};
					
					let key = match k.get (c).unwrap () {
						Value::String (s) => s,
						_ => panic! ("GetTabUp only supports string keys"),
					};
					
					
					let value = env.get (key).unwrap ();
					
					r [a] = value.clone ();
				},
				Instruction::GetI (a, b, c) => {
					let a = usize::try_from (*a).unwrap ();
					let b = usize::try_from (*b).unwrap ();
					let c = usize::try_from (*c).unwrap ();
					
					let table = r.get (b).unwrap ();
					let value = match table {
						Value::BogusArg (arg) => arg.get (c).map (|x| x.as_str().into ()).unwrap_or_default(),
						_ => unimplemented!(),
					};
					
					r [a] = value;
				},
				Instruction::Jmp (s_j) => self.program_counter += s_j,
				Instruction::LoadI (a, sbx) => {
					let a  = usize::try_from  (*a).unwrap ();
					
					r [a] = (*sbx).into ();
				},
				Instruction::LoadK (a, bx) => {
					let a  = usize::try_from  (*a).unwrap ();
					let bx = usize::try_from (*bx).unwrap ();
					
					r [a] = k [bx].clone ();
				},
				Instruction::Return (a, b, _c) => {
					let a = usize::try_from (*a).unwrap ();
					let b = usize::try_from (*b).unwrap ();
					
					return self.registers [a..(a + b - 1)].to_vec();
				},
				Instruction::VarArgPrep (_) => (),
				_ => (),
			}
			
			self.program_counter += 1;
		}
		
		panic! ("Hit max iterations before chunk returned");
	}
}

fn main() {
	let arg: Vec <_> = std::env::args ().collect ();
	
	let chunk = Chunk {
		instructions: vec! [
			Instruction::VarArgPrep (0),
			Instruction::GetTabUp (1, 0, 0),
			Instruction::GetI (1, 1, 1),
			Instruction::EqK (1, 1, 0),
			Instruction::Jmp (6),
			Instruction::GetTabUp (1, 0, 2),
			Instruction::LoadK (2, 3),
			Instruction::Call (1, 2, 1),
			Instruction::LoadI (1, 0),
			Instruction::Return (1, 2, 1),
			Instruction::Jmp (5),
			Instruction::GetTabUp (1, 0, 2),
			Instruction::LoadK (2, 4),
			Instruction::Call (1, 2, 1),
			Instruction::LoadI (1, 1),
			Instruction::Return (1, 2, 1),
			Instruction::Return (1, 1, 1),
		],
		constants: vec! [
			"arg",
			"93",
			"print",
			"it's 93",
			"it's not 93",
		].into_iter ().map (|s| Value::from (s)).collect (),
	};
	
	let env = BTreeMap::from_iter ([
		("arg", Value::BogusArg (arg.clone ())),
		("print", Value::BogusPrint),
	].map (|(k, v)| (k.to_string (), v)).into_iter ());
	
	let upvalues = vec! [
		Value::BogusEnv (env),
	];
	
	let mut vm = VirtualMachine::default ();
	println! ("Returned: {:?}", vm.execute_chunk (&chunk, &upvalues));
}

#[cfg (test)]
mod tests {
	use super::*;
	
	#[test]
	fn is_93 () {
		/*
		if arg [1] == "93" then
			print "it's 93"
			return 0
		else
			print "it's not 93"
			return 1
		end
		*/
		
		let chunk = Chunk {
			instructions: vec! [
				Instruction::VarArgPrep (0),
				Instruction::GetTabUp (1, 0, 0),
				Instruction::GetI (1, 1, 1),
				Instruction::EqK (1, 1, 0),
				Instruction::Jmp (6),
				Instruction::GetTabUp (1, 0, 2),
				Instruction::LoadK (2, 3),
				Instruction::Call (1, 2, 1),
				Instruction::LoadI (1, 0),
				Instruction::Return (1, 2, 1),
				Instruction::Jmp (5),
				Instruction::GetTabUp (1, 0, 2),
				Instruction::LoadK (2, 4),
				Instruction::Call (1, 2, 1),
				Instruction::LoadI (1, 1),
				Instruction::Return (1, 2, 1),
				Instruction::Return (1, 1, 1),
			],
			constants: vec! [
				"arg",
				"93",
				"print",
				"it's 93",
				"it's not 93",
			].into_iter ().map (|s| Value::from (s)).collect (),
		};
		
		for (arg, expected) in [
			(vec! ["_exe_name"], vec! [1.into ()]),
			(vec! ["_exe_name", "93"], vec! [0.into ()]),
			(vec! ["_exe_name", "94"], vec! [1.into ()]),
		] {
			let arg = arg.into_iter ().map (|s| s.to_string ()).collect ();
			
			let env = BTreeMap::from_iter ([
				("arg", Value::BogusArg (arg)),
				("print", Value::BogusPrint),
			].map (|(k, v)| (k.to_string (), v)).into_iter ());
			
			let upvalues = vec! [
				Value::BogusEnv (env),
			];
			
			let mut vm = VirtualMachine::default ();
			let actual = vm.execute_chunk (&chunk, &upvalues);
			
			assert_eq! (actual, expected);
		}
	}
}