// Copyright Marcus Del Favero 2025
// Licensed under the GNU AGPLv3 with an exception, see `README.md` for details
use std::{ops::{Add, Sub, Mul, Div, Rem}, path::PathBuf};

use glam::{UVec2, IVec2};
use rand::SeedableRng;

use super::{Op, BlockMap};

use crate::world::WorldRng;
use crate::utils::rng::Float;

pub struct Vm {
	pub(super) stack: Vec<BlockMap>,
	pub(super) config_dir: PathBuf,
	pub(super) grid_size: UVec2,
	pub(super) block_count: usize,
}

impl Vm {
	pub fn new(grid_size: UVec2, config_path: PathBuf) -> Vm {
		let mut config_dir = config_path;
		config_dir.pop();

		Vm {
			stack: Vec::new(),
			config_dir,
			grid_size,
			block_count: grid_size.element_product() as usize,
		}
	}

	pub fn generate(mut self, ops: &[Op]) -> Result<BlockMap, String> {
		self.run(ops)?;
		self.stack.pop().ok_or_else(|| String::from("cannot return block map, stack is empty"))
	}

	pub(super) fn run(&mut self, ops: &[Op]) -> Result<(), String> {
		for (i, op) in ops.iter().enumerate() {
			self.execute_op(op).map_err(|err| format!("failed executing operator {op:?} at index {i}: {err}"))?;
		}
		Ok(())
	}

