export module oar;

import coral;
import coral.files;
import coral.functional;
import coral.io;
import coral.stack;

using coral::closure;
using coral::expected;
using coral::file_reader;
using coral::fs;
using coral::io_error;
using coral::path;
using coral::u8;
using coral::u64;
using coral::usize;

/**
 * Length of the full magic signature at the beginning of an Oar file.
 */
constexpr static usize signature_length {4};

/**
 * Hardcoded signature magic value that this implementation of Oar expects when reading archives.
 */
constexpr static u8 signature_magic[signature_length] {'o', 'a', 'r', 1};

/**
 * Oar file header format.
 */
union header {
	struct {
		u8 signature[signature_length];

		coral::u32 entry_count;
	} layout;

	u8 bytes[512];

	static constexpr bool is_sizeof(usize value) {
		return value == sizeof(header);
	}

	expected<bool, io_error> read(coral::reader & archive_reader) {
		return archive_reader.read(this->bytes).map<bool>(is_sizeof).map<bool>([&](bool is_valid) -> bool {
			return is_valid && coral::equals(this->layout.signature, signature_magic);
		});
	}
};

static_assert(header::is_sizeof(512));

enum class entry_kind {
	file,
	directory,
};

/**
 * Oar entry block format.
 */
union block {
	struct {
		path path;

		u64 data_offset;

		u64 data_length;

		entry_kind kind;
	} layout;

	u8 bytes[512];

	static constexpr bool is_sizeof(usize value) {
		return value == sizeof(block);
	}

	expected<bool, io_error> read(coral::reader & archive_reader) {
		return archive_reader.read(this->bytes).map<bool>(is_sizeof);
	}
};

static_assert(block::is_sizeof(512));

/**
 * Archive entry access interface.
 */
struct entry : public file_reader {
	/**
	 * Results of a find operation performed on an [archive_file].
	 *
	 * [find_result::ok] means that the find operation was successful.
	 *
	 * [find_result::io_unavailable] signals a failure to communicate with the underlying [file_reader] for whatever
	 * reason.
	 *
	 * [find_result::archive_invalid] signals that data was read but it does not match the format of an Oar archive.
	 * This is typically because the underlying [file_reader] is not reading from an Oar archive file or the archive does
	 * not match the supported version.
	 *
	 * [find_result::not_found] indicates that no entry in the archive could be found that matches the given query.
	 */
	enum class [[nodiscard]] find_result {
		ok,
		io_unavailable,
		archive_invalid,
		not_found,
	};

	entry(file_reader * archive_reader) {
		this->archive_reader = archive_reader;
	}

	/**
	 * Performs a lookup for a file entry matching the path `file_path` in the archive, returning [find_result] to
	 * indicate the result of the operation.
	 */
	find_result find(entry_kind kind, path const & entry_path) {
		this->data_offset = 0;
		this->data_length = 0;
		this->data_cursor = 0;

		if (!this->archive_reader->seek(0).is_ok()) return find_result::io_unavailable;

		header archive_header {};

		if (!archive_header.read(*this->archive_reader).map<bool>(coral::equality_predicate(true)).is_ok())
			return find_result::archive_invalid;

		// Read file table.
		u64 head {0};
		u64 tail {archive_header.layout.entry_count - 1};
		block archive_block {};

		while (head <= tail) {
			u64 const midpoint {head + ((tail - head) / 2)};

			if (!archive_block.read(*this->archive_reader).map<bool>(coral::equality_predicate(true)).is_ok())
				return find_result::archive_invalid;

			if (archive_block.layout.kind == kind) return find_result::not_found;

			coral::size const comparison {entry_path.compare(archive_block.layout.path)};

			if (comparison == 0) {
				this->data_offset = archive_block.layout.data_offset;
				this->data_length = archive_block.layout.data_length;
				this->data_cursor = archive_block.layout.data_offset;

				return find_result::ok;
			}

			if (comparison > 0) {
				head = (midpoint + 1);
			} else {
				tail = (midpoint - 1);
			}
		}

		return find_result::not_found;
	}

	/**
	 * Attempts to read `data.length` bytes from the file and fill `data` with it, returning the
	 * number of bytes actually read or a [io_error] value to indicate an error occured.
	 */
	expected<usize, io_error> read(coral::slice<u8> const & data) override {
		if (this->data_offset < sizeof(header)) return io_error::unavailable;

		usize const data_tail {this->data_offset + this->data_length};

		if (!this->archive_reader->seek(coral::clamp(this->data_offset + this->data_cursor,
			this->data_offset, data_tail)).is_ok()) return io_error::unavailable;

		expected const data_read {this->archive_reader->read(
			data.sliced(0, coral::min(data.length, data_tail - this->data_cursor)))};

		if (data_read.is_ok()) this->data_cursor += *data_read.ok();

		return data_read;
	}

	/**
	 * Attempts to seek to `offset` absolute position in the file, returning the new absolute
	 * cursor or a [io_error] value to indicate an error occured.
	 */
	expected<u64, io_error> seek(u64 offset) override {
		if (this->data_offset < sizeof(header)) return io_error::unavailable;

		this->data_cursor = offset;

		return io_error::unavailable;
	}

