use bytes;
use endian;
use hash::crc32;
use hash;
use io;
use memio;

export type reader = struct {
	vtable: io::stream,
	src: io::handle,
	status: reader_status,
	length: size,
	ctype: u32,
	crc: crc32::state,
};

export type reader_status = enum u8 {
	// cursor is at the start of a chunk header (or at EOF)
	NEXT,
	// cursor is at the end of a chunk header
	HEADER_READ,
	// cursor is within chunk data
	READING,
};

// Creates a new PNG decoder. Reads and verifies the PNG magic before returning
// the reader object. Call [[nextchunk]] to read the next chunk from the input.
export fn newreader(src: io::handle) (reader | error) = {
	let buf: [MAGIC_LEN]u8 = [0...];
	match (io::readall(src, buf[..])?) {
	case io::EOF =>
		return invalid;
	case size =>
		yield;
	};
	if (!bytes::equal(buf, magic)) {
		return invalid;
	};
	return reader {
		vtable = &reader_vtable,
		src = src,
		status = reader_status::NEXT,
		length = 0,
		ctype = 0,
		crc = crc32::crc32(&crc32::ieee_table),
	};
};

@test fn newreader() void = {
	const src = memio::fixed([]);
	assert(newreader(&src) as error is invalid);

	const src = memio::fixed(magic);
	assert(newreader(&src) is reader);
};

// Starts decoding a new chunk from the reader, returning its type. The contents
// of the chunk can be read by calling [[io::read]] on the reader until it
// returns [[io::EOF]].
//
// However, in an ideal scenario, the caller will not read directly from the
// chunk, but instead will select a chunk-type-aware decoder based on the
// returned chunk type, and use that to read the chunk's contents.
//
// Calling this function repeatedly will return the current chunk type with no
// side effects until [[io::read]] is called on the reader, after which the
// user must read the contents of the chunk until [[io::EOF]] is returned before
// calling nextchunk again. If the checksum fails verification, [[invalid]] is
// returned from [[io::read]].
export fn nextchunk(src: *reader) (u32 | io::EOF | error) = {
	assert(src.status != reader_status::READING,
		"Must finish previous chunk before calling nextchunk again");

	if (src.status == reader_status::NEXT) {
		let buf: [8]u8 = [0...];
		match (io::readall(src.src, buf[..])?) {
		case io::EOF =>
			return io::EOF;
		case size =>
			yield;
		};
		src.status = reader_status::HEADER_READ;
		src.length = endian::begetu32(buf[..4]);
		src.ctype = endian::begetu32(buf[4..]);
		src.crc = crc32::crc32(&crc32::ieee_table);
		hash::write(&src.crc, buf[4..]);
	};

	return src.ctype;
};

const reader_vtable: io::vtable = io::vtable {
	reader = &read,
	...
};

// Returns the type of the chunk being read.
export fn chunk_type(src: *reader) u32 = {
	assert(src.status != reader_status::NEXT,
		"Must call nextchunk before calling chunk_type");
	return src.ctype;
};

// Returns the remaining length of the chunk being read in bytes, not including
// the header or checksum (that is, it returns the length of the chunk data).
export fn chunk_length(src: *reader) size = {
	assert(src.status != reader_status::NEXT,
		"Must call nextchunk before calling chunk_length");
	return src.length;
};

fn read(st: *io::stream, buf: []u8) (size | io::EOF | io::error) = {
	let st = st: *reader;
	assert(st.vtable == &reader_vtable);

	if (st.status == reader_status::NEXT) {
		return io::EOF;
	};
	st.status = reader_status::READING;

	if (st.length == 0) {
		let ckbuf: [4]u8 = [0...];
		match (io::readall(st.src, ckbuf[..])?) {
		case io::EOF =>
			return wraperror(invalid);
		case size =>
			yield;
		};
		st.status = reader_status::NEXT;
		const want = endian::begetu32(ckbuf);
		const have = crc32::sum32(&st.crc);
		if (want != have) {
			return wraperror(invalid);
		};
		return io::EOF;
	};

	const max = if (len(buf) < st.length) {
		yield len(buf);
	} else {
		yield st.length;
	};

	const z = match (io::read(st.src, buf[..max])?) {
	case io::EOF =>
		// Missing checksum
		return wraperror(invalid);
	case let z: size =>
		yield z;
	};

	hash::write(&st.crc, buf[..z]);
	st.length -= z;
	return z;
};

@test fn nextchunk() void = {
	const src = memio::fixed(magic);
	const read = newreader(&src) as reader;
	assert(nextchunk(&read) is io::EOF);

	const src = memio::fixed(simple_png);
	const read = newreader(&src) as reader;
	assert(nextchunk(&read) as u32 == IHDR);
	let buf: [32]u8 = [0...];
	assert(io::read(&read, buf) as size == 13);
	assert(io::read(&read, buf) is io::EOF);
	assert(bytes::equal(buf[..13], simple_png[16..16+13]));

	const src = memio::fixed(invalid_chunk);
	const read = newreader(&src) as reader;
	nextchunk(&read) as u32;
	assert(io::read(&read, buf) as size == 13);
	assert(io::read(&read, buf) is io::error);
};