use heck::ToSnakeCase as _;
use resvg::{tiny_skia, usvg};
use std::{
	collections::BTreeMap,
	env,
	fs::{self, File},
	hash::{DefaultHasher, Hash, Hasher},
	io::{self, Write as _},
	path::{Path, PathBuf}
};

const TILE_SIZE: u32 = 64;

#[derive(Default)]
struct Assets {
	/// Assets directly contained within this asset group, mapping their name to the name
	/// of the constant storing their png.
	assets: BTreeMap<String, String>,

	/// Sound assets.
	sound_assets: BTreeMap<String, String>,

	/// Asset groups contained within  this asset group, mapping their name to the assets
	/// that group contains.
	groups: BTreeMap<String, Box<Assets>>
}

struct AssetsWriter {
	file: File,
	root: Assets
}

impl AssetsWriter {
	fn new<P: AsRef<Path>>(path: P) -> Self {
		let mut file = File::create(path).expect("Failed to create assets file");
		writeln!(file, "// @generated").unwrap();
		Self {
			file,
			root: Assets::default()
		}
	}

	fn asset_name<P: AsRef<Path>>(canonical_path: P) -> String {
		let mut hasher = DefaultHasher::new();
		canonical_path.as_ref().hash(&mut hasher);
		let hash = hasher.finish();
		format!("ASSET_{hash:X}")
	}

	fn add_png<P: AsRef<Path> + Copy>(
		&mut self,
		canonical_path: P,
		png: tiny_skia::Pixmap
	) -> String {
		let const_name = Self::asset_name(canonical_path);

		let out_dir = env::var_os("OUT_DIR").unwrap();
		let out_dir: PathBuf = out_dir.into();
		png.save_png(out_dir.join(format!("{const_name}.png")))
			.expect("Failed to save png");

		writeln!(self.file, "// {}", canonical_path.as_ref().display()).unwrap();
		writeln!(
			self.file,
			"const {const_name}: &[u8] = include_bytes!(\"{const_name}.png\");"
		)
		.unwrap();

		const_name
	}

	fn add_ogg<P: AsRef<Path> + Copy>(&mut self, canonical_path: P) -> String {
		let const_name = Self::asset_name(canonical_path);
		writeln!(self.file, "// {}", canonical_path.as_ref().display()).unwrap();
		writeln!(
			self.file,
			"const {const_name}: &[u8] = include_bytes!(\"{}\");",
			canonical_path.as_ref().display()
		)
		.unwrap();
		const_name
	}

	fn finish(mut self) {
		fn write_assets_struct(
			file: &mut File,
			root: &Assets,
			indent: &str
		) -> io::Result<()> {
			for (group_name, group) in &root.groups {
				writeln!(file, "{indent}mod {group_name} {{")?;
				writeln!(file, "{indent}\t#[allow(clippy::wildcard_imports)]")?;
				writeln!(file, "{indent}\tuse super::*;")?;
				write_assets_struct(file, group, &format!("{indent}\t"))?;
				writeln!(file, "}}")?;
			}

			writeln!(file, "{indent}#[allow(dead_code)]")?;
			writeln!(file, "{indent}pub struct Assets {{")?;
			for asset_name in root.assets.keys() {
				writeln!(
					file,
					"{indent}\tpub {}: comfy::TextureHandle,",
					asset_name.to_snake_case()
				)?;
			}
			for asset_name in root.sound_assets.keys() {
				writeln!(
					file,
					"{indent}\tpub {}: comfy::Sound,",
					asset_name.to_snake_case()
				)?;
			}
			for group_name in root.groups.keys() {
				writeln!(
					file,
					"{indent}\tpub {group_name}: &'static {group_name}::Assets,"
				)?;
			}
			writeln!(file, "{indent}}}")?;

			writeln!(file, "{indent}impl Assets {{")?;
			writeln!(
				file,
				"{indent}\tpub fn load(_ctx: &mut comfy::EngineContext<'_>) {{"
			)?;
			for asset_const_name in root.assets.values() {
				writeln!(file, "{indent}\t\t_ctx.load_texture_from_bytes({asset_const_name:?}, {asset_const_name});")?;
			}
			for asset_const_name in root.sound_assets.values() {
				writeln!(file, "{indent}\t\tcomfy::load_sound_from_bytes({asset_const_name:?}, {asset_const_name},")?;
				writeln!(
					file,
					"{indent}\t\t\tcomfy::StaticSoundSettings::new().loop_region(..));"
				)?;
			}
			for group_name in root.groups.keys() {
				writeln!(file, "{indent}\t\t{group_name}::Assets::load(_ctx);")?;
			}
			writeln!(file, "{indent}\t}}")?;
			writeln!(file, "{indent}}}")?;

			writeln!(
				file,
				"{indent}pub static ASSETS: comfy::Lazy<Assets> = comfy::Lazy::new(|| Assets {{"
			)?;
			for (asset_name, asset_const_name) in &root.assets {
				writeln!(
					file,
					"{indent}\t{}: comfy::texture_id({asset_const_name:?}),",
					asset_name.to_snake_case()
				)?;
			}
			for (asset_name, asset_const_name) in &root.sound_assets {
				writeln!(
					file,
					"{indent}\t{}: comfy::sound_id({asset_const_name:?}),",
					asset_name.to_snake_case()
				)?;
			}
			for group_name in root.groups.keys() {
				writeln!(
					file,
					"{indent}\t{group_name}: comfy::Lazy::force(&{group_name}::ASSETS),"
				)?;
			}
			writeln!(file, "{indent}}});")?;

			Ok(())
		}

		write_assets_struct(&mut self.file, &self.root, "").unwrap();
	}
}

