/*******************************************************************************
 * Copyright (c) 2007, 2012 Association for Decentralized Information Management
 * in Industry THTH ry.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 *
 * Contributors:
 *     VTT Technical Research Centre of Finland - initial API and implementation
 *******************************************************************************/
package org.simantics.compressions.impl;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.IntBuffer;
import java.nio.channels.WritableByteChannel;

/**
 * @author Hannu Niemistö
 * @author Tuukka Lehtonen
 */
public abstract class CompressingOutputStream extends OutputStream {

    protected static final int    BLOCK_SIZE = 1024 * 1024;

    protected OutputStream        stream;
    protected WritableByteChannel channel;
    protected ByteBuffer          compressed;
    protected byte[]              compressedArray;
    protected ByteBuffer          uncompressed;
    protected IntBuffer           header;
    protected int                 bytes      = 0;

    public CompressingOutputStream(File file) throws FileNotFoundException {
        this(new FileOutputStream(file));
    }

    public CompressingOutputStream(FileOutputStream stream) {
        this(stream, stream.getChannel());
    }

    public CompressingOutputStream(OutputStream stream) {
        this(stream, null);
    }

    public CompressingOutputStream(OutputStream stream, WritableByteChannel channel) {
        this.stream = stream;
        this.channel = channel;
        uncompressed = allocateBuffer(BLOCK_SIZE);
    }

    @Override
    public void write(int b) throws IOException {
        if(uncompressed.remaining() == 0)
            flushBuffer();
        uncompressed.put((byte)b);
        ++bytes;
    }

    @Override
    public void write(byte[] b, int off, int len) throws IOException {
        bytes += len;
        while(len > 0) {
            int s = Math.min(len, uncompressed.remaining());
            uncompressed.put(b, off, s);
            len -= s;
            off += s;
            if(uncompressed.remaining() == 0)
                flushBuffer();
        }
    }

    @Override
    public void close() throws IOException {
        try {
            flushBuffer();

            if (channel != null) {
                header.position(0);
                header.put(0);
                header.put(0);
                compressed.position(0);
                compressed.limit(8);
                channel.write(compressed);
            } else if (stream != null) {
                stream.write(new byte[8]);
            }
            //System.out.println("Wrote " + bytes + " uncompressed bytes.");
        } finally {
            if (channel != null) {
                channel.close();
                channel = null;
            }
            if (stream != null) {
                stream.close();
                stream = null;
            }
        }
    }

    private void flushBuffer() throws IOException {
        int uncompressedSize = uncompressed.position();
        uncompressed.position(0);
        int maxCompressedSize = compressBound(uncompressedSize) + 8;
        if(compressed == null || compressed.capacity() < maxCompressedSize) {
            compressed = allocateBuffer(maxCompressedSize);
            compressed.order(ByteOrder.LITTLE_ENDIAN);
            header = compressed.asIntBuffer();
            if (channel == null)
                compressedArray = new byte[maxCompressedSize];
        }
        int compressedSize = compress(uncompressed, 0, uncompressedSize, 
                compressed, 8);
        header.position(0);
        header.put(compressedSize);
        header.put(uncompressedSize);
        compressed.position(0);
        compressed.limit(compressedSize+8);
        if (channel != null) {
            channel.write(compressed);
        } else {
            compressed.get(compressedArray, 0, compressedSize+8);
            stream.write(compressedArray, 0, compressedSize+8);
        }
    }

    protected abstract int compressBound(int inputSize);

    protected abstract int compress(ByteBuffer uncompressed, int uncompressedOffset, int uncompressedSize,
            ByteBuffer compressed, int compressedOffset) throws IOException;

    protected abstract ByteBuffer allocateBuffer(int capacity);

}