local PNGImage = {}

--
-- libPNGimage by TehSomeLuigi
-- Revision:
PNGImage.rev = 1
--
-- A library to load, edit and save PNGs for OpenComputers
--

--[[
	
	Feel free to use however you wish.
	This header must however be preserved should this be redistributed, even
	if in a modified form.
	
	This software comes with no warranties whatsoever.
	
	2014 TehSomeLuigi

]]--


--local DEFLATE = require("libDEFLATE")
local DeflateLua = require("deflatelua")
local CRC32Lua = require("crc32lua")

local bit = require("bit32")

local PNGImagemetatable = {}
PNGImagemetatable.__index = PNGImage




PNGImage.ColourTypes = {
	Greyscale=0,
	Truecolour=2,
	IndexedColour=3,
	GreyscaleAlpha=4,
	TruecolourAlpha=6
}






local band = bit.band
local rshift = bit.rshift



-- Unpack 32-bit unsigned integer (most-significant-byte, MSB, first)
-- from byte string.
local function __unpack_msb_uint32(s)
	local a,b,c,d = s:byte(1,#s)
	local num = (((a*256) + b) * 256 + c) * 256 + d
	return num
end

local function __write_msb_uint32(fh, int)
	-- MSB B2 B1 LSB
	
	local msb = rshift(band(0xFF000000, int), 24)
	local b2 = rshift(band(0x00FF0000, int), 16)
	local b1 = rshift(band(0x0000FF00, int), 8)
	local lsb = band(0x000000FF, int)
	
	fh:write(string.char(msb))
	fh:write(string.char(b2))
	fh:write(string.char(b1))
	fh:write(string.char(lsb))
end

local function __wbuf_msb_uint32(buf, int)
	-- MSB B2 B1 LSB
	
	local msb = rshift(band(0xFF000000, int), 24)
	local b2 = rshift(band(0x00FF0000, int), 16)
	local b1 = rshift(band(0x0000FF00, int), 8)
	local lsb = band(0x000000FF, int)
	
	buf(msb)
	buf(b2)
	buf(b1)
	buf(lsb)
end

local function __sep_msb_uint32(int)
	-- MSB B2 B1 LSB
	
	local msb = rshift(band(0xFF000000, int), 24)
	local b2 = rshift(band(0x00FF0000, int), 16)
	local b1 = rshift(band(0x0000FF00, int), 8)
	local lsb = band(0x000000FF, int)
	
	return msb, b2, b1, lsb
end

local function __pack_msb_uint16(int)
	return string.char(rshift(band(int, 0xFF00), 8), band(int, 0xFF))
end

local function __pack_lsb_uint16(int)
	return string.char(band(int, 0xFF), rshift(band(int, 0xFF00), 8))
end

-- Read 32-bit unsigned integer (most-significant-byte, MSB, first) from file.
local function __read_msb_uint32(fh)
	return __unpack_msb_uint32(fh:read(4))
end

-- Read unsigned byte (integer) from file
local function __read_byte(fh)
	return fh:read(1):byte()
end



local function getBitWidthPerPixel(ihdr)
	if ihdr.color_type == PNGImage.ColourTypes.Greyscale then -- Greyscale
		return ihdr.bit_depth
	end
	if ihdr.color_type == PNGImage.ColourTypes.Truecolour then -- Truecolour
		return ihdr.bit_depth * 3
	end
	if ihdr.color_type == PNGImage.ColourTypes.IndexedColour then -- Indexed-colour
		return ihdr.bit_depth
	end
	if ihdr.color_type == PNGImage.ColourTypes.GreyscaleAlpha then -- Greyscale + Alpha
		return ihdr.bit_depth * 2
	end
	if ihdr.color_type == PNGImage.ColourTypes.TruecolourAlpha then -- Truecolour + Alpha
		return ihdr.bit_depth * 4
	end
end

local function getByteWidthPerScanline(ihdr)
	return math.ceil((ihdr.width * getBitWidthPerPixel(ihdr)) / 8)
end

local outssmt = {}

function outssmt:__call(write)
	self.str = self.str .. string.char(write)
end

function outssmt.OutStringStream()
	local outss = {str=""}
	setmetatable(outss, outssmt)
	return outss
end




local function __parse_IHDR(fh, len)
	if len ~= 13 then
		error("PNG IHDR Corrupt - should be 13 bytes long")
	end
	
	local ihdr = {}
	
	ihdr.width = __read_msb_uint32(fh)
	ihdr.height = __read_msb_uint32(fh)
	ihdr.bit_depth = __read_byte(fh)
	ihdr.color_type = __read_byte(fh)
	ihdr.compression_method = __read_byte(fh)
	ihdr.filter_method = __read_byte(fh)
	ihdr.interlace_method = __read_byte(fh)
	
	--[[
	print("width=", ihdr.width)
	print("height=", ihdr.height)
	print("bit_depth=", ihdr.bit_depth)
	print("color_type=", ihdr.color_type)
	print("compression_method=", ihdr.compression_method)
	print("filter_method=", ihdr.filter_method)
	print("interlace_method=", ihdr.interlace_method)
	]]--
	
	return ihdr
end

--[[
local function __parse_IDAT(fh, len, commeth)
	if commeth ~= 0 then
		error("Only zlib/DEFLATE compression supported")
	end
	
	local d, msg = DEFLATE.inflate(fh, len);
	
	if not d then
		return nil, msg
	end
	
	local oh = io.open('dump.dat', 'wb')
	oh:write(d.dat)
	oh:close()
	
	return true
end
]]--

local function __parse_IDAT(fh, len, commeth, outss)
	if commeth ~= 0 then
		error("Only zlib/DEFLATE compression supported")
	end
	
	
	--local oh = io.open('dump.dat', 'wb')
	--oh:write(d.dat)
	
	--local ph = io.open('pass.dat', 'wb')
	
	local input = fh:read(len)
	
	--ph:write(input)
	--ph:close()
	
	local cfg = {input=input, output=outss, disable_crc=true}
	
	DeflateLua.inflate_zlib(cfg)
	
	--oh:close()
	
--	if not d then
--		return nil, msg
--	end
	
	return true
end



local function getPNGStdByteAtXY(ihdr, oss, x, y)
	local bpsl = getByteWidthPerScanline(ihdr) -- don't include filterType byte -- we don't store that after it has been read
	if (x <= 0) or (y <= 0) then
		return 0 -- this is what the spec says we should return when the coordinate is out of bounds -- in this part of the code, the coordinates are ONE-BASED like in good Lua
	end
	local offset_by_y = (y - 1) * bpsl
	-- now read it!
	local idx = offset_by_y + x
	return oss.str:sub(idx, idx):byte()
end


local function __paeth_predictor(a, b, c)
	local p = a + b - c
	local pa = math.abs(p - a)
	local pb = math.abs(p - b)
	local pc = math.abs(p - c)
	if pa <= pb and pa <= pc then
		return a
	elseif pb <= pc then
		return b
	else
		return c
	end
end



local function __parse_IDAT_effective_bytes(outss, ihdr)
	local bpsl = getByteWidthPerScanline(ihdr)
	local bypsl = math.ceil(getBitWidthPerPixel(ihdr) / 8)
	
	if outss.str:len() == 0 then
		error("Empty string: outss")
	end
	
	local bys = DeflateLua.stringToBytestream(outss.str)
	
	if not bys then
		error("Did not get a bytestream from string", bys, outss)
	end
	
	local out2 = outssmt.OutStringStream() -- __callable table with metatable that stores what you give it
	
	local y = 0
	
	-- x the byte being filtered;
	-- a the byte corresponding to x in the pixel immediately before the pixel containing x (or the byte immediately before x, when the bit depth is less than 8);
	-- b the byte corresponding to x in the previous scanline;
	-- c the byte corresponding to b in the pixel immediately before the pixel containing b (or the byte immediately before b, when the bit depth is less than 8).
	
	while true do
		local filterType = bys:read()
		
		if filterType == nil then
			break
		end
		
		y = y + 1
		
		for x = 1, bpsl do
			--[..c..][..b..]
			--[..a..][..x <--- what is being processed (x)
			local a = getPNGStdByteAtXY(ihdr, out2, x - bypsl, y)
			local b = getPNGStdByteAtXY(ihdr, out2, x, y - 1)
			local c = getPNGStdByteAtXY(ihdr, out2, x - bypsl, y - 1)
			
			local outVal = 0
			
			if filterType == 0 then
				-- Recon(x) = Filt(x)
				outVal = bys:read()
			elseif filterType == 1 then
				-- Recon(x) = Filt(x) + Recon(a)
				outVal = bys:read() + a
			elseif filterType == 2 then
				-- Recon(x) = Filt(x) + Recon(b)
				outVal = bys:read() + b
			elseif filterType == 3 then
				-- Recon(x) = Filt(x) + floor((Recon(a) + Recon(b)) / 2)
				outVal = bys:read() + math.floor((a + b) / 2)
			elseif filterType == 4 then
				-- Recon(x) = Filt(x) + PaethPredictor(Recon(a), Recon(b), Recon(c))
				outVal = bys:read() + __paeth_predictor(a, b, c)
			else
				error("Unsupported Filter Type: " .. tostring(filterType))
			end
			
			outVal = outVal % 256
			
			out2(outVal)
		end
	end
	
	return out2
end



local function __newPNGImage()
	local pngi = {}
	setmetatable(pngi, PNGImagemetatable)
	return pngi
end

function PNGImage.newFromFile(fn)
	local fh = io.open(fn, 'rb')
	if not fh then
		error("Could not open PNG file")
	end
	return PNGImage.newFromFileHandle(fh)
end

function PNGImage.newFromScratch(width, height, bkcol)
	local pngi = __newPNGImage()
	
	local width = tonumber(width)
	local height = tonumber(height)
	
	if (not width) or (width < 1) or (math.floor(width) ~= width) then
		error("Invalid param #1 (width) to PNGImage.newFromScratch - integer (>=0) expected")
	end
	if (not height) or (height < 1) or (math.floor(height) ~= height) then
		error("Invalid param #2 (height) to PNGImage.newFromScratch - integer (>=0) expected")
	end
	
	local bkg = {0, 0, 0, 0} -- Transparency
	
	if type(bkcol) == "table" then
		if #bkcol == 3 then
			bkg = {bkcol[1], bkcol[2], bkcol[3], 255} -- Opaque colour
		elseif #bkcol == 4 then
			bkg = {bkcol[1], bkcol[2], bkcol[3], bkcol[4]} -- Defined
		else
			error("Invalid format for param #3 (bkcol) to PNGImage.newFromScratch: nil or table expected, but the table must be of format {r, g, b} or {r, g, b, a} -- Invalid table")
		end
	elseif bkcol ~= nil then
		error("Invalid format for param #3 (bkcol) to PNGImage.newFromScratch: nil or table expected, but the table must be of format {r, g, b} or {r, g, b, a} -- Parameter is not nil or table")
	end
	
	bkg[1] = tonumber(bkg[1])
	if bkg[1] == nil then
		error("PNGImage.newFromScratch: bkg[R] is not numeric")
	end
	bkg[2] = tonumber(bkg[2])
	if bkg[2] == nil then
		error("PNGImage.newFromScratch: bkg[G] is not numeric")
	end
	bkg[3] = tonumber(bkg[3])
	if bkg[3] == nil then
		error("PNGImage.newFromScratch: bkg[B] is not numeric")
	end
	bkg[4] = tonumber(bkg[4])
	if bkg[4] == nil then
		error("PNGImage.newFromScratch: bkg[A] is not numeric")
	end
	
	pngi.ihdr = {
		width = width,
		height = height,
		bit_depth = 8,
		color_type = 6,
		compression_method = 0,
		filter_method = 0,
		interlace_method = 0
	}
	
	-- do this for every pixel.. i.e string.rep(str, w*h)
	pngi.data = string.byte(bkg[1], bkg[2], bkg[3], bkg[4]):rep(width * height)
	
	return pngi
end

function PNGImage.newFromFileHandle(fh)
	local pngi = __newPNGImage()
	local expecting = "\137\080\078\071\013\010\026\010"
	if fh:read(8) ~= expecting then -- check the 8-byte PNG header exists
		error("Not a PNG file")
	end
	
	local ihdr
	
	local outss = outssmt.OutStringStream()
	
	while true do
		local len = __read_msb_uint32(fh)
		local stype = fh:read(4)
		
		if stype == 'IHDR' then
			ihdr, msg = __parse_IHDR(fh, len)
		elseif stype == 'IDAT' then
			local res, msg = __parse_IDAT(fh, len, ihdr.compression_method, outss)
		else
			fh:read(len) -- dummy read
		end
		
		local crc = __read_msb_uint32(fh)
		
		-- print("chunk:", "type=", stype, "len=", len, "crc=", crc)
		
		if stype == 'IEND' then
			break
		end
	end
	
	fh:close()
	
	
	if ihdr.filter_method ~= 0 then
		error("Unsupported Filter Method: " .. ihdr.filter_method)
	end
	
	if ihdr.interlace_method ~= 0 then
		error("Unsupported Interlace Method (Interlacing is currently unsupported): " .. ihdr.interlace_method)
	end
	
	if ihdr.color_type ~= PNGImage.ColourTypes.TruecolourAlpha and ihdr.color_type ~= PNGImage.ColourTypes.Truecolour then
		error("Currently, only Truecolour and Truecolour+Alpha images are supported.")
	end
	
	if ihdr.bit_depth ~= 8 then
		error("Currently, only images with a bit depth of 8 are supported.")
	end
	
	--[[
	local oh = io.open('before-decode.dat', 'wb')
	oh:write(outss.str)
	oh:close()
	]]--
	
	-- now parse the IDAT chunks
	local out2 = __parse_IDAT_effective_bytes(outss, ihdr)
	
	if ihdr.color_type == PNGImage.ColourTypes.Truecolour then
		-- add an alpha layer so it effectively becomes RGBA, not RGB
		local inp = out2.str
		out2 = outssmt.OutStringStream()
		
		for i=1, ihdr.width*ihdr.height do
			local b = ((i - 1)*3) + 1
			out2(inp:byte(b)) -- R
			out2(inp:byte(b + 1)) -- G
			out2(inp:byte(b + 2)) -- B
			out2(255) -- A
		end
	end
	
	pngi.ihdr = ihdr
	pngi.data = out2.str
	
	--[[
	local oh = io.open('effective.dat', 'wb')
	oh:write(out2.str)
	oh:close()
	]]--
	
	return pngi
end


-- Warning: Co-ordinates are Zero-based but strings are 1-based
function PNGImage:getByteOffsetForPixel(x, y)
	return (((y * self.ihdr.width) + x) * 4) + 1
end

function PNGImage:getPixel(x, y)
	local off = self:getByteOffsetForPixel(x, y)
	return self.data:byte(off, off + 3)
end

function PNGImage:setPixel(x, y, col)
	local off = self:getByteOffsetForPixel(x, y)
	self.data = table.concat({self.data:sub(1, off - 1), string.char(col[1], col[2], col[3], col[4]), self.data:sub(off + 4)})
end

function PNGImage:lineXAB(ax, y, bx, col)
	for x=ax, bx do
		self:setPixel(x, y, col)
	end
end

function PNGImage:lineYAB(x, ay, by, col)
	for y=ay, by do
		self:setPixel(x, y, col)
	end
end

function PNGImage:lineRectangleAB(ax, ay, bx, by, col)
	self:lineXAB(ax, ay, bx, col)
	self:lineXAB(ax, by, bx, col)
	self:lineYAB(ax, ay, by, col)
	self:lineYAB(bx, ay, by, col)
end

function PNGImage:fillRectangleAB(ax, ay, bx, by, col)
	for x=ax, bx do
		for y=ay, by do
			self:setPixel(x, y, col)
		end
	end
end

function PNGImage:saveToFile(fn)
	local fh = io.open(fn, 'wb')
	if not fh then
		error("Could not open for writing: " .. fn)
	end
	self:saveToFileHandle(fh)
	fh:close()
end

function PNGImage:getSize()
	return self.ihdr.width, self.ihdr.height
end

function PNGImage:generateRawIDATData(outbuf)
	for y = 0, self.ihdr.height - 1 do
		outbuf(0) -- filter type is 0 (Filt(x) = Orig(x))
		for x = 0, self.ihdr.width - 1 do
			local r, g, b, a = self:getPixel(x, y)
			outbuf(r)
			outbuf(g)
			outbuf(b)
			outbuf(a)
		end
	end
end

local ZLIB_LITERAL_LIMIT = 65535

local function __raw_to_literalZLIB(inbuf)
	local outstr = string.char(8, 29) -- zlib headers
	
	while inbuf.str:len() > 0 do
		if inbuf.str:len() > ZLIB_LITERAL_LIMIT then
			outstr = outstr .. string.char(0) -- LITERAL[00] FINAL[0]
		else
			outstr = outstr .. string.char(1) -- LITERAL[00] FINAL[1]
		end
		local min = math.min(ZLIB_LITERAL_LIMIT, inbuf.str:len())
		outstr = table.concat({outstr, __pack_msb_uint16(min), __pack_msb_uint16(bit.bnot(min)), inbuf.str:sub(1, min)})
		inbuf.str = inbuf.str:sub(min + 1)
	end
	
	local adler32 = 1
	
	for i = 1, outstr:len() do
		adler32 = DeflateLua.adler32(outstr:byte(i), adler32)
	end
	
	outstr = table.concat({outstr, string.char(__sep_msb_uint32(adler32))})
	
	return outstr
end

function PNGImage:saveToFileHandle(fh)
	-- basic PNG 'encoder'
	-- most likely temporary
	local expecting = "\137\080\078\071\013\010\026\010"
	fh:write(expecting)
	
	local outbuf = outssmt.OutStringStream()
	
	__wbuf_msb_uint32(outbuf, self.ihdr.width)
	__wbuf_msb_uint32(outbuf, self.ihdr.height)
	
	outbuf(8) -- bit depth
	outbuf(6) -- colour type
	outbuf(0) -- compression method
	outbuf(0) -- filter method
	outbuf(0) -- interlace method
	
	__write_msb_uint32(fh, outbuf.str:len()) -- length field
	fh:write("IHDR") -- name of chunk
	fh:write(outbuf.str) -- chunk data
	__write_msb_uint32(fh, CRC32Lua.crc32_string("IHDR" .. outbuf.str)) -- chunk data CRC32
	outbuf.str = "" -- reset buffer
	
	-- now onto the IDAT
	
	self:generateRawIDATData(outbuf)
	
	local zlibstr = __raw_to_literalZLIB(outbuf)
	
	local tmpstr = ""
	while zlibstr:len() > 0 do
		local min = math.min(ZLIB_LITERAL_LIMIT, zlibstr:len())
		tmpstr = zlibstr:sub(1, min)
		zlibstr = zlibstr:sub(min + 1)
		
		__write_msb_uint32(fh, tmpstr:len()) -- length field
		fh:write("IDAT") -- name of chunk
		fh:write(tmpstr) -- chunk data
		__write_msb_uint32(fh, CRC32Lua.crc32_string("IDAT" .. tmpstr)) -- chunk data CRC32
		tmpstr = "" -- reset
	end
	
	__write_msb_uint32(fh, 0) -- length field
	fh:write("IEND") -- name of chunk
	-- fh:write(tmpstr) -- chunk data
	__write_msb_uint32(fh, CRC32Lua.crc32_string("IEND")) -- chunk data CRC32
end



return PNGImage