/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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
 *
 *      http://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.
 */

package org.apache.sanselan.formats.ico;

import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

import org.apache.sanselan.ImageReadException;
import org.apache.sanselan.ImageWriteException;
import org.apache.sanselan.Sanselan;
import org.apache.sanselan.common.BinaryOutputStream;
import org.apache.sanselan.util.Debug;
import org.apache.sanselan.util.IOUtils;

public class IcoRoundtripTest extends IcoBaseTest
{
	// 16x16 test image
	private static final int[][] image = {
		{0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0},
		{0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0},
		{0,0,0,0,1,0,0,0,0,0,1,1,0,0,0,0},
		{0,0,0,1,1,0,0,0,0,1,0,0,1,0,0,0},
		{0,0,1,0,1,0,0,0,1,0,0,0,0,1,0,0},
		{0,0,0,0,1,0,0,0,1,0,0,0,0,0,0,0},
		{0,0,0,0,1,0,0,0,1,0,0,0,0,0,0,0},
		{0,0,0,0,1,0,0,0,1,0,0,0,0,0,0,0},
		{0,0,0,0,1,0,0,0,1,0,1,1,1,0,0,0},
		{0,0,0,0,1,0,0,0,1,1,0,0,0,1,0,0},
		{0,0,0,0,1,0,0,0,1,0,0,0,0,1,0,0},
		{0,0,0,0,1,0,0,0,1,0,0,0,0,1,0,0},
		{0,0,0,0,1,0,0,0,0,1,0,0,0,1,0,0},
		{0,0,0,0,1,0,0,0,0,0,1,1,1,0,0,0},
		{0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0},
		{0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0}
	};

	private Map generatorMap = new HashMap();

	public IcoRoundtripTest()
	{
		generatorMap.put(Integer.valueOf(1), new GeneratorFor1BitBitmaps());
		generatorMap.put(Integer.valueOf(4), new GeneratorFor4BitBitmaps());
		generatorMap.put(Integer.valueOf(8), new GeneratorFor8BitBitmaps());
		generatorMap.put(Integer.valueOf(16), new GeneratorFor16BitBitmaps());
		generatorMap.put(Integer.valueOf(24), new GeneratorFor24BitBitmaps());
		generatorMap.put(Integer.valueOf(32), new GeneratorFor32BitBitmaps());
	}

	private static interface BitmapGenerator
	{
		byte[] generateBitmap(int foreground, int background, int paletteSize)
				throws IOException, ImageWriteException;
	}

	private class GeneratorFor1BitBitmaps implements BitmapGenerator
	{
		public byte[] generateBitmap(int foreground, int background, int paletteSize)
				throws IOException, ImageWriteException
		{
			ByteArrayOutputStream byteArrayStream = new ByteArrayOutputStream();
			BinaryOutputStream bos = new BinaryOutputStream(byteArrayStream,
					BinaryOutputStream.BYTE_ORDER_LITTLE_ENDIAN);
			// Palette
			bos.write3Bytes(background);
			bos.write(0);
			bos.write3Bytes(foreground);
			bos.write(0);
			for (int i = 2; i < paletteSize; i++)
				bos.write4Bytes(0);
			// Image
			for (int y = 15; y >= 0; y--)
			{
				for (int x = 0; x < 16; x += 8)
				{
					bos.write(
						((0x1 & image[y][x]) << 7) |
						((0x1 & image[y][x+1]) << 6) |
						((0x1 & image[y][x+2]) << 5) |
						((0x1 & image[y][x+3]) << 4) |
						((0x1 & image[y][x+4]) << 3) |
						((0x1 & image[y][x+5]) << 2) |
						((0x1 & image[y][x+6]) << 1) |
						((0x1 & image[y][x+7]) << 0));
				}
				// Pad to multiple of 32 bytes
				bos.write(0);
				bos.write(0);
			}
			// Mask
			for (int y = image.length - 1; y >= 0; y--)
			{
				bos.write(0);
				bos.write(0);
				// Pad to 4 bytes:
				bos.write(0);
				bos.write(0);
			}
			bos.flush();
			return byteArrayStream.toByteArray();
		}
	}

