// Copyright 2020 The Wuffs Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use "std/crc32"
use "std/zlib"

pub status "#bad checksum"
pub status "#bad chunk"
pub status "#bad filter"
pub status "#bad header"
pub status "#missing palette"
pub status "#unsupported PNG file"

pri status "#internal error: inconsistent workbuf length"
pri status "#internal error: zlib decoder did not exhaust its input"

pub const DECODER_WORKBUF_LEN_MAX_INCL_WORST_CASE : base.u64 = 0

pub struct decoder? implements base.image_decoder(
	// The 0x00FF_FFFF limit is arbitrary (the PNG spec says 0x7FFF_FFFF) but
	// it means that (width * height * bytes_per_pixel) doesn't overflow.
	width  : base.u32[..= 0x00FF_FFFF],
	height : base.u32[..= 0x00FF_FFFF],

	// bytes_per_row doesn't include the 1 byte for the per-row filter.
	bytes_per_row : base.u64[..= 0x07FF_FFF8],

	workbuf_wi     : base.u64,
	workbuf_length : base.u64[..= 0x0007_FFFF_F100_0007],

	// Call sequence states:
	//  - 0x00: initial state.
	//  - 0x03: image config decoded.
	//  - 0x04: frame config decoded.
	//  - 0xFF: end-of-data, usually after (the non-animated) frame decoded.
	//
	// State transitions:
	//
	//  - 0x00 -> 0x03: via DIC
	//  - 0x00 -> 0x04: via DFC with implicit DIC
	//  - 0x00 -> 0xFF: via DF  with implicit DIC and DFC
	//
	//  - 0x03 -> 0x04: via DFC
	//  - 0x03 -> 0xFF: via DF  with implicit DFC
	//
	//  - 0x04 -> 0xFF: via DFC
	//  - 0x04 -> 0xFF: via DF
	//
	//  - ???? -> 0x03: via RF  for ???? > 0x00
	//
	// Where:
	//  - DF  is decode_frame
	//  - DFC is decode_frame_config, implicit means nullptr args.dst
	//  - DIC is decode_image_config, implicit means nullptr args.dst
	//  - RF  is restart_frame
	call_sequence : base.u8,

	ignore_checksum : base.bool,

	depth           : base.u8[..= 16],
	color_type      : base.u8,
	filter_distance : base.u8[..= 8],

	seen_palette : base.bool,

	dst_pixfmt : base.u32,
	src_pixfmt : base.u32,

	chunk_type       : base.u32,
	chunk_type_array : array[4] base.u8,
	chunk_length     : base.u64,

	frame_config_io_position : base.u64,

	swizzler : base.pixel_swizzler,
	util     : base.utility,
)(
	crc32 : crc32.ieee_hasher,
	zlib  : zlib.decoder,

	dst_palette : array[4 * 256] base.u8,
	src_palette : array[4 * 256] base.u8,
)

pub func decoder.set_quirk_enabled!(quirk: base.u32, enabled: base.bool) {
	if args.quirk == base.QUIRK_IGNORE_CHECKSUM {
		this.ignore_checksum = args.enabled
		this.zlib.set_quirk_enabled!(quirk: args.quirk, enabled: args.enabled)
	}
}