fn main() {
	println!("cargo::rerun-if-changed=build.rs");
	println!("cargo::rerun-if-changed=assets/");

	let out_dir = env::var_os("OUT_DIR").unwrap();
	let assets = PathBuf::from(out_dir).join("assets.rs");
	println!("cargo::warning=Writing assets to {}", assets.display());
	println!("cargo::rustc-env=ASSETS_RS={}", assets.display());
	let mut writer = AssetsWriter::new(assets);
	process_dir("assets", &mut writer, &mut Vec::new());
	writer.finish();
}

fn process_dir<P: AsRef<Path> + Copy>(
	dir: P,
	writer: &mut AssetsWriter,
	groups: &mut Vec<String>
) {
	for entry in fs::read_dir(dir).expect("Cannot read dir") {
		let entry = entry.expect("Cannot read dir entry");
		let path = entry.path();
		if entry
			.metadata()
			.expect("Cannot read dir entry metadata")
			.is_dir()
		{
			groups.push(
				path.file_name()
					.unwrap()
					.to_str()
					.expect("Non-UTF8 file names aren't allowed")
					.into()
			);
			process_dir(&path, writer, groups);
			groups.pop();
		} else if path.extension().map(|ext| ext == "svg").unwrap_or(false) {
			process_svg(&path, dir, writer, groups);
		} else if path.extension().map(|ext| ext == "ogg").unwrap_or(false) {
			process_ogg(&path, writer, groups);
		}
	}
}

fn process_svg<P: AsRef<Path> + Copy, Q: AsRef<Path>>(
	file: P,
	dir: Q,
	writer: &mut AssetsWriter,
	groups: &[String]
) {
	let bytes = fs::read(file).expect("Cannot read svg file");
	let tree = usvg::Tree::from_data(&bytes, &usvg::Options {
		resources_dir: Some(dir.as_ref().to_owned()),
		..Default::default()
	})
	.unwrap_or_else(|err| {
		panic!("Cannot parse svg file {}: {err}", file.as_ref().display())
	});
	let mut pixmap = tiny_skia::Pixmap::new(TILE_SIZE, TILE_SIZE).unwrap();
	let transform = tiny_skia::Transform::from_scale(
		TILE_SIZE as f32 / tree.size().width(),
		TILE_SIZE as f32 / tree.size().height()
	);
	resvg::render(&tree, transform, &mut pixmap.as_mut());

	let const_name = writer.add_png(
		&file
			.as_ref()
			.canonicalize()
			.expect("Failed to canonicalize"),
		pixmap
	);
	let mut group = &mut writer.root;
	for group_name in groups {
		if !group.groups.contains_key(group_name) {
			group.groups.insert(group_name.to_owned(), Box::default());
		}
		group = group.groups.get_mut(group_name).unwrap();
	}
	group.assets.insert(
		file.as_ref()
			.file_stem()
			.expect("File doesn't have a stem")
			.to_str()
			.expect("Non-UTF8 file names aren't allowed")
			.into(),
		const_name
	);
}

fn process_ogg<P: AsRef<Path> + Copy>(
	file: P,
	writer: &mut AssetsWriter,
	groups: &[String]
) {
	let const_name = writer.add_ogg(
		&file
			.as_ref()
			.canonicalize()
			.expect("Failed to canonicalize")
	);
	let mut group = &mut writer.root;
	for group_name in groups {
		if !group.groups.contains_key(group_name) {
			group.groups.insert(group_name.to_owned(), Box::default());
		}
		group = group.groups.get_mut(group_name).unwrap();
	}
	group.sound_assets.insert(
		file.as_ref()
			.file_stem()
			.expect("File doesn't have a stem")
			.to_str()
			.expect("Non-UTF8 file names aren't allowed")
			.into(),
		const_name
	);
}