// Copyright Marcus Del Favero 2025
// Licensed under the GNU AGPLv3 with an exception, see `README.md` for details
use std::{f32::consts::{TAU, SQRT_2}, thread, sync::mpsc::{self, Sender}, io::{self, Write}, collections::VecDeque};

use rand::{RngCore, SeedableRng};
use rand_chacha::ChaCha8Rng;
use glam::{Vec2, Vec3, U64Vec3};
use png::ColorType;

use super::{image::Image, layer::Layer};

use crate::utils;

const PATH: &str = "../data/textures/checkpoint.png";
const SIZE: usize = 512;
const ITERATIONS: u64 = 10_000_000_000; // The more iterations, the less noisy the image, but the longer generation takes
const PROGRESS_PERIOD: u64 = 5_000_000; // Number of iterations before updating the progress
const JOBS: usize = 4;

fn job(index: usize, mut iterations: u64, progress: Sender<u64>) -> Layer<U64Vec3> {
	const X: f32 = 0.875;
	const Y: f32 = 1.0 / SQRT_2;
	const SQRT_3: f32 = 1.7320508;

	const COLOUR_SCALE: f32 = 1.75;
	const COLOUR_BASE: f32 = 4096.0;
	const HISTORY_SIZE: usize = 6;

	let colour_average = Vec3::splat(COLOUR_BASE * (0..HISTORY_SIZE).map(|i| COLOUR_SCALE.powi(i as i32)).sum::<f32>() / 3.0);
	let colour_average_int = colour_average.as_u64vec3();

	let triangle_points = [
		Vec2::new(-0.5 * X, -SQRT_3 / 4.0 * X),
		Vec2::new(0.5 * X, -SQRT_3 / 4.0 * X),
		Vec2::new(0.0, SQRT_3 / 4.0 * X),
	];

	let mut layer = Layer::<U64Vec3>::new(SIZE, SIZE);
	let mut rng = ChaCha8Rng::seed_from_u64(641892749 + index as u64);

	let mut point = triangle_points[0];
	let mut stack = Vec::new();
	let mut history = VecDeque::new();
	let mut chose_circle = false;

	while iterations > 0 {
		let count = iterations.min(PROGRESS_PERIOD);
		iterations -= count;
		for _ in 0..count {
			// Reusing the same randomly generated number for performance
			let r = rng.next_u32();
			let amount = if r & 0xc0000000 != 0 || chose_circle {
				chose_circle = false;

				if let Some(p) = stack.pop() {
					point = p;
				} else {
					let i = r % 3;
					let next_point = triangle_points[i as usize];
					let f = r & 0x20000000 == 0;
					point = point.lerp(next_point, if f { 0.525 } else { -0.025 });

					if f {
						history.push_back(i);
						if history.len() > HISTORY_SIZE {
							history.pop_front();
						}
					}
				}

				let mut amount = Vec3::ZERO;
				let mut scl = COLOUR_BASE;
				for &i in &history {
					amount[i as usize] += scl;
					scl *= COLOUR_SCALE;
				}
				amount.as_u64vec3()
			} else if r & 0xf0000000 != 0 {
				stack.push(point);
				let next_point = match r & 0x7 {
					0 => Vec2::new(-Y, -Y),
					1 => Vec2::new(Y, -Y),
					2 => Vec2::new(-Y, Y),
					3 => Vec2::new(Y, Y),

					4 => Vec2::new(1.0, 0.0),
					5 => Vec2::new(-1.0, 0.0),
					6 => Vec2::new(0.0, 1.0),
					_ => Vec2::new(0.0, -1.0),
				};

				point = point.lerp(next_point, 0.75);

				let mut amount = Vec3::ZERO;
				let mut scl = COLOUR_BASE;
				for &i in &history {
					amount[i as usize] += scl;
					scl *= COLOUR_SCALE;
				}
				colour_average.lerp(amount, 1.0 / (stack.len() + 1).pow(2) as f32).as_u64vec3()
			} else {
				chose_circle = true;
				stack.push(point);
				let angle = r as f32 * const { TAU / 0x0fffffff as f32 };
				let next_point = Vec2::from_angle(angle);
				point = point.lerp(next_point, 0.625);
				colour_average_int
			};

			layer.point((point + Vec2::ONE) / 2.0 * SIZE as f32, amount);
		}
		progress.send(count).unwrap();
	}

	layer
}

pub fn generate() {
	if utils::exists(PATH) { return; }

	let mut threads = Vec::with_capacity(JOBS);
	let iterations_per_job = ITERATIONS / JOBS as u64;
	assert_eq!(ITERATIONS % JOBS as u64, 0);
	let (sender, receiver) = mpsc::channel();
	for i in 0..JOBS {
		let sender = sender.clone();
		threads.push(thread::spawn(move || job(i, iterations_per_job, sender)));
	}
	drop(sender);

	let mut iterations = 0;
	while let Ok(count) = receiver.recv() {
		iterations += count;
		let percent = iterations as f32 / ITERATIONS as f32 * 100.0;
		print!("\x1b[2K\r   Progress: {percent:.2}%");
		let _ = io::stdout().flush();
	}
	println!();

	let pixels = threads.into_iter().map(|th| th.join().unwrap()).reduce(Layer::<U64Vec3>::add).unwrap().into_pixels();
	let max = pixels.iter().copied().map(U64Vec3::max_element).max().unwrap() as f32;

	Image {
		width: SIZE,
		height: SIZE,
		pixels: pixels.into_iter().flat_map(|px| {
			((px.as_vec3() / max).powf(0.375) * 255.0).to_array().map(|x| x as u8).into_iter()
		}).collect(),
	}.save(PATH, ColorType::Rgb);
}
