package com.jotuntech.sketcher.client;

import java.awt.Point;
import java.awt.image.BufferedImage;
import java.awt.image.DataBufferInt;
import java.io.DataOutput;
import java.io.IOException;
import java.io.RandomAccessFile;
import java.util.Arrays;
import java.util.Map;

import com.jotuntech.sketcher.common.BitmapLayer;
import com.jotuntech.sketcher.common.BitmapTile;
import com.jotuntech.sketcher.common.Layer;
import com.jotuntech.sketcher.common.TwoWayHashMap;

public class PSDEncoder {
	private RandomAccessFile raf;
	private BufferedImage image;
	private BitmapLayer[] layers;
	private long[][] channelLengthPositions;

	public PSDEncoder(RandomAccessFile raf, BufferedImage image, TwoWayHashMap<Integer, Layer> layers) throws IOException {
		this.raf = raf;
		this.image = image;

		int layersLength = 0;
		for(Layer l : layers.values()) {
			if(l instanceof BitmapLayer) {
				++layersLength;
			}
		}

		this.layers = new BitmapLayer[layersLength];
		int layerIndex = 0;
		for(Layer l : layers.values()) {
			if(l instanceof BitmapLayer) {
				this.layers[layerIndex++] = (BitmapLayer) l;
			}
		}

		channelLengthPositions = new long[this.layers.length][4];
	}


	public void encode() throws IOException {
		/* Signature */
		raf.write(new byte[] {'8', 'B', 'P', 'S'});

		/* Version */
		raf.writeShort(1);

		/* Reserved */
		raf.write(new byte[] {0, 0, 0, 0, 0, 0});

		/* Number of channels */
		raf.writeShort(4);

		/* Image height */
		raf.writeInt(image.getHeight());

		/* Image width */
		raf.writeInt(image.getWidth());

		/* Bits per channel */
		raf.writeShort(8);

		/* Color mode: RGB color */
		raf.writeShort(3);

		/* Length of color mode data section */
		raf.writeInt(0);

		/* Length of image resources section */
		raf.writeInt(0);

		/* Layer and mask information */
		writeLayerAndMaskInfo(raf);

		/* Image data */
		writeImagePixelData(raf);
	}

	private void writeLayerAndMaskInfo(RandomAccessFile raf) throws IOException {
		long lengthPointer = raf.getFilePointer();
		raf.writeInt(0);

		/* Layer info */
		writeLayerInfo(raf);

		/* Mask info */
		writeMaskInfo(raf);

		/* Go back and write length */
		long endPointer = raf.getFilePointer();
		raf.seek(lengthPointer);
		raf.writeInt((int)(endPointer - lengthPointer - 4));
		raf.seek(endPointer);
	}


	private void writeLayerInfo(RandomAccessFile odo) throws IOException {
		System.err.println("Writing layer info.");

		long lengthPointer = raf.getFilePointer();
		odo.writeInt(0);

		/* Layers struct */
		writeLayersStruct(odo);

		/* Pixel data */
		writeAllLayersPixelData(odo);

		/* Go back and write length */
		long endPointer = raf.getFilePointer();
		raf.seek(lengthPointer);
		raf.writeInt((int)(endPointer - lengthPointer - 4));
		raf.seek(endPointer);
	}

	private void writeLayersStruct(RandomAccessFile odo) throws IOException {
		System.err.println("Writing layers struct.");

		/* Layer count */
		odo.writeShort(layers.length);

		/* Layers */
		for(int layerIndex = 0; layerIndex < layers.length; layerIndex++) {
			writeLayer(odo, layerIndex);
		}
	}

	private void writeLayer(RandomAccessFile odo, int layerIndex) throws IOException {
		/* Get layer */
		BitmapLayer l = layers[layerIndex];

		System.err.println("Writing layer: " + l.getName());

		/* Layer top */
		odo.writeInt(0);

		/* Layer left */
		odo.writeInt(0);

		/* Layer bottom */
		odo.writeInt(image.getHeight());

		/* Layer right */
		odo.writeInt(image.getWidth());

		/* Number of channels */
		odo.writeShort(4);

		/* Channel length info */
		writeChannelLengthInfo(odo, layerIndex);

		/* Blend mode signature */
		odo.write(new byte[] {'8', 'B', 'I', 'M'});

		/* Blend mode key */
		odo.write(new byte[] {'n', 'o', 'r', 'm'});

		/* Opacity */
		odo.writeByte(Math.round(l.getOpacity() * 255f));

		/* Clipping: Base */
		odo.writeByte(0);

		/* Flags: Transparency not protected and visible */
		odo.writeByte(0x80);

		/* Filler */
		odo.writeByte(0);

		/* Prepare layer name */
		String name = l.getName();
		int origLength = name.length();
		int moduloLength = (origLength + 1) % 4;
		int padLength = 0;
		if(moduloLength > 0) {
			padLength = 4 - moduloLength;
			for(int i = 0; i < padLength; i++) {
				name += "\000";
			}
		}

		/* Extra data size */
		odo.writeInt(4 + 44 + 1 + origLength + padLength);

		/* Layer mask data size */
		odo.writeInt(0);

		/* Layer blending ranges */
		writeLayerBlendingRanges(odo);

		odo.writeByte(origLength);
		odo.writeBytes(name);
	}

