use io;
use memio;

export type plte_reader = struct {
	src: *reader,
};

// Returns a new PLTE reader for a [[chunk_reader]].
export fn new_plte_reader(src: *reader) (plte_reader | invalid) = {
	assert(chunk_type(src) == PLTE,
		"Attempted to create PLTE reader for non-PLTE chunk");
	const length = chunk_length(src);
	if (length % 3 != 0) {
		return invalid;
	};
	const ncolor = length / 3;
	if (ncolor < 1 && ncolor > 256) {
		return invalid;
	};
	return plte_reader {
		src = src,
	};
};

// Returns the number of colors defined by this palette. The return value will
// be between 1 and 256 (inclusive), so should the caller wish to statically
// allocate a suitable buffer, [256]u32 will suffice for all cases.
export fn plte_ncolors(pr: *plte_reader) size = {
	return chunk_length(pr.src) / 3;
};

// Reads the palette from a PLTE chunk. The caller must provide a buffer equal
// in length to the return value of [[plte_ncolors]], which will be filled with
// RGB colors, such that the most significant byte is red and the least
// significant byte is unused (set to 0xFF). Returns the number of colors in the
// palette.
export fn plte_read(pr: *plte_reader, buf: []u32) (size | error) = {
	const ncolor = plte_ncolors(pr);
	assert(len(buf) == ncolor, "Invalid buffer size for plte_read");

	static let rbuf: [3 * 256]u8 = [0...];
	for (let i = 0z; i < ncolor) {
		let n = (ncolor - i) * 3;
		if (n > len(rbuf)) {
			n = len(rbuf);
		};
		match (io::readall(pr.src, rbuf[..n])?) {
		case io::EOF =>
			return invalid;
		case size =>
			yield;
		};
		for (let q = 0z; q < n / 3; q += 1) {
			const r = rbuf[q * 3 + 0]: u32;
			const g = rbuf[q * 3 + 1]: u32;
			const b = rbuf[q * 3 + 2]: u32;
			buf[i + q] = (r << 24) | (g << 16) | (b << 8) | 0xFF;
		};
		i += n / 3;
	};
	// required so that the checksum gets read
	assert(io::read(pr.src, rbuf[..1])? is io::EOF);
	return ncolor;
};

@test fn plte_reader() void = {
	const src = memio::fixed(palette_sample);
	const read = newreader(&src) as reader;
	assert(nextchunk(&read) as u32 == IHDR);
	io::copy(io::empty, &read)!;

	for (true) {
		const ctype = nextchunk(&read) as u32;
		if (ctype == PLTE) {
			break;
		};
		io::copy(io::empty, &read)!;
	};

	const plte = new_plte_reader(&read)!;

	let buf: []u32 = alloc([0...], plte_ncolors(&plte));
	defer free(buf);
	plte_read(&plte, buf)!;

	const expected: [_]u32 = [
		0x2200FFFF, 0x00FFFFFF, 0x8800FFFF, 0x22FF00FF, 0x0099FFFF,
		0xFF6600FF, 0xDD00FFFF, 0x77FF00FF, 0xFF0000FF, 0x00FF99FF,
		0xDDFF00FF, 0xFF00BBFF, 0xFFBB00FF, 0x0044FFFF, 0x00FF44FF,
	];
	assert(len(expected) == len(buf));
	for (let i = 0z; i < len(expected); i += 1) {
		assert(buf[i] == expected[i]);
	};
};