package ar.com.hjg.pngj;

import java.util.zip.DataFormatException;
import java.util.zip.Inflater;

/**
 * A set of IDAT-like chunks which, concatenated, form a zlib stream.
 * <p>
 * The inflated stream is intented to be read as a sequence of "rows", of which
 * the caller knows the lengths (not
 * necessary equal) and number.
 * <p>
 * Eg: For IDAT non-interlaced images, a row has bytesPerRow + 1 filter byte<br>
 * For interlaced images, the lengths are variable.
 * <p>
 * This class can work in sync (polled) mode or async (callback) mode. But for
 * callback mode the method
 * processRowCallback() must be overriden
 * <p>
 * See {@link IdatSet}, which is mostly used and has a slightly simpler use.<br>
 * See <code>DeflatedChunkSetTest</code> for example of use.
 */
public class DeflatedChunksSet {

	protected byte[] row; // a "row" here means a raw (uncopressed filtered) part of the IDAT stream,
							// normally a image row (or subimage row for interlaced) plus a filter byte
	private int rowfilled; // effective/valid length of row
	private int rowlen; // what amount of bytes is to be interpreted as a complete "row". can change
						// (for interlaced)
	private int rown; // only coincide with image row if non-interlaced - incremented by
						// setNextRowSize()

	/*
	 * States WAITING_FOR_INPUT ROW_READY WORK_DONE TERMINATED
	 * 
	 * processBytes() is externally called, prohibited in READY (in DONE it's ignored)
	 * 
	 * WARNING: inflater.finished() != DONE (not enough, not neccesary) DONE means that we have already uncompressed all
	 * the data of interest.
	 * 
	 * In non-callback mode, prepareForNextRow() is also externally called, in
	 * 
	 * Flow: - processBytes() calls inflateData() - inflateData() : if buffer is filled goes to READY else if !
	 * inf.finished goes to WAITING else if any data goes to READY (incomplete data to be read) else goes to DONE - in
	 * Callback mode, after going to READY, n=processCallback() is called and then prepareForNextRow(n) is called. - in
	 * Polled mode, prepareForNextRow(n) must be called from outside (after checking state=READY) - prepareForNextRow(n)
	 * goes to DONE if n==0 calls inflateData() again - end() goes to DONE
	 */
	private enum State {
		WAITING_FOR_INPUT, // waiting for more input
		ROW_READY, // ready for consumption (might be less than fully filled), ephemeral for CALLBACK
		// mode
		WORK_DONE, // all data of interest has been read, but we might accept still more trailing chunks
		// (we'll ignore them)
		TERMINATED; // we are done, and also won't accept more IDAT chunks

		public boolean isDone() {
			return this == WORK_DONE || this == TERMINATED;
		} // the caller has already uncompressed all the data of interest or EOF

		public boolean isTerminated() {
			return this == TERMINATED;
		} // we dont accept more chunks
	}

	State state = State.WAITING_FOR_INPUT; // never null

	private Inflater inf;
	private final boolean infOwn; // true if we own the inflater (we created it)

	private DeflatedChunkReader curChunk;

	private boolean callbackMode = true;
	private long nBytesIn = 0; // count the total compressed bytes that have been fed
	private long nBytesOut = 0; // count the total uncompressed bytes
	int chunkNum = -1; // incremented at each new chunk start
	int firstChunqSeqNum = -1; // expected seq num for first chunk. used only for fDAT (APNG)

	/**
	 * All IDAT-like chunks that form a same DeflatedChunksSet should have the
	 * same id
	 */
	public final String chunkid;

	/**
	 * @param initialRowLen
	 *            Length in bytes of first "row" (see description)
	 * @param maxRowLen
	 *            Max length in bytes of "rows"
	 * @param inflater
	 *            Can be null. If not null, must be already reset (and it must
	 *            be closed/released by caller!)
	 */
	public DeflatedChunksSet(String chunkid, int initialRowLen, int maxRowLen, Inflater inflater, byte[] buffer) {
		this.chunkid = chunkid;
		this.rowlen = initialRowLen;
		if (initialRowLen < 1 || maxRowLen < initialRowLen)
			throw new PngjException("bad inital row len " + initialRowLen);
		if (inflater != null) {
			this.inf = inflater;
			infOwn = false;
		} else {
			this.inf = new Inflater();
			infOwn = true; // inflater is own, we will release on close()
		}
		this.row = buffer != null && buffer.length >= initialRowLen ? buffer : new byte[maxRowLen];
		rown = -1;
		this.state = State.WAITING_FOR_INPUT;
		try {
			prepareForNextRow(initialRowLen);
		} catch (RuntimeException e) {
			close();
			throw e;
		}
	}