	private class GeneratorFor4BitBitmaps implements BitmapGenerator
	{
		public byte[] generateBitmap(int foreground, int background, int paletteSize)
				throws IOException, ImageWriteException
		{
			ByteArrayOutputStream byteArrayStream = new ByteArrayOutputStream();
			BinaryOutputStream bos = new BinaryOutputStream(byteArrayStream,
					BinaryOutputStream.BYTE_ORDER_LITTLE_ENDIAN);
			// Palette
			bos.write3Bytes(background);
			bos.write(0);
			bos.write3Bytes(foreground);
			bos.write(0);
			for (int i = 2; i < paletteSize; i++)
				bos.write4Bytes(0);
			// Image
			for (int y = 15; y >= 0; y--)
			{
				for (int x = 0; x < 16; x += 2)
				{
					bos.write(((0xf & image[y][x]) << 4) |
						(0xf & image[y][x+1]));
				}
			}
			// Mask
			for (int y = image.length - 1; y >= 0; y--)
			{
				bos.write(0);
				bos.write(0);
				// Pad to 4 bytes:
				bos.write(0);
				bos.write(0);
			}
			bos.flush();
			return byteArrayStream.toByteArray();
		}
	}

	private class GeneratorFor8BitBitmaps implements BitmapGenerator
	{
		public byte[] generateBitmap(int foreground, int background, int paletteSize)
				throws IOException, ImageWriteException
		{
			ByteArrayOutputStream byteArrayStream = new ByteArrayOutputStream();
			BinaryOutputStream bos = new BinaryOutputStream(byteArrayStream,
					BinaryOutputStream.BYTE_ORDER_LITTLE_ENDIAN);
			// Palette
			bos.write3Bytes(background);
			bos.write(0);
			bos.write3Bytes(foreground);
			bos.write(0);
			for (int i = 2; i < paletteSize; i++)
				bos.write4Bytes(0);
			// Image
			for (int y = 15; y >= 0; y--)
			{
				for (int x = 0; x < 16; x++)
				{
					bos.write(image[y][x]);
				}
			}
			// Mask
			for (int y = image.length - 1; y >= 0; y--)
			{
				bos.write(0);
				bos.write(0);
				// Pad to 4 bytes:
				bos.write(0);
				bos.write(0);
			}
			bos.flush();
			return byteArrayStream.toByteArray();
		}
	}

	private class GeneratorFor16BitBitmaps implements BitmapGenerator
	{
		public byte[] generateBitmap(int foreground, int background, int paletteSize)
				throws IOException, ImageWriteException
		{
			ByteArrayOutputStream byteArrayStream = new ByteArrayOutputStream();
			BinaryOutputStream bos = new BinaryOutputStream(byteArrayStream,
					BinaryOutputStream.BYTE_ORDER_LITTLE_ENDIAN);
			// Palette
			for (int i = 0; i < paletteSize; i++)
				bos.write4Bytes(0);
			// Image
			for (int y = 15; y >= 0; y--)
			{
				for (int x = 0; x < 16; x++)
				{
					if (image[y][x] == 1)
						bos.write2Bytes((0x1f & (foreground >> 3)) |
								((0x1f & (foreground >> 11)) << 5) |
								((0x1f & (foreground >> 19)) << 10));
					else
						bos.write2Bytes((0x1f & (background >> 3)) |
								((0x1f & (background >> 11)) << 5) |
								((0x1f & (background >> 19)) << 10));
				}
			}
			// Mask
			for (int y = image.length - 1; y >= 0; y--)
			{
				bos.write(0);
				bos.write(0);
				// Pad to 4 bytes:
				bos.write(0);
				bos.write(0);
			}
			bos.flush();
			return byteArrayStream.toByteArray();
		}
	}