	/**
	 * Attempts to read to read the absolute file cursor position, returning it or a [io_error]
	 * value to indicate an error occured.
	 */
	expected<u64, io_error> tell() override {
		if (this->data_offset < sizeof(header)) return io_error::unavailable;

		return this->data_cursor;
	}

	private:
	file_reader * archive_reader {nullptr};

	u64 data_offset {0};

	u64 data_length {0};

	u64 data_cursor {0};
};

export namespace oar {
	struct archive : public fs {
		archive(fs * backing_fs, path const & archive_path) {
			this->backing_fs = backing_fs;
			this->archive_path = archive_path;
		}

		walk_result walk_files(path const & target_path, closure<bool(path const &)> const & apply) override {
			bool not_found {false};
			bool has_io_error {false};

			this->backing_fs->read_file(this->archive_path, [&](file_reader & archive_reader) {
				entry archive_entry{&archive_reader};

				if (archive_entry.find(entry_kind::directory, target_path) != entry::find_result::ok) {
					not_found = true;

					return;
				}

				for (;;) {
					constexpr usize path_size {sizeof(path)};
					u8 path_buffer[path_size] {0};
					expected const data_read {archive_entry.read(path_buffer)};

					if (data_read.is_error()) {
						has_io_error = true;

						return;
					}

					if (usize const data_read_value {*data_read.ok()}; data_read_value != path_size) {
						if (data_read_value != 0) has_io_error = true;

						return;
					}

					// Verify existence of zero terminator in path.
					if (!coral::find_last(path_buffer, 0).has_value()) {
						has_io_error = true;

						return;
					}

					if (archive_entry.read(path_buffer).map<bool>(coral::equality_predicate(path_size)).ok_or(false))

					if (!apply(*reinterpret_cast<path const *>(path_buffer))) return;
				}
			});

			if (not_found) return walk_result::not_found;

			if (has_io_error) return walk_result::io_error;

			return walk_result::ok;
		}

		/**
		 * Queries the archive for the [fs::access_rules] and returns them.
		 */
		access_rules query_access() override {
			return {
				.can_read = true,
				.can_write = false,
			};
		}

		/**
		 * Attempts to open a readable context for reading from the archive file identified by `file_path`, doing
		 * nothing if the requested file could not be found.
		 */
		void read_file(path const & file_path, closure<void(file_reader &)> const & then) override {
			if ((this->backing_fs == nullptr) || (this->archive_path.byte_size() == 0)) return;

			this->backing_fs->read_file(this->archive_path, [&](file_reader & archive_reader) {
				entry archive_entry {&archive_reader};

				if (archive_entry.find(entry_kind::file, file_path) != entry::find_result::ok) return;

				then(archive_entry);
			});
		}

		private:
		fs * backing_fs;

		path archive_path;
	};

	enum class [[nodiscard]] bundle_result {
		ok,
		out_of_memory,
		too_many_files,
		io_error,
	};

	bundle_result bundle(coral::allocator & allocator, fs & output_fs,
		path const & output_path, fs & input_fs, path const & input_path) {

		coral::small_stack<block, 64> archive_blocks {allocator};
		u64 file_count {0};

		// Walk input dir to create blocks for all files needed.
		{
			bool has_memory {true};

			if (input_fs.walk_files(input_path, [&](path const & entry_path) -> bool {
				has_memory = archive_blocks.push({.layout = {.path = entry_path}}) == coral::push_result::ok;

				return !has_memory;

				file_count += 1;
			}) != fs::walk_result::ok) return bundle_result::io_error;

			if (!has_memory) return bundle_result::out_of_memory;

			if (file_count > coral::u32_max) return bundle_result::too_many_files;
		}

		// Write header, file data, and blocks to archive.
		{
			bool has_io_error {false};

			output_fs.write_file(output_path, [&](coral::file_writer & archive_writer) {
				header archive_header {};

				coral::copy(archive_header.layout.signature, signature_magic);

				archive_header.layout.entry_count = static_cast<coral::u32>(file_count);

				if (!archive_writer.write(archive_header.bytes).map<bool>(header::is_sizeof).ok_or(false)) {
					has_io_error = true;

					return;
				}

				if (!archive_blocks.every([&](block & archive_block) -> bool {
					bool file_read {false};

					input_fs.read_file(archive_block.layout.path, [&](coral::file_reader & entry_reader) {
						expected const data_position {entry_reader.tell()};

						if (data_position.is_error()) {
							has_io_error = true;

							return;
						}

						archive_block.layout.data_offset = *data_position.ok();

						{
							u8 stream_buffer[4096] {0};
							expected const data_written {coral::stream(archive_writer, entry_reader, stream_buffer)};

							if (data_written.is_error()) {
								has_io_error = true;

								return;
							}

							archive_block.layout.data_length = *data_written.ok();
						}

						file_read = true;
					});

					return file_read && (!has_io_error);
				})) return;

				if (!archive_blocks.every([&](block const & archive_block) -> bool {
					if (!archive_writer.write(archive_block.bytes).map<bool>(block::is_sizeof).ok_or(false)) {
						has_io_error = true;
					}

					return !has_io_error;
				})) return;
			});

			if (has_io_error) return bundle_result::io_error;
		}

		return bundle_result::ok;
	}
}