	public DeflatedChunksSet(String chunkid, int initialRowLen, int maxRowLen) {
		this(chunkid, initialRowLen, maxRowLen, null, null);
	}

	protected void appendNewChunk(DeflatedChunkReader cr) {
		// all chunks must have same id
		if (!this.chunkid.equals(cr.getChunkRaw().id))
			throw new PngjInputException("Bad chunk inside IdatSet, id:" + cr.getChunkRaw().id + ", expected:" + this.chunkid);
		this.curChunk = cr;
		chunkNum++;
		if (firstChunqSeqNum >= 0)
			cr.setSeqNumExpected(chunkNum + firstChunqSeqNum);
	}

	/**
	 * Feeds the inflater with the compressed bytes
	 * 
	 * In poll mode, the caller should not call repeatedly this, without
	 * consuming first, checking
	 * isDataReadyForConsumer()
	 * 
	 * @param buf
	 * @param off
	 * @param len
	 */
	protected void processBytes(byte[] buf, int off, int len) {
		nBytesIn += len;
		// PngHelperInternal.LOGGER.info("processing compressed bytes in chunkreader : " + len);
		if (len < 1 || state.isDone())
			return;
		if (state == State.ROW_READY)
			throw new PngjInputException("this should only be called if waitingForMoreInput");
		if (inf.needsDictionary() || !inf.needsInput())
			throw new RuntimeException("should not happen");
		inf.setInput(buf, off, len);
		// PngHelperInternal.debug("entering processs bytes, state=" + state +
		// " callback="+callbackMode);
		if (isCallbackMode()) {
			while (inflateData()) {
				int nextRowLen = processRowCallback();
				prepareForNextRow(nextRowLen);
				if (isDone())
					processDoneCallback();
			}
		} else
			inflateData();
	}

	/*
	 * This never inflates more than one row This returns true if this has resulted in a row being ready and preprocessed
	 * with preProcessRow (in callback mode, we should call immediately processRowCallback() and
	 * prepareForNextRow(nextRowLen)
	 */
	private boolean inflateData() {
		try {
			// PngHelperInternal.debug("entering inflateData bytes, state=" + state +
			// " callback="+callbackMode);
			if (state == State.ROW_READY)
				throw new PngjException("invalid state");// assert
			if (state.isDone())
				return false;
			int ninflated = 0;
			if (row == null || row.length < rowlen)
				row = new byte[rowlen]; // should not happen
			if (rowfilled < rowlen && !inf.finished()) {
				try {
					ninflated = inf.inflate(row, rowfilled, rowlen - rowfilled);
				} catch (DataFormatException e) {
					throw new PngjInputException("error decompressing zlib stream ", e);
				}
				rowfilled += ninflated;
				nBytesOut += ninflated;
			}
			State nextstate = null;
			if (rowfilled == rowlen)
				nextstate = State.ROW_READY; // complete row, process it
			else if (!inf.finished())
				nextstate = State.WAITING_FOR_INPUT;
			else if (rowfilled > 0)
				nextstate = State.ROW_READY; // complete row, process it
			else {
				nextstate = State.WORK_DONE; // eof, no more data
			}
			state = nextstate;
			if (state == State.ROW_READY) {
				preProcessRow();
				return true;
			}
		} catch (RuntimeException e) {
			close();
			throw e;
		}
		return false;
	}

	/**
	 * Called automatically in all modes when a full row has been inflated.
	 */
	protected void preProcessRow() {

	}

	/**
	 * Callback, must be implemented in callbackMode
	 * <p>
	 * This should use {@link #getRowFilled()} and {@link #getInflatedRow()} to
	 * access the row.
	 * <p>
	 * Must return byes of next row, for next callback.
	 */
	protected int processRowCallback() {
		throw new PngjInputException("not implemented");
	}

	/**
	 * Callback, to be implemented in callbackMode
	 * <p>
	 * This will be called once to notify state done
	 */
	protected void processDoneCallback() {}

	/**
	 * Inflated buffer.
	 * 
	 * The effective length is given by {@link #getRowFilled()}
	 */
	public byte[] getInflatedRow() {
		return row;
	}

