// Copyright Marcus Del Favero 2025
// Licensed under the GNU AGPLv3 with an exception, see `README.md` for details
use std::{mem, ops::AddAssign};

use glam::{IVec2, Vec2, Vec4, U64Vec3};

use super::image::Image;

#[derive(Clone)]
pub struct Layer<T> {
	width: usize,
	height: usize,
	pixels: Box<[T]>,
}

pub struct Kernel(Box<[f32]>);

impl<T: Default + Clone> Layer<T> {
	pub fn clear(&mut self) {
		self.pixels.fill(T::default());
	}
}

impl<T: Copy> Layer<T> {
	/**
	 * Gets the pixel at the provided position, which is clamped if out of
	 * bounds.
	 */
	fn get(&self, x: isize, y: isize) -> T {
		let (x, y) = (x.clamp(0, self.width as isize - 1) as usize, y.clamp(0, self.height as isize - 1) as usize);
		self.pixels[x + y * self.width]
	}
}

impl<T> Layer<T> {
	/**
	 * Sets the pixel at the provided position, panicking of out of bounds.
	 */
	fn set(&mut self, x: usize, y: usize, val: T) {
		assert!(x < self.width && y < self.height);
		self.pixels[x + y * self.width] = val;
	}

	pub fn flip_vertical(&mut self) {
		for y in 0..self.height / 2 {
			for x in 0..self.width {
				let flipped_y = self.height - 1 - y;
				self.pixels.swap(x + y * self.width, x + flipped_y * self.width);
			}
		}
	}

	pub fn into_pixels(self) -> Vec<T> {
		self.pixels.into_vec()
	}
}

impl<T: AddAssign<T>> Layer<T> {
	pub fn pixel(&mut self, x: usize, y: usize, amount: T) {
		assert!(x < self.width && y < self.height);
		self.pixels[x + y * self.width] += amount;
	}
}

impl Layer<U64Vec3> {
	pub fn new(width: usize, height: usize) -> Layer<U64Vec3> {
		Layer {
			width,
			height,
			pixels: vec![U64Vec3::ZERO; width * height].into_boxed_slice(),
		}
	}

	pub fn point(&mut self, pos: Vec2, amount: U64Vec3) {
		if pos.x < 0.0 || pos.y < 0.0 { return; }
		let (x, y) = (pos.x as usize, pos.y as usize);
		if x >= self.width || y >= self.height { return; }
		self.pixels[x + y * self.width] += amount;
	}

	pub fn add(mut self, other: Layer<U64Vec3>) -> Layer<U64Vec3> {
		assert!(self.width == other.width && self.height == other.height);
		for (a, b) in self.pixels.iter_mut().zip(other.pixels.into_iter()) {
			*a += b;
		}
		self
	}
}

impl Layer<Vec4> {
	pub fn new(width: usize, height: usize, colour: Vec4) -> Layer<Vec4> {
		Layer {
			width,
			height,
			pixels: vec![colour; width * height].into_boxed_slice(),
		}
	}

	/**
	 * Replaces all RGB values with black, keeping the alpha value.
	 */
	pub fn blacken(&mut self) {
		for pixel in &mut self.pixels {
			for i in 0..3 {
				pixel[i] = 0.0;
			}
		}
	}

	pub fn gaussian_blur_kernel(stddev: f32) -> Kernel {
		assert!((0.0..1000.0).contains(&stddev)); // Shouldn't be too big

		let half = (stddev * 3.0).ceil() as usize;
		let size = half * 2 + 1;
		let mut data = Vec::with_capacity(size);

		for i in 0..half {
			let x = (half - i) as f32;
			data.push((-x.powi(2) / (2.0 * stddev.powi(2))).exp());
		}

		let scale = 1.0 / (data.iter().sum::<f32>() * 2.0 + 1.0);
		for v in &mut data { *v *= scale; }

		data.push(scale); // Adds the centre
		for i in (0..half).rev() { data.push(data[i]); }

		let kernel = data.into_boxed_slice();
		assert!(kernel.len() % 2 == 1);
		assert!((kernel.iter().sum::<f32>() - 1.0).abs() < 1e-4);
		Kernel(kernel)
	}

