use std::collections::BTreeMap;

#[derive (Debug)]
pub enum Instruction {
	Add (u8, u8, u8),
	
	Call (u8, u8, u8),
	Closure (u8, i32),
	
	// Equals Constant?
	EqK (u8, u8, u8),
	
	// Get Table, Upvalue
	GetTabUp (u8, u8, u8),
	
	// Get Immediate?
	GetI (u8, u8, u8),
	
	// Jump
	Jmp (i32),
	
	LoadFalse (u8),
	
	// Load Integer?
	LoadI (u8, i32),
	
	// Load Constant
	LoadK (u8, i32),
	
	LoadTrue (u8),
	
	// MetaMethod, Binary
	MmBin (u8, u8, u8),
	
	Move (u8, u8),
	
	Not (u8, u8),
	
	// (A, B, _C) Return B - 1 registers starting with  A
	Return (u8, u8, u8),
	
	Return0,
	
	// Return just one register
	Return1 (u8),
	
	Test (u8, i32),
	
	VarArgPrep (i32),
}

#[derive (Clone, Debug, PartialEq)]
pub enum Value {
	Nil,
	Boolean (bool),
	Float (f64),
	String (String),
	
	// These are all bogus, I haven't figured out how to implement
	// tables and function pointers yet
	
	BogusArg (Vec <String>),
	BogusClosure (usize),
	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::Float (f64::try_from (x).unwrap ())
	}
}

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

impl Value {
	fn as_float (&self) -> Option <f64> {
		match self {
			Self::Float (x) => Some (*x),
			_ => None,
		}
	}
	
	fn is_truthy (&self) -> bool {
		// And this is something Lua does better than JS and Python.
		
		match self {
			Value::Nil => false,
			Value::Boolean (false) => false,
			_ => true,
		}
	}
}

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

pub struct Chunk {
	pub blocks: Vec <Block>,
}

#[derive (Debug)]
struct StackFrame {
	// i32 makes it a little easier to implement jumps
	// Starts at 0 right after OP_CALL
	
	program_counter: i32,
	
	// Starts from 0 for main and 1 for the first closure
	block_idx: usize,
	
	register_offset: usize,
}

pub struct State {
	registers: Vec <Value>,
	pub debug_print: bool,
}

impl Default for State {
	fn default () -> Self {
		Self {
			registers: vec! [Value::Nil; 256],
			debug_print: false,
		}
	}
}