pub func decoder.decode_image_config?(dst: nptr base.image_config, src: base.io_reader) {
	var magic         : base.u64
	var mark          : base.u64
	var checksum_have : base.u32
	var checksum_want : base.u32
	var status        : base.status

	if this.call_sequence <> 0 {
		return base."#bad call sequence"
	}

	magic = args.src.read_u64le?()
	if magic <> '\x89PNG\x0D\x0A\x1A\x0A'le {
		return "#bad header"
	}
	magic = args.src.read_u64le?()
	if magic <> '\x00\x00\x00\x0DIHDR'le {
		return "#bad header"
	}
	this.crc32.reset!()
	this.chunk_type_array[0] = 'I'
	this.chunk_type_array[1] = 'H'
	this.chunk_type_array[2] = 'D'
	this.chunk_type_array[3] = 'R'
	this.crc32.update_u32!(x: this.chunk_type_array[..])

	while true {
		mark = args.src.mark()
		status =? this.decode_ihdr?(src: args.src)
		if not this.ignore_checksum {
			checksum_have = this.crc32.update_u32!(x: args.src.since(mark: mark))
		}
		if status.is_ok() {
			break
		}
		yield? status
	} endwhile

	// Verify CRC-32 checksum.
	checksum_want = args.src.read_u32be?()
	if (not this.ignore_checksum) and (checksum_have <> checksum_want) {
		return "#bad checksum"
	}

	// Read up until an IDAT chunk.
	//
	// By default, libpng "warns and discards" when seeing ancillary chunk
	// checksum failures (as opposed to critical chunk checksum failures) but
	// it still continues to decode the image. Wuffs' decoder is similar,
	// simply always ignoring ancillary chunks' CRC-32 checksums.
	//
	// https://github.com/glennrp/libpng/blob/dbe3e0c43e549a1602286144d94b0666549b18e6/png.h#L1436
	//
	// We've already seen the IHDR chunk. We're not expecting an IEND chunk. An
	// IDAT chunk breaks the loop. The only other possible critical chunk is a
	// PLTE chunk. We verify PLTE checksums here but ignore other checksums.
	while true {
		this.chunk_length = args.src.read_u32be_as_u64?()
		this.chunk_type = args.src.read_u32le?()
		if (not this.ignore_checksum) and
			((this.chunk_type == 'IDAT'le) or (this.chunk_type == 'PLTE'le)) {
			this.crc32.reset!()
			this.chunk_type_array[0] = ((this.chunk_type >> 0) & 0xFF) as base.u8
			this.chunk_type_array[1] = ((this.chunk_type >> 8) & 0xFF) as base.u8
			this.chunk_type_array[2] = ((this.chunk_type >> 16) & 0xFF) as base.u8
			this.chunk_type_array[3] = ((this.chunk_type >> 24) & 0xFF) as base.u8
			this.crc32.update_u32!(x: this.chunk_type_array[..])
		}

		if this.chunk_type == 'IDAT'le {
			break
		}

		while true {
			mark = args.src.mark()
			status =? this.decode_other_chunk?(src: args.src)
			if (not this.ignore_checksum) and (this.chunk_type == 'PLTE'le) {
				checksum_have = this.crc32.update_u32!(x: args.src.since(mark: mark))
			}
			if status.is_ok() {
				break
			}
			yield? status
		} endwhile
		checksum_want = args.src.read_u32be?()
		if (not this.ignore_checksum) and (this.chunk_type == 'PLTE'le) and
			(checksum_have <> checksum_want) {
			return "#bad checksum"
		}

	} endwhile

	if (this.color_type == 3) and (not this.seen_palette) {
		return "#missing palette"
	}

	this.frame_config_io_position = args.src.position()

	if args.dst <> nullptr {
		args.dst.set!(
			pixfmt: this.dst_pixfmt,
			pixsub: 0,
			width: this.width,
			height: this.height,
			first_frame_io_position: this.frame_config_io_position,
			first_frame_is_opaque: false)
	}

	this.call_sequence = 3
}

pri func decoder.decode_ihdr?(src: base.io_reader) {
	var a32 : base.u32
	var a8  : base.u8

	a32 = args.src.read_u32be?()
	if a32 >= 0x8000_0000 {
		return "#bad header"
	} else if a32 >= 0x0100_0000 {
		return "#unsupported PNG file"
	}
	this.width = a32

	a32 = args.src.read_u32be?()
	if a32 >= 0x8000_0000 {
		return "#bad header"
	} else if a32 >= 0x0100_0000 {
		return "#unsupported PNG file"
	}
	this.height = a32

	// Depth.
	a8 = args.src.read_u8?()
	if a8 > 16 {
		return "#bad header"
	}
	this.depth = a8

	// Color.
	this.color_type = args.src.read_u8?()
	if this.color_type == 0 {
		this.dst_pixfmt = base.PIXEL_FORMAT__Y
		this.src_pixfmt = base.PIXEL_FORMAT__Y
		if this.depth == 8 {
			this.filter_distance = 1
			this.bytes_per_row = (this.width as base.u64) * 1
		} else {
			return "#unsupported PNG file"
		}
	} else if this.color_type == 2 {
		this.dst_pixfmt = base.PIXEL_FORMAT__BGR
		this.src_pixfmt = base.PIXEL_FORMAT__RGB
		if this.depth == 8 {
			this.filter_distance = 3
			this.bytes_per_row = (this.width as base.u64) * 3
		} else {
			return "#unsupported PNG file"
		}
	} else if this.color_type == 3 {
		// TODO: s/BINARY/NONPREMUL/ and decode the tRNS chunk.
		this.dst_pixfmt = base.PIXEL_FORMAT__INDEXED__BGRA_BINARY
		this.src_pixfmt = base.PIXEL_FORMAT__INDEXED__BGRA_BINARY
		if this.depth == 8 {
			this.filter_distance = 1
			this.bytes_per_row = (this.width as base.u64) * 1
		} else {
			return "#unsupported PNG file"
		}
	} else if this.color_type == 6 {
		this.dst_pixfmt = base.PIXEL_FORMAT__BGRA_NONPREMUL
		this.src_pixfmt = base.PIXEL_FORMAT__RGBA_NONPREMUL
		if this.depth == 8 {
			this.filter_distance = 4
			this.bytes_per_row = (this.width as base.u64) * 4
		} else {
			return "#unsupported PNG file"
		}
	} else {
		return "#unsupported PNG file"
	}
	this.workbuf_length = (this.height as base.u64) * (1 + this.bytes_per_row)
	this.choose_filter_implementations!()

	// Compression.
	a8 = args.src.read_u8?()
	if a8 <> 0 {
		return "#bad header"
	}

	// Filter.
	a8 = args.src.read_u8?()
	if a8 <> 0 {
		return "#bad header"
	}

	// Interlace.
	a8 = args.src.read_u8?()
	if a8 == 0 {
		// No-op.
	} else if a8 == 1 {
		return "#unsupported PNG file"
	} else {
		return "#bad header"
	}
}