	pub fn gaussian_blur(&mut self, kernel: &Kernel) {
		let mut old = self.clone();

		let off = -(kernel.0.len() as isize / 2);

		// Does this on each axis using a 1D kernel
		// First does the x-axis, then the y-axis
		for y in 0..self.height {
			for x in 0..self.width {
				let mut val = Vec4::ZERO;
				for (i, &k) in kernel.0.iter().enumerate() {
					val += old.get(x as isize + off + i as isize, y as isize) * k;
				}
				self.set(x, y, val);
			}
		}

		mem::swap(&mut old, self);

		for x in 0..self.width {
			for y in 0..self.height {
				let mut val = Vec4::ZERO;
				for (i, &k) in kernel.0.iter().enumerate() {
					val += old.get(x as isize, y as isize + off + i as isize) * k;
				}
				self.set(x, y, val);
			}
		}
	}

	/**
	 * Adds the layer `other` on top of this layer, with the corner of `layer`
	 * being positioned at `offset` in `self`.
	 */
	pub fn add(&mut self, other: &Layer<Vec4>, offset: IVec2) {
		let (off_x, off_y) = (offset.x as isize, offset.y as isize);
		for y in off_y..off_y + other.height as isize {
			for x in off_x..off_x + other.width as isize {
				let (src, dst) = (other.get(x - off_x, y - off_y), self.get(x, y));
				let mut colour = Vec4::ZERO;

				// https://en.wikipedia.org/wiki/Alpha_compositing#Description
				colour.w = src.w + dst.w * (1.0 - src.w);
				for i in 0..3 { colour[i] = (src[i] * src.w + dst[i] * dst.w * (1.0 - src.w)) / colour.w; }
				self.set(x as usize, y as usize, colour);
			}
		}
	}
}

impl Layer<f32> {
	pub fn new(width: usize, height: usize) -> Layer<f32> {
		Layer {
			width,
			height,
			pixels: vec![0.0; width * height].into_boxed_slice(),
		}
	}

	pub fn normalise(&mut self) {
		let max = *self.pixels.iter().max_by(|a, b| a.total_cmp(b)).unwrap();
		for v in &mut self.pixels { *v /= max; }
	}

	pub fn pow(&mut self, x: f32) {
		for v in &mut self.pixels { *v = v.powf(x); }
	}

	pub fn darken(&mut self) {
		let fsize = Vec2::new(self.width as f32, self.height as f32);
		let centre = fsize / 2.0 - 0.5;

		let mut i = 0;
		for y in 0..self.height {
			for x in 0..self.width {
				let dist = ((Vec2::new(x as f32, y as f32) - centre) / (fsize / 2.0)).length();
				self.pixels[i] *= if dist >= 1.0 {
					0.0
				} else {
					1.0 - dist * dist
				};
				i += 1;
			}
		}
	}
}

impl From<Image> for Layer<Vec4> {
	fn from(image: Image) -> Layer<Vec4> {
		assert_eq!(image.pixels.len(), image.width * image.height * 4, "Expected RGBA image");

		let mut pixels = Vec::with_capacity(image.width * image.height);
		for pixel in image.pixels.chunks_exact(4) {
			let mut c = Vec4::ZERO;
			for i in 0..4 { c[i] = pixel[i] as f32 / 255.0; }
			pixels.push(c);
		}

		Layer {
			width: image.width,
			height: image.height,
			pixels: pixels.into_boxed_slice(),
		}
	}
}

impl From<Layer<Vec4>> for Image {
	fn from(layer: Layer<Vec4>) -> Image {
		let mut pixels = Vec::with_capacity(layer.pixels.len() * 4);
		for pixel in &layer.pixels {
			for i in 0..4 { pixels.push(f32_to_u8(pixel[i])); }
		}

		Image {
			width: layer.width,
			height: layer.height,
			pixels: pixels.into_boxed_slice(),
		}
	}
}

impl From<Layer<f32>> for Image {
	fn from(layer: Layer<f32>) -> Image {
		let mut pixels = Vec::with_capacity(layer.pixels.len());
		for pixel in layer.pixels {
			pixels.push(f32_to_u8(pixel));
		}

		Image {
			width: layer.width,
			height: layer.height,
			pixels: pixels.into_boxed_slice(),
		}
	}
}

fn f32_to_u8(x: f32) -> u8 {
	u8::try_from((x * 255.0).round() as i32).expect("Layer values must be between 0 and 1")
}