	private void writeChannelLengthInfo(RandomAccessFile raf, int layerIndex) throws IOException {
		System.err.println("Writing channel length info.");

		/* Channel ID: Transparency Mask */
		raf.writeShort(-1);

		/* Channel length */
		channelLengthPositions[layerIndex][0] = raf.getFilePointer();
		raf.writeInt(0);

		/* Channel ID: Red */
		raf.writeShort(0);

		/* Channel length */
		channelLengthPositions[layerIndex][1] = raf.getFilePointer();
		raf.writeInt(0);

		/* Channel ID: Green */
		raf.writeShort(1);

		/* Channel length */
		channelLengthPositions[layerIndex][2] = raf.getFilePointer();
		raf.writeInt(0);

		/* Channel ID: Blue */
		raf.writeShort(2);

		/* Channel length */
		channelLengthPositions[layerIndex][3] = raf.getFilePointer();
		raf.writeInt(0);
	}

	private void writeLayerBlendingRanges(DataOutput odo) throws IOException {
		System.err.println("Writing layer blending ranges.");

		/* Length */
		odo.writeInt(40);

		/* Composite gray blend source */
		odo.writeInt(0x0000FFFF);

		/* Composite gray blend destination */
		odo.writeInt(0x0000FFFF);

		/* First channel source range */
		odo.writeInt(0x0000FFFF);

		/* First channel destination range */
		odo.writeInt(0x0000FFFF);

		/* Second channel source range */
		odo.writeInt(0x0000FFFF);

		/* Second channel destination range */
		odo.writeInt(0x0000FFFF);

		/* Third channel source range */
		odo.writeInt(0x0000FFFF);

		/* Third channel destination range */
		odo.writeInt(0x0000FFFF);

		/* Fourth channel source range */
		odo.writeInt(0x0000FFFF);

		/* Fourth channel destination range */
		odo.writeInt(0x0000FFFF);
	}

	private void writeAllLayersPixelData(RandomAccessFile odo) throws IOException {
		System.out.println("Writing pixel layer data.");
		for(int layerIndex = 0; layerIndex < layers.length; layerIndex++) {
			writeLayerPixelData(odo, layerIndex);
		}
	}

	private void writeMaskInfo(RandomAccessFile raf) throws IOException {
		System.err.println("Writing mask info.");

		/* Mask info length */
		long lengthPointer = raf.getFilePointer();
		raf.writeInt(0);

		/* Overlay color space */
		raf.writeShort(0);

		/* Color components */
		raf.writeShort(0);
		raf.writeShort(0);
		raf.writeShort(0);
		raf.writeShort(0);

		/* Opacity */
		raf.writeShort(0);

		/* Kind: Use value stored per layer */
		raf.writeByte(128);

		/* Filler */
		raf.writeByte(0);

		/* Go back and write length */
		long endPointer = raf.getFilePointer();
		raf.seek(lengthPointer);
		raf.writeInt((int)(endPointer - lengthPointer - 4));
		raf.seek(endPointer);
	}