	private class GeneratorFor24BitBitmaps implements BitmapGenerator
	{
		public byte[] generateBitmap(int foreground, int background, int paletteSize)
				throws IOException, ImageWriteException
		{
			ByteArrayOutputStream byteArrayStream = new ByteArrayOutputStream();
			BinaryOutputStream bos = new BinaryOutputStream(byteArrayStream,
					BinaryOutputStream.BYTE_ORDER_LITTLE_ENDIAN);
			// Palette
			for (int i = 0; i < paletteSize; i++)
				bos.write4Bytes(0);
			// Image
			for (int y = 15; y >= 0; y--)
			{
				for (int x = 0; x < 16; x++)
				{
					if (image[y][x] == 1)
						bos.write3Bytes(0xffffff & foreground);
					else
						bos.write3Bytes(0xffffff & background);
				}
			}
			// Mask
			for (int y = image.length - 1; y >= 0; y--)
			{
				bos.write(0);
				bos.write(0);
				// Pad to 4 bytes:
				bos.write(0);
				bos.write(0);
			}
			bos.flush();
			return byteArrayStream.toByteArray();
		}
	}

	private class GeneratorFor32BitBitmaps implements BitmapGenerator
	{
		public byte[] generateBitmap(int foreground, int background, int paletteSize)
				throws IOException, ImageWriteException
		{
			return generate32bitRGBABitmap(foreground, background, paletteSize, true);
		}

		public byte[] generate32bitRGBABitmap(int foreground, int background,
				int paletteSize, boolean writeMask) throws IOException,
				ImageWriteException
		{
			ByteArrayOutputStream byteArrayStream = new ByteArrayOutputStream();
			BinaryOutputStream bos = new BinaryOutputStream(byteArrayStream,
					BinaryOutputStream.BYTE_ORDER_LITTLE_ENDIAN);
			// Palette
			for (int i = 0; i < paletteSize; i++)
				bos.write4Bytes(0);
			// Image
			for (int y = 15; y >= 0; y--)
			{
				for (int x = 0; x < 16; x++)
				{
					if (image[y][x] == 1)
						bos.write4Bytes(foreground);
					else
						bos.write4Bytes(background);
				}
			}
			// Mask
			if (writeMask)
			{
				for (int y = image.length - 1; y >= 0; y--)
				{
					bos.write(0);
					bos.write(0);
					// Pad to 4 bytes:
					bos.write(0);
					bos.write(0);
				}
			}
			bos.flush();
			return byteArrayStream.toByteArray();
		}
	}

	private void writeICONDIR(BinaryOutputStream bos, int reserved, int type, int count)
			throws IOException, ImageWriteException
	{
		bos.write2Bytes(reserved);
		bos.write2Bytes(type);
		bos.write2Bytes(count);
	}

	private void writeICONDIRENTRY(BinaryOutputStream bos, int width, int height,
			int colorCount, int reserved, int planes, int bitCount, int bytesInRes)
			throws IOException, ImageWriteException
	{
		bos.write(width);
		bos.write(height);
		bos.write(colorCount);
		bos.write(reserved);
		bos.write2Bytes(planes);
		bos.write2Bytes(bitCount);
		bos.write4Bytes(bytesInRes);
		bos.write4Bytes(22); // image comes immediately after this
	}

	private void writeBITMAPINFOHEADER(BinaryOutputStream bos, int width, int height,
			int colorPlanes, int bitCount, int compression, int colorsUsed,
			int colorsImportant) throws IOException, ImageWriteException
	{
		// BITMAPINFOHEADER
		bos.write4Bytes(40); // biSize, always 40 for BITMAPINFOHEADER
		bos.write4Bytes(width); // biWidth
		bos.write4Bytes(height); // biHeight
		bos.write2Bytes(colorPlanes); // biPlanes
		bos.write2Bytes(bitCount); // bitCount
		bos.write4Bytes(compression); // biCompression
		bos.write4Bytes(0); // biSizeImage, can be 0 for uncompressed
		bos.write4Bytes(0); // X pixels per metre
		bos.write4Bytes(0); // Y pixels per metre
		bos.write4Bytes(colorsUsed); // colors used, ignored
		bos.write4Bytes(colorsImportant); // colors important
	}

