use bufio;
use bytes;
use compress::zlib;
use io;

export type decoder = struct {
	vtable: io::stream,
	src: io::handle,
	inflate: zlib::reader,
	ihdr: *ihdr,
	filter: (filter | void),
	bpp: size, // bytes per pixel, rounded up
	buffered: size,
	cr: []u8,
	pr: []u8,
};

// Returns a new image decoder, from which raw pixel data may be read via
// [[io::read]]. The input should be a compressed IDAT data stream, normally
// obtained via an [[idat_reader]].
//
// The user must provide a pixel buffer of suitable size to store two scanlines
// for image filtering, or larger for interlaced images; see [[decoder_bufsiz]]
// to determine the appropriate length.
//
// The header is borrowed from the input and should remain valid until the
// decoder is no longer in use.
//
// The user must call [[io::close]] afterwards to free resources allocated for
// the decompressor. The call never returns an error.
export fn newdecoder(
	src: io::handle,
	ihdr: *ihdr,
	buf: []u8,
) (decoder | error) = {
	assert(len(buf) == decoder_bufsiz(ihdr),
		"Incorrect buffer length provided for PNG decoder");
	bytes::zero(buf);
	return decoder {
		vtable = &decoder_vtable,
		src = src,
		inflate = zlib::decompress(src)?,
		ihdr = ihdr,
		filter = void,
		bpp = bytesperpixel(ihdr),
		buffered = 0,
		cr = buf[..len(buf) / 2],
		pr = buf[len(buf) / 2..],
	};
};

const decoder_vtable: io::vtable = io::vtable {
	reader = &decoder_read,
	closer = &decoder_close,
	...
};

fn decoder_read(st: *io::stream, buf: []u8) (size | io::EOF | io::error) = {
	let dec = st: *decoder;
	assert(dec.vtable == &decoder_vtable);

	if (dec.buffered != 0) {
		return decoder_copy(dec, buf);
	};

	if (dec.filter is void) {
		const ft = match (bufio::read_byte(&dec.inflate)?) {
		case io::EOF =>
			if (!(io::read(dec.src, &[0u8])? is io::EOF)) {
				// Extra data following zlib stream
				return wraperror(invalid);
			};
			return io::EOF;
		case let b: u8 =>
			yield b: filter;
		};
		if (ft > filter::PAETH) {
			return wraperror(unsupported);
		};
		dec.filter = ft;
	};

	if (dec.ihdr.interlace == interlace::ADAM7) {
		abort(); // TODO
	};

	// Read one scanline
	match (io::readall(&dec.inflate, dec.cr)) {
	case io::EOF =>
		return wraperror(invalid);
	case let err: io::error =>
		if (err is io::underread) {
			return wraperror(invalid);
		};
		return err;
	case size => void;
	};

	applyfilter(dec);
	dec.buffered = len(dec.cr);
	return decoder_copy(dec, buf);
};

fn decoder_close(st: *io::stream) (void | io::error) = {
	let dec = st: *decoder;
	assert(dec.vtable == &decoder_vtable);
	io::close(&dec.inflate)?;
};

// Returns the size of a buffer suitable to store decoder state for this image.
// The computed size is twice the length of a scanline in bytes, unless the
// image is interlaced, in which case the buffer size is the full length
// necessary to store the complete image.
export fn decoder_bufsiz(ihdr: *ihdr) size = {
	let sampledepth = ihdr.bitdepth: uint;
	const samples: uint = switch (ihdr.colortype) {
	case colortype::GRAYSCALE =>
		yield 1;
	case colortype::RGB =>
		yield 3;
	case colortype::PLTE =>
		yield 1;
	case colortype::GRAYALPHA =>
		yield 2;
	case colortype::RGBA =>
		yield 4;
	};
	const width = ihdr.width: uint, height = ihdr.height: uint;

	let scanline = width * samples * sampledepth; // in bits
	if (scanline % 8 != 0) {
		// Pad to nearest byte
		scanline += 8 - (scanline % 8);
	};

	if (ihdr.interlace == interlace::ADAM7) {
		return scanline * height;
	};

	return (scanline / 8) * 2; // Two scanlines
};

fn bytesperpixel(ihdr: *ihdr) size = {
	let sampledepth = ihdr.bitdepth: uint;
	const samples: uint = switch (ihdr.colortype) {
	case colortype::GRAYSCALE =>
		yield 1;
	case colortype::RGB =>
		yield 3;
	case colortype::PLTE =>
		yield 1;
	case colortype::GRAYALPHA =>
		yield 2;
	case colortype::RGBA =>
		yield 4;
	};
	return (sampledepth * samples + 7) / 8;
};

// Copies pending pixel data into a user buffer.
fn decoder_copy(dec: *decoder, buf: []u8) size = {
	let max = dec.buffered;
	if (len(buf) < max) {
		max = len(buf);
	};
	assert(max > 0);
	buf[..max] = dec.cr[..max];
	dec.cr[..len(dec.cr)-max] = dec.cr[max..];
	dec.buffered -= max;
	return max;
};