	/**
	 * Should be called after the previous row was processed
	 * <p>
	 * Pass 0 or negative to signal that we are done (not expecting more bytes)
	 * <p>
	 * This resets {@link #rowfilled}
	 * <p>
	 * The
	 */
	public void prepareForNextRow(int len) {
		rowfilled = 0;
		rown++;
		if (len < 1) {
			rowlen = 0;
			done();
		} else if (inf.finished()) {
			rowlen = 0;
			done();
		} else {
			state = State.WAITING_FOR_INPUT;
			rowlen = len;
			if (!callbackMode)
				inflateData();
		}
	}

	/**
	 * In this state, the object is waiting for more input to deflate.
	 * <p>
	 * Only in this state it's legal to feed this
	 */
	public boolean isWaitingForMoreInput() {
		return state == State.WAITING_FOR_INPUT;
	}

	/**
	 * In this state, the object is waiting the caller to retrieve inflated data
	 * <p>
	 * Effective length: see {@link #getRowFilled()}
	 */
	public boolean isRowReady() {
		return state == State.ROW_READY;
	}

	/**
	 * In this state, all relevant data has been uncompressed and retrieved
	 * (exceptionally, the reading has ended
	 * prematurely).
	 * <p>
	 * We can still feed this object, but the bytes will be swallowed/ignored.
	 */
	public boolean isDone() {
		return state.isDone();
	}

	public boolean isTerminated() {
		return state.isTerminated();
	}

	/**
	 * This will be called by the owner to report us the next chunk to come. We
	 * can make our own internal changes and
	 * checks. This returns true if we acknowledge the next chunk as part of
	 * this set
	 */
	public boolean ackNextChunkId(String id) {
		if (state.isTerminated())
			return false;
		else if (id.equals(chunkid)) {
			return true;
		} else {
			if (!allowOtherChunksInBetween(id)) {
				if (state.isDone()) {
					if (!isTerminated())
						terminate();
					return false;
				} else {
					throw new PngjInputException("Unexpected chunk " + id + " while " + chunkid + " set is not done");
				}
			} else
				return true;
		}
	}

	protected void terminate() {
		close();
	}

	/**
	 * This should be called when discarding this object, or for aborting.
	 * Secure, idempotent Don't use this just to
	 * notify this object that it has no more work to do, see {@link #done()}
	 */
	public void close() {
		try {
			if (!state.isTerminated()) {
				state = State.TERMINATED;
			}
			if (infOwn && inf != null) {
				inf.end();// we end the Inflater only if we created it
				inf = null;
			}
		} catch (Exception e) {}
	}

	/**
	 * Forces the DONE state, this object won't uncompress more data. It's still
	 * not terminated, it will accept more IDAT
	 * chunks, but will ignore them.
	 */
	public void done() {
		if (!isDone())
			state = State.WORK_DONE;
	}

	/**
	 * Target size of the current row, including filter byte. <br>
	 * should coincide (or be less than) with row.length
	 */
	public int getRowLen() {
		return rowlen;
	}

	/** This the amount of valid bytes in the buffer */
	public int getRowFilled() {
		return rowfilled;
	}

	/**
	 * Get current (last) row number.
	 * <p>
	 * This corresponds to the raw numeration of rows as seen by the deflater.
	 * Not the same as the real image row, if
	 * interlaced.
	 * 
	 */
	public int getRown() {
		return rown;
	}

	/**
	 * Some IDAT-like set can allow other chunks in between (APGN?).
	 * <p>
	 * Normally false.
	 * 
	 * @param id
	 *            Id of the other chunk that appeared in middel of this set.
	 * @return true if allowed
	 */
	public boolean allowOtherChunksInBetween(String id) {
		return false;
	}

	/**
	 * Callback mode = async processing
	 */
	public boolean isCallbackMode() {
		return callbackMode;
	}

	public void setCallbackMode(boolean callbackMode) {
		this.callbackMode = callbackMode;
	}

	/** total number of bytes that have been fed to this object */
	public long getBytesIn() {
		return nBytesIn;
	}

	/** total number of bytes that have been uncompressed */
	public long getBytesOut() {
		return nBytesOut;
	}

	@Override
	public String toString() {
		StringBuilder sb = new StringBuilder("idatSet : " + curChunk.getChunkRaw().id + " state=" + state + " rows=" + rown + " bytes=" + nBytesIn + "/" + nBytesOut);
		return sb.toString();
	}

}