	public void testNormalIcons() throws IOException, ImageWriteException, ImageReadException
	{
		final int foreground = 0xFFF000E0;
		final int background = 0xFF102030;
		for (Iterator it = generatorMap.entrySet().iterator(); it.hasNext(); )
		{
			Map.Entry entry = (Map.Entry) it.next();
			int bitDepth = ((Integer)entry.getKey()).intValue();
			BitmapGenerator bitmapGenerator = (BitmapGenerator) entry.getValue();

			ByteArrayOutputStream baos = new ByteArrayOutputStream();
			BinaryOutputStream bos = new BinaryOutputStream(baos,
					BinaryOutputStream.BYTE_ORDER_LITTLE_ENDIAN);
			byte[] bitmap = bitmapGenerator.generateBitmap(foreground, background,
					(bitDepth <= 8) ? (1 << bitDepth) : 0);
			writeICONDIR(bos, 0, 1, 1);
			writeICONDIRENTRY(bos, 16, 16, 0, 0, 1, bitDepth, 40 + bitmap.length);
			writeBITMAPINFOHEADER(bos, 16, 2*16, 1, bitDepth, 0, 0, 0);
			bos.write(bitmap);
			bos.flush();
			writeAndReadImageData("16x16x" + bitDepth, baos.toByteArray(), foreground, background);
		}
	}

	public void testBadICONDIRENTRYIcons() throws IOException, ImageWriteException, ImageReadException
	{
		final int foreground = 0xFFF000E0;
		final int background = 0xFF102030;
		// Windows ignores the ICONDIRENTRY values when parsing the ICO file.
		for (Iterator it = generatorMap.entrySet().iterator(); it.hasNext(); )
		{
			Map.Entry entry = (Map.Entry) it.next();
			int bitDepth = ((Integer)entry.getKey()).intValue();
			BitmapGenerator bitmapGenerator = (BitmapGenerator) entry.getValue();

			ByteArrayOutputStream baos = new ByteArrayOutputStream();
			BinaryOutputStream bos = new BinaryOutputStream(baos,
					BinaryOutputStream.BYTE_ORDER_LITTLE_ENDIAN);
			byte[] bitmap = bitmapGenerator.generateBitmap(foreground, background,
					(bitDepth <= 8) ? (1 << bitDepth) : 0);
			writeICONDIR(bos, 0, 1, 1);
			writeICONDIRENTRY(bos, 3 /* width, should be 16 */,
					4 /* height, should be 16 */,
					7 /* colorCount, should be 2 or 0 */,
					20 /* reserved, should be 0 */,
					11 /* planes, should be 1 or 0 */,
					19 /* bitCount, should be bitDepth */,
					40 + bitmap.length);
			writeBITMAPINFOHEADER(bos, 16, 2*16, 1, bitDepth, 0, 0, 0);
			bos.write(bitmap);
			bos.flush();
			writeAndReadImageData("16x16x" + bitDepth + "-corrupt-icondirentry",
					baos.toByteArray(), foreground, background);
		}
	}

	public void testColorsUsed() throws IOException, ImageWriteException, ImageReadException
	{
		final int foreground = 0xFFF000E0;
		final int background = 0xFF102030;
		for (Iterator it = generatorMap.entrySet().iterator(); it.hasNext(); )
		{
			Map.Entry entry = (Map.Entry) it.next();
			int bitDepth = ((Integer)entry.getKey()).intValue();
			BitmapGenerator bitmapGenerator = (BitmapGenerator) entry.getValue();

			ByteArrayOutputStream baos = new ByteArrayOutputStream();
			BinaryOutputStream bos = new BinaryOutputStream(baos,
					BinaryOutputStream.BYTE_ORDER_LITTLE_ENDIAN);
			byte[] bitmap = bitmapGenerator.generateBitmap(foreground, background, 2);
			writeICONDIR(bos, 0, 1, 1);
			writeICONDIRENTRY(bos, 3, 4, 7, 20, 11, 19, 40 + bitmap.length);
			writeBITMAPINFOHEADER(bos, 16, 2*16, 1, bitDepth, 0, 2, 0);
			bos.write(bitmap);
			bos.flush();
			writeAndReadImageData("16x16x" + bitDepth + "-custom-palette",
					baos.toByteArray(), foreground, background);
		}
	}

