/* Copyright (c) 2014, David A. Clunie DBA Pixelmed Publishing. All rights reserved. */

package com.pixelmed.codec.jpeg;

import java.awt.Rectangle;
import java.awt.Shape;

import java.io.ByteArrayOutputStream;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.InputStream;
import java.io.IOException;
import java.io.OutputStream;

import java.util.HashMap;
import java.util.Map;
import java.util.Vector;

/**
 * <p>A class to parse a JPEG bitstream.</p>
 *
 * <p>Includes the ability to selectively redact blocks and leave other blocks alone, to permit "lossless" redaction.</p>
 *
 * <p>Development of this class was supported by funding from MDDX Research and Informatics.</p>
 *
 * @author	dclunie
 */
public class Parse {

	private static final String identString = "@(#) $Header: /userland/cvs/codec/com/pixelmed/codec/jpeg/Parse.java,v 1.6 2014/03/29 21:58:58 dclunie Exp $";
	
	private static int getLargestSamplingFactor(int[] factors) {
		int largest = 0;
		for (int factor : factors) {
			if (factor > largest) {
				largest = factor;
			}
		}
		return largest;
	}
	
	private static final void writeMarkerAndLength(OutputStream out, int marker,int length) throws IOException {
		out.write(0xff);
		out.write(marker&0xff);
		out.write((length>>>8)&0xff);
		out.write(length&0xff);
	}
	
	private static final void writeVariableLengthMarkerSegment(OutputStream out, int marker,int length,byte[] b) throws IOException {
		writeMarkerAndLength(out,marker,length);
		out.write(b,0,length-2);
	}
	
	// follows pattern of dicom3tools appsrc/misc/jpegdump.cc
	
