use errors;
use io;
use types;
use rt;
use fmt;

use mbr;
use sector;

export type nogpt = !void;

export type gpt = struct {
	mbr_sec: sector::sector,
	mbr: *mbr::bootsector,
	primary: block,
	backup: block
};

export type block = struct {
	header_sec: sector::sector,
	entries_sec: sector::sector,
	header: *header,
	entries: *[128]entry
};

fn block_commit(self: *block) (void | io::error) = {
	sector::commit(&self.header_sec)?;
	sector::commit(&self.entries_sec)?;
};

fn block_fetch(self: *block) (void | io::error) = {
	sector::fetch(&self.header_sec)?;
	sector::fetch(&self.entries_sec)?;
};

fn block_drop(self: *block) void = {
	sector::finish(&self.header_sec);
	sector::finish(&self.entries_sec);
};

export type header = struct {
	@offset(0) signature: u64,
	@offset(8) revision: u32,
	@offset(12) header_size: u32,
	@offset(16) header_crc32: u32,
	@offset(20) reserved0: u32,
	@offset(24) header_lba: u64,
	@offset(32) backup_header_lba: u64,
	@offset(40) first_lba: u64,
	@offset(48) last_lba: u64,
	@offset(56) disk_guid: [2]u64,
	@offset(72) entries_lba: u64,
	@offset(80) entries_len: u32,
	@offset(84) entry_size: u32,
	@offset(88) entries_crc32: u32
};

fn header_validate(self: *header) bool = {
	return self.signature == 0x5452415020494645;
};

export type entry = struct {
	@offset(0) part_type: [2]u64,
	@offset(16) part: [2]u64,
	@offset(32) lba_begin: u64,
	@offset(40) lba_end: u64,
	@offset(48) attributes: u64,
	@offset(56) name: [72]u8
};

export fn from(fd: io::file) (*gpt | nogpt | io::error) = {
	// Get the primary header
	const prim_header_sec = sector::map(fd, 1, 1);
	sector::fetch(&prim_header_sec)?;

	const prim_header = prim_header_sec.buf: *[*]u8: *header;
	if (!header_validate(prim_header)) {
		sector::finish(&prim_header_sec);
		return nogpt;
	};

	// Get the primary entries
	const prim_entries_sec = sector::map(fd, prim_header.entries_lba, 128*128/512);
	sector::fetch(&prim_entries_sec)?;

	// Get the backup header
	const back_header_sec = sector::map(fd, prim_header.backup_header_lba, 1);
	sector::fetch(&back_header_sec)?;

	const back_header = back_header_sec.buf: *[*]u8: *header;
	if (!header_validate(back_header)) {
		sector::finish(&prim_header_sec);
		sector::finish(&back_header_sec);
		return nogpt;
	};

	// Get the backup entries
	const back_entries_sec = sector::map(fd, back_header.entries_lba, 128*128/512);
	sector::fetch(&back_entries_sec)?;

	const mbr_sec = sector::map(fd, 0, 1);
	sector::fetch(&mbr_sec)?;

	// Allocate the GPT structure
	let self = alloc(gpt {
		mbr_sec = mbr_sec,
		mbr = mbr_sec.buf: *[*]u8: *mbr::bootsector,
		primary = block {
			header_sec = prim_header_sec,
			entries_sec = prim_entries_sec,
			header = prim_header,
			entries = prim_entries_sec.buf: *[*]u8: *[128]entry
		},
		backup = block {
			header_sec = back_header_sec,
			entries_sec = back_entries_sec,
			header = back_header,
			entries = back_entries_sec.buf: *[*]u8: *[128]entry
		}
	});

	return self;
};

export fn create(fd: io::file, freeno: size) *gpt = {
	const entries_begin = 2z;
	const free_begin = entries_begin + 128*128/512;
	const bentries_begin = free_begin + freeno + 1;
	const bheader_begin = bentries_begin + 128*128/512;

	const mbr_sec = sector::map(fd, 0, 1);
	const mbr = mbr_sec.buf: *[*]u8: *mbr::bootsector;

	*mbr = mbr::bootsector {
		magic = 0xaa55,
		...
	};

	const prim_header_sec = sector::map(fd, 1, 1);
	const prim_header = prim_header_sec.buf: *[*]u8: *header;
	*prim_header = header {
		signature = 0x5452415020494645,
		revision = 0x10000,
		header_size = 92,
		header_lba = 1,
		backup_header_lba = bheader_begin,
		first_lba = free_begin,
		last_lba = bentries_begin - 1,
		entries_lba = entries_begin,
		entries_len = 128,
		entry_size = 128,
		...
	};

	const prim_entries_sec = sector::map(fd, entries_begin, 128*128/512);
	const prim_entries = prim_entries_sec.buf: *[*]u8: *[128]entry;

	const back_header_sec = sector::map(fd, bheader_begin, 1);
	const back_header = back_header_sec.buf: *[*]u8: *header;

	const back_entries_sec = sector::map(fd, bentries_begin, 128*128/512);
	const back_entries = back_entries_sec.buf: *[*]u8: *[128]entry;
	
	let self = alloc(gpt {
		mbr_sec = mbr_sec,
		mbr = mbr,
		primary = block {
			header_sec = prim_header_sec,
			entries_sec = prim_entries_sec,
			header = prim_header,
			entries = prim_entries
		},
		backup = block {
			header_sec = back_header_sec,
			entries_sec = back_entries_sec,
			header = back_header,
			entries = back_entries
		}
	});

	mkbackup(self);

	return self;
};

