path: root/src/file.rs

use std::{
	ffi::OsStr,
	path::Path,
	sync::{
		Arc,
		atomic::{AtomicBool, Ordering},
		mpsc::{Receiver, Sender},
	},
	time::{Duration, Instant, SystemTime},
};

use rustc_hash::FxHashMap;
use smol::{Task, block_on, unblock};

use crate::script::{WasmRunner, wasm::WasmModule};

type FileTask = Task<std::io::Result<WasmModule>>;

pub struct FileManager {
	watcher: Notifier,
	updated_files: Receiver<Box<Path>>,
	active_tasks: FxHashMap<Box<Path>, FileTask>,
	cache: FxHashMap<Box<Path>, Arc<std::io::Result<WasmModule>>>,
}

struct Notifier {
	paths: Arc<spin::Mutex<Vec<Box<Path>>>>,
	cancel_token: Arc<AtomicBool>,
}

fn compile_data(path: &Path, data: &[u8], compiler: &WasmRunner) -> std::io::Result<WasmModule> {
	let extension = path
		.extension()
		.and_then(OsStr::to_str)
		.map(|ext| ext.to_lowercase());
	if extension.as_deref() == Some("wasm") {
		return compiler.compile(data).map_err(std::io::Error::other);
	}
	Err(std::io::Error::other("Unrecognized file extension"))
}

impl Notifier {
	pub fn new(channel: Sender<Box<Path>>, interval: Duration) -> Self {
		let paths = Arc::new(spin::Mutex::new(Vec::new()));
		let paths_2 = paths.clone();
		let cancel_token = Arc::new(AtomicBool::new(false));
		let cancel_token_2 = cancel_token.clone();
		std::thread::spawn(move || {
			loop {
				let last_check = SystemTime::now();

				std::thread::sleep(interval);
				if cancel_token_2.load(Ordering::Acquire) {
					break;
				}
				let paths = paths_2.lock().clone();
				for path in paths {
					if let Ok(metadata) = std::fs::metadata(&path)
						&& let Ok(mtime) = metadata.modified()
						&& mtime > last_check
						&& let Err(e) = channel.send(path)
					{
						log::error!("{e}");
					}
				}
			}
		});

		Self {
			paths,
			cancel_token,
		}
	}

	pub fn watch(&mut self, path: Box<Path>) {
		let mut paths = self.paths.lock();
		if paths.contains(&path) {
			return;
		}

		paths.push(path);
	}
}

impl FileManager {
	pub fn new() -> Self {
		let (sender, updated_files) = std::sync::mpsc::channel();
		let watcher = Notifier::new(sender, Duration::from_secs(10));
		Self {
			active_tasks: FxHashMap::default(),
			cache: FxHashMap::default(),
			watcher,
			updated_files,
		}
	}

	pub fn load_file(
		&mut self,
		path: Box<Path>,
		wasm_compiler: &WasmRunner,
	) -> Arc<std::io::Result<WasmModule>> {
		if let Some(task) = self.active_tasks.get(&path)
			&& task.is_finished()
		{
			log::info!("loaded {}", path.to_string_lossy());
			let task = self.active_tasks.remove(&path).unwrap();
			let data = block_on(task);
			self.cache.insert(path.clone(), Arc::new(data));
		} else if !self.cache.contains_key(&path) {
			let read_time = Instant::now();
			let data = std::fs::read(path.as_ref()).map(Vec::into_boxed_slice);
			log::info!("Load time: {} ms", read_time.elapsed().as_millis());
			let compile_time = Instant::now();
			let data = data.and_then(|data| compile_data(path.as_ref(), &data, wasm_compiler));
			log::info!("Compile time: {} ms", compile_time.elapsed().as_millis());
			self.cache.entry(path.clone()).insert_entry(Arc::new(data));
		}

		self.watcher.watch(path.clone());
		self.cache.get(&path).unwrap().clone()
	}

	pub fn start_loading_file(&mut self, path: Box<Path>, wasm_compiler: &WasmRunner) {
		let boxed_path = path.clone();
		let compiler = wasm_compiler.clone();
		let task = unblock(move || {
			compile_data(
				&boxed_path,
				&std::fs::read(&boxed_path).map(Vec::into_boxed_slice)?,
				&compiler,
			)
		});
		self.watcher.watch(path.clone());
		self.active_tasks.insert(path, task);
	}

	pub fn reload_files(&mut self, wasm_compiler: &WasmRunner) {
		while let Ok(path) = self.updated_files.try_recv() {
			log::info!("reloading {}", path.to_string_lossy());
			self.start_loading_file(path, wasm_compiler);
		}
	}

	pub fn loaded_files(&self) -> Box<[Box<Path>]> {
		self.active_tasks
			.iter()
			.filter(|&(_, v)| v.is_finished())
			.map(|(k, _)| k.clone())
			.chain(self.cache.keys().cloned())
			.collect()
	}

	pub fn is_file_loaded(&self, path: impl AsRef<Path>) -> bool {
		self.active_tasks
			.get(path.as_ref())
			.is_some_and(|task| task.is_finished())
			|| self.cache.contains_key(path.as_ref())
	}

	pub fn unload_file(&mut self, path: impl AsRef<Path>) {
		self.active_tasks.remove(path.as_ref());
		self.cache.remove(path.as_ref());
	}
}

impl Default for FileManager {
	fn default() -> Self {
		Self::new()
	}
}

impl Drop for Notifier {
	fn drop(&mut self) {
		self.cancel_token.store(true, Ordering::Release);
	}
}