impl State {
	pub fn upvalues_from_args <I: Iterator <Item = String>> (args: I) -> Vec <Value>
	{
		let arg = args.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)));
		
		vec! [
			Value::BogusEnv (env),
		]
	}
	
	pub fn execute_chunk (&mut self, chunk: &Chunk, upvalues: &[Value]) 
	-> Vec <Value> {
		let max_iters = 2000;
		
		let mut stack = vec! [
			StackFrame {
				program_counter: 0,
				block_idx: 0,
				register_offset: 0,
			},
		];
		
		for _ in 0..max_iters {
			let stack_idx = stack.len () - 1;
			let frame = stack.get_mut (stack_idx).unwrap ();
			let block = chunk.blocks.get (frame.block_idx).unwrap ();
			
			let mut next_pc = frame.program_counter;
			
			let pc = usize::try_from (frame.program_counter).expect ("program_counter is not a valid usize");
			let instruction = match block.instructions.get (pc) {
				Some (x) => x,
				None => {
					dbg! (&stack);
					panic! ("program_counter went out of bounds");
				}
			};
			
			let r = &mut self.registers [frame.register_offset..];
			let k = &block.constants;
			
			match instruction {
				Instruction::Add (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 v_b = r [b].as_float ().unwrap ();
					let v_c = r [c].as_float ().unwrap ();
					
					r [a] = (v_b + v_c).into ();
				},
				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 values for now
					
					let a = usize::try_from (*a).unwrap ();
					let v_a = r.get (a).unwrap ();
					
					match v_a {
						Value::BogusClosure (idx) => {
							let block_idx = frame.block_idx;
							let target_block = idx + 1;
							
							let current_frame = &stack [stack.len () - 1];
							
							stack.push (StackFrame {
								program_counter: 0,
								block_idx: target_block,
								register_offset: current_frame.register_offset + a + 1,
							});
							
							if self.debug_print {
								println! ("Inst {block_idx}:{pc} calls {target_block}:0");
								let stack_depth = stack.len ();
								println! ("stack_depth: {stack_depth}");
							}
							
							// Skip the PC increment at the bottom of the loop
							continue;
						},
						Value::BogusPrint => {
							// In real Lua, print is a function inside
							// the runtime. Here it's bogus.
							
							// assert_eq! (*b, 2);
							assert_eq! (*c, 1);
							
							println! ("{:?}", r.get (a + 1).unwrap ());
						},
						_ => panic! ("Cannot call value {a:?}. backtrace: {stack:?}"),
					}
				},
				Instruction::Closure (a, b) => {
					let a = usize::try_from (*a).unwrap ();
					let b = usize::try_from (*b).unwrap ();
					
					r [a] = Value::BogusClosure (b);
				},
				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) => next_pc += 1,
						(false, 1) => next_pc += 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) => next_pc += s_j,
				Instruction::LoadFalse (a) => {
					let a  = usize::try_from  (*a).unwrap ();
					r [a] = Value::Boolean (false);
				},
				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::LoadTrue (a) => {
					let a  = usize::try_from  (*a).unwrap ();
					r [a] = Value::Boolean (true);
				},
				Instruction::MmBin (a, b, _c) => {
					let a = usize::try_from (*a).unwrap ();
					let b = usize::try_from (*b).unwrap ();
					
					let a = &r [a];
					let b = &r [b];
					
					if a.as_float().is_some() && b.as_float().is_some () {
						// No need for metamethods
					}
					else {
						panic! ("Not sure how to implememtn OP_MMBIN for these 2 values {a:?}, {b:?}");
					}
				},
				Instruction::Move (a, b) => {
					let a = usize::try_from (*a).unwrap ();
					let b = usize::try_from (*b).unwrap ();
					
					r [a] = r [b].clone ();
				},
				Instruction::Not (a, b) => {
					let a = usize::try_from (*a).unwrap ();
					let b = usize::try_from (*b).unwrap ();
					
					r [a] = Value::Boolean (! r [b].is_truthy());
				}
				Instruction::Return (a, b, _c) => {
					let a = usize::try_from (*a).unwrap ();
					let b = usize::try_from (*b).unwrap ();
					
					return r [a..(a + b - 1)].to_vec();
				},
				Instruction::Return1 (a) => {
					let a = usize::try_from (*a).unwrap ();
					let popped_frame = stack.pop ().unwrap ();
					
					self.registers [popped_frame.register_offset - 1] = r [a].clone ();
					
					let stack_idx = stack.len () - 1;
					let frame = stack.get (stack_idx).unwrap ();
					let new_block = frame.block_idx;
					next_pc = frame.program_counter;
					
					if self.debug_print {
						let old_block = popped_frame.block_idx;
						let old_pc = popped_frame.program_counter;
						println! ("Inst {old_block}:{old_pc} returns to inst {new_block}:{next_pc}");
						let stack_depth = stack.len ();
						println! ("stack_depth: {stack_depth}");
					}
				},
				Instruction::Test (a, _k) => {
					let a = usize::try_from (*a).unwrap ();
					
					let a = r.get (a).unwrap ();
					
					if self.debug_print {
						println! ("Test {a:?}");
					}
					
					if a.is_truthy() {
						next_pc += 1;
					}
				},
				Instruction::VarArgPrep (_) => (),
				x => panic! ("Unimplemented instruction {x:?}"),
			}
			
			next_pc += 1;
			{
				let stack_idx = stack.len () - 1;
				let frame = stack.get_mut (stack_idx).unwrap ();
				frame.program_counter = next_pc;
			}
		}
		
		panic! ("Hit max iterations before block returned");
	}
}