	public void writeImagePixelData(RandomAccessFile odo) throws IOException {
		System.out.println("Writing image data.");

		int[] pixels = ((DataBufferInt)image.getRaster().getDataBuffer()).getData();
		byte[] scanline = new byte[image.getWidth()];

		/* Compression: RLE */
		odo.writeShort(1);

		/* Placeholder scanline bytecounts */
		long redCountsPosition = raf.getFilePointer();
		for(int i = 0; i < image.getHeight(); i++) {
			odo.writeShort(0);
		}
		long greenCountsPosition = raf.getFilePointer();
		for(int i = 0; i < image.getHeight(); i++) {
			odo.writeShort(0);
		}
		long blueCountsPosition = raf.getFilePointer();
		for(int i = 0; i < image.getHeight(); i++) {
			odo.writeShort(0);
		}
		long alphaCountsPosition = raf.getFilePointer();
		for(int i = 0; i < image.getHeight(); i++) {
			odo.writeShort(0);
		}

		PackBitsOutputStream pbos = new PackBitsOutputStream(raf);

		/* Image data - Red */
		for(int y = 0; y < image.getHeight(); y++) {
			for(int x = 0; x < scanline.length; x++) {
				scanline[x] = (byte) ((pixels[y * scanline.length + x] >> 16) & 0xFF);
			}
			long scanlineStart = raf.getFilePointer();
			pbos.write(scanline);
			pbos.flush();
			long scanlineStop = raf.getFilePointer();
			raf.seek(redCountsPosition + y * 2);
			raf.writeShort((short)(scanlineStop - scanlineStart));
			raf.seek(scanlineStop);
		}

		/* Image data - Green */
		for(int y = 0; y < image.getHeight(); y++) {
			for(int x = 0; x < scanline.length; x++) {
				scanline[x] = (byte) ((pixels[y * scanline.length + x] >> 8) & 0xFF);
			}
			long scanlineStart = raf.getFilePointer();
			pbos.write(scanline);
			pbos.flush();
			long scanlineStop = raf.getFilePointer();
			raf.seek(greenCountsPosition + y * 2);
			raf.writeShort((short)(scanlineStop - scanlineStart));
			raf.seek(scanlineStop);
		}

		/* Image data - Blue */
		for(int y = 0; y < image.getHeight(); y++) {
			for(int x = 0; x < scanline.length; x++) {
				scanline[x] = (byte) (pixels[y * scanline.length + x] & 0xFF);
			}
			long scanlineStart = raf.getFilePointer();
			pbos.write(scanline);
			pbos.flush();
			long scanlineStop = raf.getFilePointer();
			raf.seek(blueCountsPosition + y * 2);
			raf.writeShort((short)(scanlineStop - scanlineStart));
			raf.seek(scanlineStop);
		}

		/* Image data - Alpha */
		for(int y = 0; y < image.getHeight(); y++) {
			for(int x = 0; x < scanline.length; x++) {
				scanline[x] = (byte) ((pixels[y * scanline.length + x] >> 24) & 0xFF);
			}
			long scanlineStart = raf.getFilePointer();
			pbos.write(scanline);
			pbos.flush();
			long scanlineStop = raf.getFilePointer();
			raf.seek(alphaCountsPosition + y * 2);
			raf.writeShort((short)(scanlineStop - scanlineStart));
			raf.seek(scanlineStop);
		}
	}