	public void testZeroColorPlanes() throws IOException, ImageWriteException
	{
		final int foreground = 0xFFF000E0;
		final int background = 0xFF102030;
		for (Iterator it = generatorMap.entrySet().iterator(); it.hasNext(); )
		{
			Map.Entry entry = (Map.Entry) it.next();
			int bitDepth = ((Integer)entry.getKey()).intValue();
			BitmapGenerator bitmapGenerator = (BitmapGenerator) entry.getValue();

			ByteArrayOutputStream baos = new ByteArrayOutputStream();
			BinaryOutputStream bos = new BinaryOutputStream(baos,
					BinaryOutputStream.BYTE_ORDER_LITTLE_ENDIAN);
			byte[] bitmap = bitmapGenerator.generateBitmap(foreground, background,
					(bitDepth <= 8) ? (1 << bitDepth) : 0);
			writeICONDIR(bos, 0, 1, 1);
			writeICONDIRENTRY(bos, 16, 16, 0, 0, 1, bitDepth, 40 + bitmap.length);
			writeBITMAPINFOHEADER(bos, 16, 2*16, 0 /* should be 1 */, bitDepth, 0, 0, 0);
			bos.write(bitmap);
			bos.flush();

			boolean threw = false;
			try
			{
				writeAndReadImageData("16x16x" + bitDepth + "-zero-colorPlanes",
						baos.toByteArray(), foreground, background);
			}
			catch (ImageReadException imageReadException)
			{
				threw = true;
			}
			assertTrue(threw);
		}
	}

	public void testBitfieldCompression() throws IOException, ImageWriteException, ImageReadException
	{
		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		BinaryOutputStream bos = new BinaryOutputStream(baos,
				BinaryOutputStream.BYTE_ORDER_LITTLE_ENDIAN);
		byte[] bitmap = new GeneratorFor32BitBitmaps().generate32bitRGBABitmap(
				0xFFFF0000, 0xFFFFFFFF, 0, true);
		writeICONDIR(bos, 0, 1, 1);
		writeICONDIRENTRY(bos, 16, 16, 0, 0, 1, 32, 40 + bitmap.length);
		writeBITMAPINFOHEADER(bos, 16, 2*16, 1, 32, 3 /* BI_BITFIELDS */, 0, 0);
		bos.write4Bytes(0x000000FF); // red mask
		bos.write4Bytes(0x0000FF00); // green mask
		bos.write4Bytes(0x00FF0000); // blue mask
		bos.write(bitmap);
		bos.flush();
		writeAndReadImageData("16x16x32-bitfield-compressed", baos.toByteArray(),
				0xFF0000FF, 0xFFFFFFFF);
	}

	public void test32bitMask() throws IOException, ImageWriteException, ImageReadException
	{
		final int foreground = 0xFFF000E0;
		final int background = 0xFF102030;
		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		BinaryOutputStream bos = new BinaryOutputStream(baos,
				BinaryOutputStream.BYTE_ORDER_LITTLE_ENDIAN);
		// For 32 bit RGBA, the AND mask can be missing:
		byte[] bitmap = new GeneratorFor32BitBitmaps().generate32bitRGBABitmap(
				foreground, background, 0, false);
		writeICONDIR(bos, 0, 1, 1);
		writeICONDIRENTRY(bos, 16, 16, 0, 0, 1, 32, 40 + bitmap.length);
		writeBITMAPINFOHEADER(bos, 16, 2*16, 1, 32, 0, 0, 0);
		bos.write(bitmap);
		bos.flush();
		writeAndReadImageData("16x16x32-no-mask", baos.toByteArray(), foreground, background);
	}