pri func decoder.choose_filter_implementations!() {
	// Filter 0 is a no-op. Filter 2, the up filter, should already vectorize
	// easily by a good optimizing C compiler.
	if this.filter_distance == 3 {
		choose filter_1 = [filter_1_distance_3_fallback]
		choose filter_3 = [filter_3_distance_3_fallback]
		choose filter_4 = [
			filter_4_distance_3_sse42,
			filter_4_distance_3_fallback]
	} else if this.filter_distance == 4 {
		choose filter_1 = [
			filter_1_distance_4_sse42,
			filter_1_distance_4_fallback]
		choose filter_3 = [
			filter_3_distance_4_sse42,
			filter_3_distance_4_fallback]
		choose filter_4 = [
			filter_4_distance_4_sse42,
			filter_4_distance_4_fallback]
	}
}

pri func decoder.decode_other_chunk?(src: base.io_reader) {
	if this.chunk_type == 'PLTE'le {
		if this.seen_palette {
			return "#bad chunk"
		}
		this.decode_plte?(src: args.src)
		this.seen_palette = true
	} else if this.chunk_type == 'tRNS'le {
		// TODO.
		args.src.skip?(n: this.chunk_length)
	} else {
		args.src.skip?(n: this.chunk_length)
	}
}

pri func decoder.decode_plte?(src: base.io_reader) {
	var num_palette_entries : base.u32[..= 256]
	var i                   : base.u32
	var argb                : base.u32

	if (this.chunk_length > 768) or ((this.chunk_length % 3) <> 0) {
		return "#bad header"
	}
	num_palette_entries = (this.chunk_length as base.u32) / 3

	while i < num_palette_entries {
		assert i < 256 via "a < b: a < c; c <= b"(c: num_palette_entries)
		// Convert from RGB (in memory order) to ARGB (in native u32 order)
		// to BGRA (in memory order).
		argb = args.src.read_u24be_as_u32?()
		argb |= 0xFF00_0000
		this.src_palette[(4 * i) + 0] = ((argb >> 0) & 0xFF) as base.u8
		this.src_palette[(4 * i) + 1] = ((argb >> 8) & 0xFF) as base.u8
		this.src_palette[(4 * i) + 2] = ((argb >> 16) & 0xFF) as base.u8
		this.src_palette[(4 * i) + 3] = ((argb >> 24) & 0xFF) as base.u8
		i += 1
	} endwhile

	// Set the remaining palette entries to opaque black.
	while i < 256 {
		this.src_palette[(4 * i) + 0] = 0x00
		this.src_palette[(4 * i) + 1] = 0x00
		this.src_palette[(4 * i) + 2] = 0x00
		this.src_palette[(4 * i) + 3] = 0xFF
		i += 1
	} endwhile
}