	public void writeLayerPixelData(RandomAccessFile raf, int layerIndex) throws IOException {
		System.out.println("Writing image data.");

		byte[] scanline = new byte[image.getWidth()];

		/* Get tiles */
		Map<Point, BitmapTile> tiles = layers[layerIndex].getTiles();

		PackBitsOutputStream pbos = new PackBitsOutputStream(raf);

		long startOfAlpha = raf.getFilePointer();
		int alphaLength;

		/* Compression: RLE */
		raf.writeShort(1);

		/* Alpha scanline counts */
		long alphaCountsPosition = raf.getFilePointer();
		for(int y = 0; y < image.getHeight(); y++) {
			raf.writeShort(0);
		}

		/* Alpha image data */
		for(int y = 0; y < image.getHeight(); y++) {
			int yModulus = y & BitmapTile.SIZE_MASK;
			for(int x = 0; x < scanline.length; x += BitmapTile.SIZE) {
				BitmapTile t = tiles.get(new Point(x >> BitmapTile.SIZE_2, y >> BitmapTile.SIZE_2));
				if(t == null) {
					Arrays.fill(scanline, x, x + BitmapTile.SIZE, (byte) 0);
				} else {
					int[] tilePixels = t.getPixels();
					for(int offset = 0; offset < BitmapTile.SIZE; offset++) {
						scanline[x + offset] = (byte) ((tilePixels[(yModulus << BitmapTile.SIZE_2) + offset] >> 24) & 0xFF);
					}
				}
			}
			long scanlineStart = raf.getFilePointer();
			pbos.write(scanline);
			pbos.flush();
			long scanlineEnd = raf.getFilePointer();
			raf.seek(alphaCountsPosition + y * 2);
			raf.writeShort((short) (scanlineEnd - scanlineStart));
			raf.seek(scanlineEnd);
		}

		long endOfAlpha = raf.getFilePointer();
		alphaLength = (int) (endOfAlpha - startOfAlpha);
		raf.seek(channelLengthPositions[layerIndex][0]);
		raf.writeInt(alphaLength);
		raf.seek(endOfAlpha);

		long startOfRed = raf.getFilePointer();
		int redLength;

		/* Compression: RLE */
		raf.writeShort(1);

		/* Red scanline counts */
		long redCountsPosition = raf.getFilePointer();
		for(int y = 0; y < image.getHeight(); y++) {
			raf.writeShort(0);
		}

		/* Red image data */
		for(int y = 0; y < image.getHeight(); y++) {
			int yModulus = y & BitmapTile.SIZE_MASK;
			for(int x = 0; x < scanline.length; x += BitmapTile.SIZE) {
				BitmapTile t = tiles.get(new Point(x >> BitmapTile.SIZE_2, y >> BitmapTile.SIZE_2));
				if(t == null) {
					Arrays.fill(scanline, x, x + BitmapTile.SIZE, (byte) 0);
				} else {
					for(int offset = 0; offset < BitmapTile.SIZE; offset++) {
						int[] tilePixels = t.getPixels();
						scanline[x + offset] = (byte) ((tilePixels[(yModulus << BitmapTile.SIZE_2) + offset] >> 16) & 0xFF);
					}
				}
			}
			long scanlineStart = raf.getFilePointer();
			pbos.write(scanline);
			pbos.flush();
			long scanlineEnd = raf.getFilePointer();
			raf.seek(redCountsPosition + y * 2);
			raf.writeShort((short) (scanlineEnd - scanlineStart));
			raf.seek(scanlineEnd);
		}

		long endOfRed = raf.getFilePointer();
		redLength = (int) (endOfRed - startOfRed);
		raf.seek(channelLengthPositions[layerIndex][1]);
		raf.writeInt(redLength);
		raf.seek(endOfRed);

		long startOfGreen = raf.getFilePointer();
		int greenLength;

		/* Compression: RLE */
		raf.writeShort(1);

		/* Green scanline counts */
		long greenCountsPosition = raf.getFilePointer();
		for(int y = 0; y < image.getHeight(); y++) {
			raf.writeShort(0);
		}

		/* Green image data */
		for(int y = 0; y < image.getHeight(); y++) {
			int yModulus = y & BitmapTile.SIZE_MASK;
			for(int x = 0; x < scanline.length; x += BitmapTile.SIZE) {
				BitmapTile t = tiles.get(new Point(x >> BitmapTile.SIZE_2, y >> BitmapTile.SIZE_2));
				if(t == null) {
					Arrays.fill(scanline, x, x + BitmapTile.SIZE, (byte) 0);
				} else {
					for(int offset = 0; offset < BitmapTile.SIZE; offset++) {
						int[] tilePixels = t.getPixels();
						scanline[x + offset] = (byte) ((tilePixels[(yModulus << BitmapTile.SIZE_2) + offset] >> 8) & 0xFF);
					}
				}
			}
			long scanlineStart = raf.getFilePointer();
			pbos.write(scanline);
			pbos.flush();
			long scanlineEnd = raf.getFilePointer();
			raf.seek(greenCountsPosition + y * 2);
			raf.writeShort((short)(scanlineEnd - scanlineStart));
			raf.seek(scanlineEnd);
		}

		long endOfGreen = raf.getFilePointer();
		greenLength = (int) (endOfGreen - startOfGreen);
		raf.seek(channelLengthPositions[layerIndex][2]);
		raf.writeInt(greenLength);
		raf.seek(endOfGreen);

		long startOfBlue = raf.getFilePointer();
		int blueLength;

		/* Compression: RLE */
		raf.writeShort(1);

		/* Blue scanline counts */
		long blueCountsPosition = raf.getFilePointer();
		for(int y = 0; y < image.getHeight(); y++) {
			raf.writeShort(0);
		}

		/* Blue image data */
		for(int y = 0; y < image.getHeight(); y++) {
			int yModulus = y & BitmapTile.SIZE_MASK;
			for(int x = 0; x < scanline.length; x += BitmapTile.SIZE) {
				BitmapTile t = tiles.get(new Point(x >> BitmapTile.SIZE_2, y >> BitmapTile.SIZE_2));
				if(t == null) {
					Arrays.fill(scanline, x, x + BitmapTile.SIZE, (byte) 0);
				} else {
					int[] tilePixels = t.getPixels();
					for(int offset = 0; offset < BitmapTile.SIZE; offset++) {
						scanline[x + offset] = (byte) (tilePixels[(yModulus << BitmapTile.SIZE_2) + offset] & 0xFF);
					}
				}
			}
			long scanlineStart = raf.getFilePointer();
			pbos.write(scanline);
			pbos.flush();
			long scanlineEnd = raf.getFilePointer();
			raf.seek(blueCountsPosition + y * 2);
			raf.writeShort((short)(scanlineEnd - scanlineStart));
			raf.seek(scanlineEnd);
		}

		long endOfBlue = raf.getFilePointer();
		blueLength = (int) (endOfBlue - startOfBlue);
		raf.seek(channelLengthPositions[layerIndex][3]);
		raf.writeInt(blueLength);
		raf.seek(endOfBlue);
	}
}