// Finds an lba for a new partition that fits the desired length
fn findfree(self: *gpt, length: u64) u64 = {
	const header = self.primary.header;
	const entries = self.primary.entries;

	let cur = header.first_lba;

	for (cur >= header.first_lba && cur + length <= header.last_lba) {
		let i = 0z;
		for (i < header.entries_len) {
			defer i += 1;
			if (entries[i].lba_begin == 0 && entries[i].lba_end == 0)
				continue;
			if (overlap((cur, length), (entries[i].lba_begin, entries[i].lba_end - entries[i].lba_begin + 1))) {
				cur = entries[i].lba_end + 1;
				break;
			};
		};

		if (i == header.entries_len)
			return cur;
	};

	return 0;
};

// Finds a GPT entry that is currently not used
// The entry pointer returned is borrowed from the arguments
fn findfreeentry(self: *gpt) nullable *entry = {
	const header = self.primary.header;
	const entries = self.primary.entries;

	for (let i = 0z; i < header.entries_len; i += 1) {
		if (entries[i].lba_begin == 0 && entries[i].lba_end == 0)
			return &entries[i];
	};

	return null;
};

// Allocates a partition with a given length (in lba)
export fn allocate(self: *gpt, length: u64) nullable *entry = {
	const addr = findfree(self, length);
	if (addr == 0)
		return null;

	const entry = match (findfreeentry(self)) {
	case let e: *entry =>
		yield e;
	case => return null;
	};

	entry.lba_begin = addr;
	entry.lba_end = addr + length - 1;

	return entry;
};

// Returns a stream to the partition associated with the entry
// in order to write in it
export fn partstream_writer(self: *gpt, entry: *entry) (io::limitstream | io::error) = {
	const fd = self.primary.header_sec.fd;
	const part = entry.lba_begin * sector::sector_length;
	const length = (entry.lba_end - entry.lba_begin + 1) * sector::sector_length;
	

	io::seek(fd, part: io::off, io::whence::SET)?;
	return io::limitwriter(fd, length);
};

// Returns a stream to the partition associated with the entry
// in order to read from it
export fn partstream_reader(self: *gpt, entry: *entry) (io::limitstream | io::error) = {
	const fd = self.primary.header_sec.fd;
	const part = entry.lba_begin * sector::sector_length;
	const length = (entry.lba_end - entry.lba_begin + 1) * sector::sector_length;

	io::seek(fd, part: io::off, io::whence::SET)?;
	return io::limitreader(fd, length);
};

// Frees the resources associated with the GPT structure.
export fn finish(self: *gpt) void = {
	sector::finish(&self.mbr_sec);
	block_drop(&self.primary);
	block_drop(&self.backup);
	free(self);
};

// Updates the data structures from the disk
export fn fetch(self: *gpt) (void | io::error) = {
	sector::fetch(&self.mbr_sec)?;
	block_fetch(&self.primary)?;
	block_fetch(&self.backup)?;
};

// Copies the current primary block into its backup
// Note: does not compute the checksums (see chksums())
export fn mkbackup(self: *gpt) void = {
	rt::memcpy(self.backup.header, self.primary.header, size(header));
	rt::memcpy(self.backup.entries, self.primary.entries, size([128]entry));
	
	self.backup.header.header_lba = self.backup.header_sec.offs;
	self.backup.header.backup_header_lba = self.primary.header_sec.offs;
	self.backup.header.entries_lba = self.backup.entries_sec.offs;

	self.backup.header.entries_crc32 = entries_crc32(self.backup.entries);
	self.backup.header.header_crc32 = header_crc32(self.backup.header);
};

// Commits changed done on the data structure to the disk
export fn commit(self: *gpt) (void | io::error) = {
	sector::commit(&self.mbr_sec)?;
	block_commit(&self.primary)?;
	block_commit(&self.backup)?;
};

// Updates the CRC32 checksums of the different blocks
export fn chksums(self: *gpt) void = {
	self.primary.header.entries_crc32 = entries_crc32(self.primary.entries);
	self.primary.header.header_crc32 = header_crc32(self.primary.header);
};