pub func decoder.decode_frame_config?(dst: nptr base.frame_config, src: base.io_reader) {
	if this.call_sequence < 3 {
		this.decode_image_config?(dst: nullptr, src: args.src)
	} else if this.call_sequence == 3 {
		if this.frame_config_io_position <> args.src.position() {
			return base."#bad restart"
		}
	} else if this.call_sequence == 4 {
		this.call_sequence = 0xFF
		return base."@end of data"
	} else {
		return base."@end of data"
	}

	if args.dst <> nullptr {
		args.dst.set!(bounds: this.util.make_rect_ie_u32(
			min_incl_x: 0,
			min_incl_y: 0,
			max_excl_x: this.width,
			max_excl_y: this.height),
			duration: 0,
			index: 0,
			io_position: this.frame_config_io_position,
			disposal: 0,
			opaque_within_bounds: false,
			overwrite_instead_of_blend: false,
			background_color: 0x0000_0000)
	}

	this.call_sequence = 4
}

pub func decoder.decode_frame?(dst: ptr base.pixel_buffer, src: base.io_reader, blend: base.pixel_blend, workbuf: slice base.u8, opts: nptr base.decode_frame_options) {
	var w               : base.io_writer
	var w_mark          : base.u64
	var r_mark          : base.u64
	var swizzler_status : base.status
	var zlib_status     : base.status
	var checksum_have   : base.u32
	var checksum_want   : base.u32

	if this.call_sequence < 4 {
		this.decode_frame_config?(dst: nullptr, src: args.src)
	} else if this.call_sequence == 4 {
		// No-op.
	} else {
		return base."@end of data"
	}

	this.workbuf_wi = 0
	while true {
		if (this.workbuf_wi > this.workbuf_length) or (
			this.workbuf_length > args.workbuf.length()) {
			return base."#bad workbuf length"
		}
		io_bind (io: w, data: args.workbuf[this.workbuf_wi .. this.workbuf_length]) {
			io_limit (io: args.src, limit: this.chunk_length) {
				w_mark = w.mark()
				r_mark = args.src.mark()
				zlib_status =? this.zlib.transform_io?(
					dst: w, src: args.src, workbuf: this.util.empty_slice_u8())
				if not this.ignore_checksum {
					this.crc32.update_u32!(x: args.src.since(mark: r_mark))
				}
				this.chunk_length ~sat-= args.src.count_since(mark: r_mark)
				this.workbuf_wi ~sat+= w.count_since(mark: w_mark)
			}
		}

		if zlib_status.is_ok() {
			break
		} else if zlib_status == base."$short write" {
			return base."#too much data"
		} else if zlib_status <> base."$short read" {
			return zlib_status
		} else if this.chunk_length == 0 {
			// Verify the non-final IDAT chunk's CRC-32 checksum.
			checksum_want = args.src.read_u32be?()
			if not this.ignore_checksum {
				checksum_have = this.crc32.update_u32!(x: this.util.empty_slice_u8())
				if checksum_have <> checksum_want {
					return "#bad checksum"
				}
			}

			// The next chunk should be another IDAT.
			this.chunk_length = args.src.read_u32be_as_u64?()
			this.chunk_type = args.src.read_u32le?()
			if this.chunk_type <> 'IDAT'le {
				return "#bad chunk"
			}

			// The 'IDAT'be is part of the next CRC-32 checksum's input.
			if not this.ignore_checksum {
				this.crc32.reset!()
				this.chunk_type_array[0] = 'I'
				this.chunk_type_array[1] = 'D'
				this.chunk_type_array[2] = 'A'
				this.chunk_type_array[3] = 'T'
				this.crc32.update_u32!(x: this.chunk_type_array[..])
			}
			continue
		} else if args.src.length() > 0 {
			return "#internal error: zlib decoder did not exhaust its input"
		}
		yield? base."$short read"
	} endwhile

	// Verify the final IDAT chunk's CRC-32 checksum.
	if not this.ignore_checksum {
		if this.chunk_length > 0 {
			// TODO: should this be a fatal error?
			return base."#too much data"
		}
		checksum_have = this.crc32.update_u32!(x: this.util.empty_slice_u8())
		checksum_want = args.src.read_u32be?()
		if checksum_have <> checksum_want {
			return "#bad checksum"
		}
	}

	if this.workbuf_wi <> this.workbuf_length {
		return base."#not enough data"
	} else if 0 < args.workbuf.length() {
		// For the top row, the Paeth filter (4) is equivalent to the Sub
		// filter (1), but the Paeth implementation is simpler if it can assume
		// that there is a previous row.
		if args.workbuf[0] == 4 {
			args.workbuf[0] = 1
		}
	}

	swizzler_status = this.swizzler.prepare!(
		dst_pixfmt: args.dst.pixel_format(),
		dst_palette: args.dst.palette_or_else(fallback: this.dst_palette[..]),
		src_pixfmt: this.util.make_pixel_format(repr: this.src_pixfmt),
		src_palette: this.src_palette[..],
		blend: args.blend)
	if not swizzler_status.is_ok() {
		return swizzler_status
	}
	swizzler_status = this.filter_and_swizzle!(dst: args.dst, workbuf: args.workbuf)
	if not swizzler_status.is_ok() {
		return swizzler_status
	}

	this.call_sequence = 0xFF
}

