package aalg.crypto;

import java.io.IOException;
import java.io.OutputStream;

/**
 * Base implementation for byte padding writer.
 * Writer normally writes all blocks till the last one that is encoded
 * by method writePaddingToStream.
 * @author Ing. Petr Ležák
 */
abstract public class BytePaddingWriter implements IPaddingWriter {

	protected class UnpaddedStream extends OutputStream {
		@Override
		public void write (final byte[] b) throws IOException {
			paddedStream.write(b);
			updateLength (b.length);
		}

		@Override
		public void write (final byte[] b, final int offset, final int length) throws IOException {
			paddedStream.write(b, offset, length);
			updateLength (length);
		}

		@Override
		public void write(final int b) throws IOException {
			paddedStream.write(b);
			updateLength (1);
		}
		
		@Override 
		public void flush() throws IOException {
			paddedStream.flush();
		}

		@Override
		public void close() throws IOException {
			writePadding();
			
			super.close();
		}
	}
	
	protected OutputStream paddedStream;
	protected UnpaddedStream unpaddedStream;
	protected int blockLength;
	protected int filledBytes;
	private boolean closed;
	
	
	private void updateLength (final int writtenBytes) {
		filledBytes = (int) (((long) writtenBytes + (long) filledBytes) % blockLength);
	}
	
	@Override
	public void initialize(final OutputStream paddedStream, final int blockLength) throws IOException {
		this.paddedStream = paddedStream;
		this.blockLength = blockLength;
		
		if (blockLength <= 0 || blockLength > Pkcs7Padding.MAX_BLOCK_LENGTH)
			throw new RuntimeException("Wrong length of block");
		
		this.unpaddedStream = new UnpaddedStream();
		this.closed = false;
	}

	@Override
	public OutputStream getSourceStream() throws IOException {
		return unpaddedStream;
	}

	@Override
	public OutputStream getTargetStream() throws IOException {
		return paddedStream;
	}

	@Override
	public void finish() throws IOException {
		unpaddedStream.close();
	}
	
	private void writePadding() throws IOException {
		if (!closed) {
			writePaddingToStream();
			closed = true;
		}
	}
	
	protected abstract void writePaddingToStream() throws IOException;
}