	/**
	 * <p>Parse a JPEG bitstream and copying to the output redacting any blocks that intersect with the specified locations.</p>
	 *
	 * @param	in				the input JPEG bitstream
	 * @param	out				the output JPEG bitsream, redacted as specified
	 * @param	redactionShapes	a Vector of Shape that are Rectangle
	 * @exception Exception			if bad things happen parsing the JPEG bit stream, caused by malformed input
	 * @exception IOException		if bad things happen reading or writing
	 */
	public static void parse(InputStream in,OutputStream out,Vector<Shape> redactionShapes) throws Exception, IOException {
		boolean dumping = false;
		boolean copying = out != null;
	
		EntropyCodedSegment ecs = null;					// lazy instantiation of EntropyCodedSegment ... wait until we have relevant marker segments for its constructor
		
		ByteArrayOutputStream byteAccumulator = null;	// recreated for first byte of each EntropyCodedSegment (at start and at each subsequent restart interval)
		
		MarkerSegmentSOS sos = null;
		MarkerSegmentSOF sof = null;
		Map<String,HuffmanTable> htByClassAndIdentifer = new HashMap<String,HuffmanTable>();
		Map<String,QuantizationTable> qtByIdentifer = new HashMap<String,QuantizationTable>();
		int restartinterval = 0;
		
		int mcuOffset = 0;
		int nMCUHorizontally = 0;
		int mcuCountPerEntropyCodedSegment = 0;

		int offset=0;
		int markerprefix = in.read();
		while (true) {
			if (markerprefix == -1) {
				if (dumping) System.err.print("End of file\n");
				break;
			}
			if (markerprefix != 0xff) {		// byte of entropy-coded segment
				if (byteAccumulator == null) {
					if (dumping) System.err.print("Offset "+Utilities.toPaddedHexString(offset,4)+" Starting new Entropy Coded Segment\n");
					byteAccumulator = new ByteArrayOutputStream();
				}
				byteAccumulator.write(markerprefix);
				++offset;
				markerprefix=in.read();
				continue;
			}
			int marker=in.read();
			if (marker == -1) {
				if (dumping) System.err.print("End of file immediately after marker flag 0xff ... presumably was padding\n");
				break;
			}
			else if (marker == 0xff) {		// 0xff byte of padding
				if (dumping) System.err.print("Offset "+Utilities.toPaddedHexString(offset,4)+" Fill byte 0xff\n");
				++offset;
				markerprefix=marker;		// the first 0xff is padding, the 2nd may be the start of a marker
				continue;
			}
			// ignore doing_jpeg2k_tilepart for now :(
			else if (marker == 0) {			// 0xff byte of entropy-coded segment ... ignore following zero byte
				if (dumping) System.err.print("Offset "+Utilities.toPaddedHexString(offset,4)+" Encoded 0xff in entropy-coded segment followed by stuffed zero byte\n");
				if (byteAccumulator == null) {
					if (dumping) System.err.print("Offset "+Utilities.toPaddedHexString(offset,4)+" Starting new Entropy Coded Segment\n");
					byteAccumulator = new ByteArrayOutputStream();
				}
				byteAccumulator.write(markerprefix);
				markerprefix=in.read();
				offset+=2;
				continue;
			}
			// ignore doing_jpegls and zero stuffed bit instead of byte for now :(

			// Definitely have a marker ...
			
			if (byteAccumulator != null) {
				// process any Entropy Coded Segment bytes accumulated so far ...
				if (ecs == null) {
					ecs = new EntropyCodedSegment(restartinterval,sos,sof,htByClassAndIdentifer,qtByIdentifer,copying,dumping);
					
					// need to figure out the sampling factors if this is the first Entropy Coded Segment, so that EntropyCodedSegment.finish() knows how many to process and where it is at
					
					int horizontalSamplesPerMCU = 8 * getLargestSamplingFactor(sof.getHorizontalSamplingFactor());
//System.err.println("horizontalSamplesPerMCU "+horizontalSamplesPerMCU);
					nMCUHorizontally = (sof.getNSamplesPerLine()-1)/horizontalSamplesPerMCU + 1;
//System.err.println("nMCUHorizontally "+nMCUHorizontally);
		
					int verticalSamplesPerMCU = 8 * getLargestSamplingFactor(sof.getVerticalSamplingFactor());
//System.err.println("verticalSamplesPerMCU "+verticalSamplesPerMCU);
					int nMCUVertically = (sof.getNLines()-1)/verticalSamplesPerMCU + 1;					// may need to update this from DNL marker :(
//System.err.println("nMCUVertically "+nMCUVertically);
		
					mcuCountPerEntropyCodedSegment = (restartinterval == 0) ? nMCUHorizontally * nMCUVertically : restartinterval;
					mcuOffset = 0;
				}
				byte[] bytesToDecompress = byteAccumulator.toByteArray();
//System.err.println("bytesToDecompress length="+bytesToDecompress.length);
				byte[] bytesToCopy = ecs.finish(bytesToDecompress,mcuCountPerEntropyCodedSegment,nMCUHorizontally,mcuOffset,redactionShapes);
				if (copying) {
					out.write(bytesToCopy);		// NB. EntropyCodedSegment.finish() has already done the zero byte stuffing after 0xff values
				}
				byteAccumulator = null;
				mcuOffset += mcuCountPerEntropyCodedSegment;
			}

			marker|=0xff00;			// convention is to express them with the leading ff, so that is what we look up
			
			if (dumping) System.err.print("Offset "+Utilities.toPaddedHexString(offset,4)+" Marker "+Utilities.toPaddedHexString(marker,4)+" "+Markers.getAbbreviation(marker)+" "+Markers.getDescription(marker)+" ");

			offset+=2;	// wait till after we have printed it to increment it
			
			if (Markers.isVariableLengthJPEGSegment(marker)) {
				int length=Utilities.read16be(in);
				if (length == -1) {
					throw new Exception("Error - variable length marker without length at Offset "+Utilities.toPaddedHexString(offset,4));
				}
				else {
					offset+=2;
					if (dumping) System.err.print("length variable "+Utilities.toPaddedHexString(length,2)+" ");
				}
				
				if (length > 2) {
					byte[] b = new byte[length-2];
					int count = in.read(b,0,length-2);
					if (count != length-2) {
						throw new Exception("Error - couldn't read variable length parameter sequence at Offset "+Utilities.toPaddedHexString(offset,4));
					}
					else {
						switch (marker) {
							case Markers.SOS:
								sos = new MarkerSegmentSOS(b,length-2);
								if (dumping) System.err.print(sos);
								if (copying) writeVariableLengthMarkerSegment(out,marker,length,b);
								break;
							case Markers.SOF0:
							case Markers.SOF1:
							case Markers.SOF2:
							case Markers.SOF3:
							case Markers.SOF5:
							case Markers.SOF6:
							case Markers.SOF7:
							case Markers.SOF9:
							case Markers.SOFA:
							case Markers.SOFB:
							case Markers.SOFD:
							case Markers.SOFE:
							case Markers.SOFF:
							case Markers.SOF55:
								sof = new MarkerSegmentSOF(b,length-2);
								if (dumping) System.err.print(sof);
								if (copying) writeVariableLengthMarkerSegment(out,marker,length,b);
								break;
							case Markers.DHT:
								MarkerSegmentDHT dht = new MarkerSegmentDHT(b,length-2);
								dht.addToMapByClassAndIdentifier(htByClassAndIdentifer);	// hokey, but sometimes multiple tables in one segment, sometimes multiple segments
								if (dumping) System.err.print(dht);
								if (copying) writeVariableLengthMarkerSegment(out,marker,length,b);
								break;
							case Markers.DQT:
								MarkerSegmentDQT dqt = new MarkerSegmentDQT(b,length-2);
								dqt.addToMapByIdentifier(qtByIdentifer);					// hokey, but sometimes multiple tables in one segment, sometimes multiple segments
								if (dumping) System.err.print(dqt);
								if (copying) writeVariableLengthMarkerSegment(out,marker,length,b);
								break;
							//case Markers.LSE
							//	break;
							case Markers.DRI:
								if (length == 4) {
									restartinterval = Utilities.extract16be(b,0);
								}
								else if (length == 5) {
									restartinterval = (int)Utilities.extract24be(b,0);
								}
								else if (length == 6) {
									restartinterval = (int)Utilities.extract32be(b,0);
								}
								else {
									throw new Exception("Illegal length "+length+" of restart interval at Offset "+Utilities.toPaddedHexString(offset,4));
								}
								if (dumping) System.err.print("\n\tDRI - Define Restart Interval = "+Utilities.toPaddedHexString(restartinterval,4)+"\n");
								if (copying) writeVariableLengthMarkerSegment(out,marker,length,b);
								break;
							case Markers.DNL:
								long numberoflines;
								if (length == 4) {
									numberoflines = Utilities.extract16be(b,0);
								}
								else if (length == 5) {
									numberoflines = Utilities.extract24be(b,0);
								}
								else if (length == 6) {
									numberoflines = Utilities.extract32be(b,0);
								}
								else {
									throw new Exception("Illegal length "+length+" of number of lines at Offset "+Utilities.toPaddedHexString(offset,4));
								}
								if (dumping) System.err.print("\n\tDNL - Define Number of Lines = "+Utilities.toPaddedHexString(numberoflines,4)+"\n");
								if (copying) writeVariableLengthMarkerSegment(out,marker,length,b);
								break;
							//case Markers.COD:
							//	break;
							//case Markers.COM:
							//	// do NOT copy COM marker segments ... may leak identity
							//	break;
							case Markers.APP0:
							case Markers.APP1:
							case Markers.APP2:
								String magic = "";
								{
									StringBuffer magicbuf = new StringBuffer();
									for (int i=0; i<b.length && b[i] != 0; ++i) {
										magicbuf.append(Character.valueOf((char)b[i]));
									}
									magic = magicbuf.toString();
								}
								if (dumping) System.err.print(magic);
								if (marker == Markers.APP0 && magic.equals("JFIF")) {
									if (dumping) System.err.print(new MarkerSegmentAPP0JFIF(b,length-2));
									//if (copying) writeVariableLengthMarkerSegment(out,marker,length,b);
								}
								// may want to consider not copying unrecognized APPn segments ... may leak identity ... copy everything for now :(
								if (copying) writeVariableLengthMarkerSegment(out,marker,length,b);
								break;
							default:
								// may want to consider not copying unrecognized segments ... may leak identity ... copy everything for now :(
								if (copying) writeVariableLengthMarkerSegment(out,marker,length,b);
								break;
						}
					}
				}
				else {
					if (dumping) System.err.print("Warning - variable length marker without \"zero\" length (really 2)");
				}
				offset+=(length-2);
			}
			else if (Markers.isNoLengthJPEGSegment(marker)) {
				if (copying) { out.write(0xff); out.write(marker&0xff);}
				if (marker == Markers.EOI) {
					// stop rather than process padding to end of file, so as not to create spurious empty EntropyCodedSegment
					if (dumping) System.err.print("\n");
					break;
				}
			}
			else {
				int length=Markers.isFixedLengthJPEGSegment(marker);
				switch (length) {
					case 0:
						break;
					case 3:
						{
							int value = in.read();
							if (value != -1) {
								offset+=1;
								if (dumping) System.err.print("length fixed 3 value "+Utilities.toPaddedHexString(value,2)+" ");
								if (copying) { writeMarkerAndLength(out,marker,length); out.write(value&0xff); }
							}
							else {
								throw new Exception("Error - fixed length 3 marker without value at Offset "+Utilities.toPaddedHexString(offset,4));
							}
						}
						break;
					case 4:
						{
							int value=Utilities.read16be(in);
							if (value != -1) {
								offset+=2;
								if (dumping) System.err.print("length fixed 4 value "+Utilities.toPaddedHexString(value,2)+" ");
								if (copying) { writeMarkerAndLength(out,marker,length); out.write((value>>>8)&0xff); out.write(value&0xff); }
							}
							else {
								throw new Exception("Error - fixed length 4 marker without value at Offset "+Utilities.toPaddedHexString(offset,4));
							}
						}
						break;
					default:
						throw new Exception("Error - fixed length marker with unexpected length "+length+" at Offset "+Utilities.toPaddedHexString(offset,4));
						//break;
				}
			}

			if (dumping) System.err.print("\n");
			markerprefix=in.read();
		}
	}
	
	/**
	 * <p>Test utility to read and write a JPEG file to check parsing is sound.</p>
	 *
	 * @param	arg	two parameters, the input file and the output file
	 */
	public static void main(String arg[]) {
		try {
			InputStream in = new FileInputStream(arg[0]);
			OutputStream out = arg.length > 1 ? new FileOutputStream(arg[1]) : null;
			long startTime = System.currentTimeMillis();
			parse(in,out,null);
			long currentTime = System.currentTimeMillis();
			long runTime = currentTime-startTime;
System.err.println("Took = "+runTime+" ms");
		}
		catch (Exception e) {
			e.printStackTrace(System.err);
		}
	}

}