pri func decoder.filter_and_swizzle!(dst: ptr base.pixel_buffer, workbuf: slice base.u8) base.status {
	var dst_pixfmt          : base.pixel_format
	var dst_bits_per_pixel  : base.u32[..= 256]
	var dst_bytes_per_pixel : base.u64[..= 32]
	var dst_bytes_per_row   : base.u64
	var dst_palette         : slice base.u8
	var tab                 : table base.u8

	var y        : base.u32
	var dst      : slice base.u8
	var filter   : base.u8
	var curr_row : slice base.u8
	var prev_row : slice base.u8

	// TODO: the dst_pixfmt variable shouldn't be necessary. We should be able
	// to chain the two calls: "args.dst.pixel_format().bits_per_pixel()".
	dst_pixfmt = args.dst.pixel_format()
	dst_bits_per_pixel = dst_pixfmt.bits_per_pixel()
	if (dst_bits_per_pixel & 7) <> 0 {
		return base."#unsupported option"
	}
	dst_bytes_per_pixel = (dst_bits_per_pixel / 8) as base.u64
	dst_bytes_per_row = (this.width as base.u64) * dst_bytes_per_pixel
	dst_palette = args.dst.palette_or_else(fallback: this.dst_palette[..])
	tab = args.dst.plane(p: 0)

	while y < this.height {
		assert y < 0xFFFF_FFFF via "a < b: a < c; c <= b"(c: this.height)
		dst = tab.row(y: y)
		if dst_bytes_per_row < dst.length() {
			dst = dst[.. dst_bytes_per_row]
		}

		if 1 > args.workbuf.length() {
			return "#internal error: inconsistent workbuf length"
		}
		filter = args.workbuf[0]
		args.workbuf = args.workbuf[1 ..]
		if this.bytes_per_row > args.workbuf.length() {
			return "#internal error: inconsistent workbuf length"
		}
		curr_row = args.workbuf[.. this.bytes_per_row]
		args.workbuf = args.workbuf[this.bytes_per_row ..]

		if filter == 0 {
			// No-op.
		} else if filter == 1 {
			this.filter_1!(curr: curr_row)
		} else if filter == 2 {
			this.filter_2!(curr: curr_row, prev: prev_row)
		} else if filter == 3 {
			this.filter_3!(curr: curr_row, prev: prev_row)
		} else if filter == 4 {
			this.filter_4!(curr: curr_row, prev: prev_row)
		} else {
			return "#bad filter"
		}

		this.swizzler.swizzle_interleaved_from_slice!(
			dst: dst,
			dst_palette: dst_palette,
			src: curr_row)

		prev_row = curr_row
		y += 1
	} endwhile

	return ok
}

pub func decoder.frame_dirty_rect() base.rect_ie_u32 {
	return this.util.make_rect_ie_u32(
		min_incl_x: 0,
		min_incl_y: 0,
		max_excl_x: this.width,
		max_excl_y: this.height)
}

pub func decoder.num_animation_loops() base.u32 {
	return 0
}

pub func decoder.num_decoded_frame_configs() base.u64 {
	if this.call_sequence > 3 {
		return 1
	}
	return 0
}

pub func decoder.num_decoded_frames() base.u64 {
	if this.call_sequence > 4 {
		return 1
	}
	return 0
}

pub func decoder.restart_frame!(index: base.u64, io_position: base.u64) base.status {
	if this.call_sequence < 3 {
		return base."#bad call sequence"
	}
	if args.index <> 0 {
		return base."#bad argument"
	}
	this.call_sequence = 3
	this.frame_config_io_position = args.io_position
	return ok
}

pub func decoder.set_report_metadata!(fourcc: base.u32, report: base.bool) {
	// TODO.
}

pub func decoder.tell_me_more?(dst: base.io_writer, minfo: nptr base.more_information, src: base.io_reader) {
	return base."#no more information"
}

pub func decoder.workbuf_len() base.range_ii_u64 {
	return this.util.make_range_ii_u64(
		min_incl: this.workbuf_length,
		max_incl: this.workbuf_length)
}