	fn execute_op(&mut self, op: &Op) -> Result<(), String> {
		match *op {
			// Binary operators
			Op::Add => self.binary_op(f32::add)?,
			Op::Sub => self.binary_op(f32::sub)?,
			Op::Mul => self.binary_op(f32::mul)?,
			Op::Div => self.binary_op(f32::div)?,
			Op::Rem => self.binary_op(f32::rem)?,
			Op::Min => self.binary_op(f32::min)?,
			Op::Max => self.binary_op(f32::max)?,

			// Generation
			Op::Constant(x) => self.stack.push(vec![x; self.block_count].into_boxed_slice()),
			Op::WhiteNoise { seed } => {
				let mut rng = WorldRng::seed_from_u64(seed);
				self.stack.push((0..self.block_count).map(|_| rng.next_f32()).collect());
			}
			Op::ValueNoise { seed, ref levels } => self.value_noise(seed, levels)?,
			Op::Maze { seed, wall_thickness, open_thickness, open_prob, close_prob } => self.maze(seed, wall_thickness, open_thickness, open_prob, close_prob)?,
			Op::Circle(pos) => {
				let mut bmap = Vec::with_capacity(self.block_count);
				for y in 0..self.grid_size.y {
					for x in 0..self.grid_size.x {
						bmap.push((UVec2::new(x, y).as_vec2() - self.grid_size.as_vec2() / 2.0 + 0.5).distance(pos));
					}
				}
				self.stack.push(bmap.into_boxed_slice());
			}
			Op::Ellipse(pos, ratio) => {
				let mut bmap = Vec::with_capacity(self.block_count);
				for y in 0..self.grid_size.y {
					for x in 0..self.grid_size.x {
						let mut diff = UVec2::new(x, y).as_vec2() - self.grid_size.as_vec2() / 2.0 + 0.5 - pos;
						diff.x *= ratio;
						bmap.push(libm::hypotf(diff.x, diff.y));
					}
				}
				self.stack.push(bmap.into_boxed_slice());
			}
			Op::Rect { mut pos, mut size, fill } => {
				let top = Vm::get_top(&mut self.stack)?;

				for i in 0..2 {
					if pos[i] < 0 {
						size[i] -= pos[i];
						pos[i] = 0;
					}

					if pos[i] + size[i] > self.grid_size[i] as i32 {
						size[i] = self.grid_size[i] as i32 - pos[i];
					}
				}

				if size.cmplt(IVec2::ZERO).any() { return Ok(()); }

				for y in 0..size.y {
					let i_begin = (pos.x + (pos.y + y) * self.grid_size.x as i32) as usize;
					let i_end = i_begin + size.x as usize;
					for i in i_begin..i_end {
						top[i] = fill;
					}
				}
			}
			Op::Gradient(dir) => {
				let mut bmap = Vec::with_capacity(self.block_count);
				for y in 0..self.grid_size.y {
					for x in 0..self.grid_size.x {
						bmap.push(UVec2::new(x, y).as_vec2().dot(dir));
					}
				}
				self.stack.push(bmap.into_boxed_slice());
			},
			Op::Image(ref path) => self.load_image(path)?,

			// Utilities
			Op::FillHoles => self.fill_holes()?,
			Op::Threshold => {
				for x in Vm::get_top(&mut self.stack)? {
					*x = if *x > 0.0 { 1.0 } else { 0.0 };
				}
			}
			Op::FlipHorizontal => {
				let top = Vm::get_top(&mut self.stack)?;
				for y in 0..self.grid_size.y {
					for x in 0..self.grid_size.x / 2 {
						let flipped_x = self.grid_size.x - 1 - x;
						top.swap((x + y * self.grid_size.x) as usize, (flipped_x + y * self.grid_size.x) as usize);
					}
				}
			},
			Op::FlipVertical => {
				let top = Vm::get_top(&mut self.stack)?;
				for y in 0..self.grid_size.y / 2 {
					for x in 0..self.grid_size.x {
						let flipped_y = self.grid_size.y - 1 - y;
						top.swap((x + y * self.grid_size.x) as usize, (x + flipped_y * self.grid_size.x) as usize);
					}
				}
			},
			Op::Shrink(count) => {
				let top = Vm::get_top(&mut self.stack)?;

				if count >= self.grid_size.min_element().div_ceil(2) as usize {
					top.fill(0.0);
					return Ok(());
				}

				let mut prev = top.clone();
				for _ in 0..count {
					let mut i = 0;
					for y in 0..self.grid_size.y {
						for x in 0..self.grid_size.x {
							if x == 0 || x == self.grid_size.x - 1 || y == 0 || y == self.grid_size.y - 1 {
								top[i] = 0.0;
							} else {
								'out: for iy in -1..=1i32 {
									for ix in -1..=1i32 {
										if prev[(i as i32 + ix + iy * self.grid_size.x as i32) as usize] <= 0.0 {
											top[i] = 0.0;
											break 'out;
										}
									}
								}
							}

							i += 1;
						}
					}

					prev.copy_from_slice(top);
				}
			},

			// Stack operations
			Op::Duplicate(index) => self.stack.push(self.stack[self.get_index(index)?].clone()),
			Op::Remove(index) => _ = self.stack.remove(self.get_index(index)?),
		}

		Ok(())
	}

	fn binary_op(&mut self, op: fn(f32, f32) -> f32) -> Result<(), String> {
		let stack_size = self.stack.len();
		let (rhs, ret) = self.stack.pop().zip(self.stack.last_mut()).ok_or_else(|| format!("expected 2 operands, got {stack_size}"))?;

		// Hopefully optimises the below indexing
		assert!(rhs.len() == self.block_count && ret.len() == self.block_count);

		for i in 0..self.block_count {
			ret[i] = op(ret[i], rhs[i]);
		}

		Ok(())
	}

	pub(super) fn get_top(stack: &mut [BlockMap]) -> Result<&mut BlockMap, String> {
		stack.last_mut().ok_or_else(|| String::from("expected operand"))
	}

	/**
	 * Gets an index the same way the stack in Lua is indexed in the C API. That
	 * is, non-negative indices are absolute and negative indices are relative
	 * to the top.
	 */
	fn get_index(&self, index: isize) -> Result<usize, String> {
		let abs_index = (if index >= 0 { index } else { index + self.stack.len() as isize }) as usize;
		if abs_index < self.stack.len() { Ok(abs_index) } else { Err(format!("index {index} is out of bounds with stack size of {}", self.stack.len())) }
	}
}