	// FIXME: get the BMP decoder to support alpha, then uncomment below
//	public void testAlphaVersusANDMask() throws IOException, ImageWriteException, ImageReadException
//	{
//		ByteArrayOutputStream baos = new ByteArrayOutputStream();
//		BinaryOutputStream bos = new BinaryOutputStream(baos,
//				BinaryOutputStream.BYTE_ORDER_LITTLE_ENDIAN);
//		byte[] bitmap = new GeneratorFor32BitBitmaps().generate32bitRGBABitmap(
//				0xFF000000, 0x00000000, 0, true);
//		writeICONDIR(bos, 0, 1, 1);
//		writeICONDIRENTRY(bos, 16, 16, 0, 0, 1, 32, 40 + bitmap.length);
//		writeBITMAPINFOHEADER(bos, 16, 2*16, 1, 32, 0, 0, 0);
//		bos.write(bitmap);
//		bos.flush();
//		// The AND mask is fully opaque, yet the fully transparent alpha should win:
//		writeAndReadImageData("16x16x32-alpha-vs-mask", baos.toByteArray(),
//				0xFF000000, 0x00000000);
//	}

	public void testFullyTransparent32bitRGBA() throws IOException, ImageWriteException, ImageReadException
	{
		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		BinaryOutputStream bos = new BinaryOutputStream(baos,
				BinaryOutputStream.BYTE_ORDER_LITTLE_ENDIAN);
		byte[] bitmap = new GeneratorFor32BitBitmaps().generate32bitRGBABitmap(
				0x00000000, 0x00FFFFFF, 0, true);
		writeICONDIR(bos, 0, 1, 1);
		writeICONDIRENTRY(bos, 16, 16, 0, 0, 1, 32, 40 + bitmap.length);
		writeBITMAPINFOHEADER(bos, 16, 2*16, 1, 32, 0, 0, 0);
		bos.write(bitmap);
		bos.flush();
		// Because every pixel is fully trasparent, ***ALPHA GETS IGNORED***:
		writeAndReadImageData("16x16x32-fully-transparent", baos.toByteArray(),
				0xFF000000, 0xFFFFFFFF);
	}

	private void writeAndReadImageData(String description, byte[] rawData,
			int foreground, int background) throws IOException,
			ImageReadException
	{
		// Uncomment to generate ICO files that can be tested with Windows:
		//File exportFile = new File("/tmp/" + description + ".ico");
		//IOUtils.writeToFile(rawData, exportFile);

		File tempFile = createTempFile("temp", ".ico");
		IOUtils.writeToFile(rawData, tempFile);

		BufferedImage dstImage = Sanselan.getBufferedImage(tempFile);

		assertNotNull(dstImage);
		assertTrue(dstImage.getWidth() == image[0].length);
		assertTrue(dstImage.getHeight() == image.length);

		verify(dstImage, foreground, background);
	}

	private void verify(BufferedImage data, int foreground, int background)
	{
		assertNotNull(data);
		assertTrue(data.getHeight() == image.length);

		for (int y = 0; y < data.getHeight(); y++)
		{
			assertTrue(data.getWidth() == image[y].length);
			for (int x = 0; x < data.getWidth(); x++)
			{
				int imageARGB = (image[y][x] == 1) ? foreground : background;
				int dataARGB = data.getRGB(x, y);

				if (imageARGB != dataARGB)
				{
					Debug.debug("x: " + x + ", y: " + y + ", image: " + imageARGB
							+ " (0x" + Integer.toHexString(imageARGB) + ")"
							+ ", data: " + dataARGB + " (0x"
							+ Integer.toHexString(dataARGB) + ")");
				}
				assertTrue(imageARGB == dataARGB);
			}
		}
	}

}