package com.conveyal.osmlib;

import com.google.common.primitives.Ints;
import com.google.protobuf.ByteString;
import com.google.protobuf.GeneratedMessageLite;
import com.google.protobuf.GeneratedMessageV3;
import org.openstreetmap.osmosis.osmbinary.Fileformat;
import org.openstreetmap.osmosis.osmbinary.Osmformat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.OutputStream;
import java.util.concurrent.SynchronousQueue;
import java.util.zip.Deflater;

/**
 * Consumes OSM entity objects and writes a stream of PBF data blocks to the specified output stream.
 * This is neither threadsafe nor reentrant! Create one instance of this encoder per encode operation.
 */
public class PBFOutput implements OSMEntitySink, Runnable {

    private static final Logger LOG = LoggerFactory.getLogger(PBFOutput.class);

    /** The underlying output stream where VEX data will be written. */
    private OutputStream downstream;

    /** The replication timestamp to record in the PBF file. Should be set before writing begins. */
    private long timestamp;

    /** Values retained from one entity to the next within a block for delta decoding. */
    private long prevId, prevFixedLat, prevFixedLon;

    /** The entity type of the current block to enforce grouping of entities by type. null means no type yet. */
    private OSMEntity.Type currEntityType = null;

    /* Per-block state */

    private final StringTable stringTable = new StringTable();

    private int nEntitiesInBlock;

    private Osmformat.PrimitiveGroup.Builder primitiveGroupBuilder;

    private Osmformat.DenseNodes.Builder denseNodesBuilder;

    private Thread writerThread = null;

    /** Construct a new PBF output encoder which writes to the given downstream OutputStream. */
    public PBFOutput(OutputStream downstream) {
        this.downstream = downstream;
    }

    /** Reset the inter-entity delta coding values and set up a new block. */
    private void beginBlock(OSMEntity.Type eType) throws IOException {
        prevId = prevFixedLat = prevFixedLon = nEntitiesInBlock = 0;
        stringTable.clear();
        primitiveGroupBuilder = Osmformat.PrimitiveGroup.newBuilder();
        if (eType == OSMEntity.Type.NODE) {
            denseNodesBuilder = Osmformat.DenseNodes.newBuilder();
        }
    }

    /** We always add one primitive group of less that 8k elements to each primitive block. */
    private void endBlock () {
        if (nEntitiesInBlock > 0) {
            if (currEntityType == OSMEntity.Type.NODE) {
                primitiveGroupBuilder.setDense(denseNodesBuilder);
            }
            // Pass the block off to the compression/writing thread
            try {
                Osmformat.PrimitiveBlock primitiveBlock = Osmformat.PrimitiveBlock.newBuilder()
                        .setStringtable(stringTable.toBuilder()).addPrimitivegroup(primitiveGroupBuilder).build();
                synchronousQueue.put(primitiveBlock);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
    }

    /** @param block is either a PrimitiveBlock or a HeaderBlock */
    private void writeOneBlob(GeneratedMessageV3 block) {

        // FIXME lotsa big copies going on here

        String blobTypeString;
        if (block instanceof Osmformat.HeaderBlock) {
            blobTypeString = "OSMHeader";
        } else if (block instanceof Osmformat.PrimitiveBlock) {
            blobTypeString = "OSMData";
        } else {
            throw new AssertionError("block must be either a header block or a primitive block.");
        }

        Fileformat.Blob.Builder blobBuilder = Fileformat.Blob.newBuilder();
        byte[] serializedBlock = block.toByteArray();
        byte[] deflatedBlock = new byte[serializedBlock.length];
        int deflatedSize = deflate(serializedBlock, deflatedBlock);
        if (deflatedSize < 0) {
            LOG.debug("Deflate did not reduce the size of a block. Saving it uncompressed.");
            blobBuilder.setRaw(ByteString.copyFrom(serializedBlock));
        } else {
            blobBuilder.setZlibData(ByteString.copyFrom(deflatedBlock, 0, deflatedSize));
            blobBuilder.setRawSize(serializedBlock.length);
        }
        byte[] serializedBlob = blobBuilder.build().toByteArray();

        Fileformat.BlobHeader blobHeader = Fileformat.BlobHeader.newBuilder()
                .setType(blobTypeString).setDatasize(serializedBlob.length).build();
        byte[] serializedBlobHeader = blobHeader.toByteArray();
        try {
            // "Returns a big-endian representation of value in a 4-element byte array"
            downstream.write(Ints.toByteArray(serializedBlobHeader.length));
            downstream.write(serializedBlobHeader);
            downstream.write(serializedBlob);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

    }

    /**
     * Called at the beginning of each node, way, or relation to enforce grouping of entities by type.
     * If the entity type has changed since the last entity (except at the beginning of the first block),
     * ends the block and starts a new one of the new type. Each block will contain entities of only a single type.
     */
    private void checkBlockTransition(OSMEntity.Type eType) throws IOException {
        if (currEntityType != eType || nEntitiesInBlock++ >= 8000) {
            if (currEntityType != null) {
                endBlock();
            }
            currEntityType = eType;
            beginBlock(eType);
        }
    }

    /**
     * Deflate the given input data buffer into the given output byte buffer.
     * Used in both PBF and VEX output.
     * @return the deflated size of the data, or -1 if deflate did not reduce the data size.
     */
    public static int deflate (byte[] input, byte[] output) {
        int pos = 0;
        // Do not compress an empty data block, it will spin forever trying to fill the zero-length output buffer.
        if (input.length > 0) {
            Deflater deflater = new Deflater(Deflater.DEFAULT_COMPRESSION, false); // include gzip header and checksum
            deflater.setInput(input, 0, input.length);
            deflater.finish(); // There will be no more input after this byte array.
            while (!deflater.finished()) {
                pos += deflater.deflate(output, pos, output.length - pos, Deflater.SYNC_FLUSH);
                if (pos >= input.length) {
                    return -1; // compressed output is bigger than buffer, store uncompressed
                }
            }
        }
        return pos;
    }

    /* OSM DATA SINK INTERFACE */

    @Override
    public void writeBegin() throws IOException {

        LOG.info("Writing PBF format...");

        // Write out a header block
        Osmformat.HeaderBlock.Builder builder = Osmformat.HeaderBlock.newBuilder();
        builder.addRequiredFeatures("DenseNodes").setWritingprogram("Vanilla Extract").build();
        if (timestamp > 0) {
            builder.setOsmosisReplicationTimestamp(timestamp);
        }
        writeOneBlob(builder.build());

        // Start another thread that will handle compression and writing in parallel.
        writerThread = new Thread(this);
        writerThread.start();

    }

    @Override
    public void setReplicationTimestamp(long secondsSinceEpoch) {
        this.timestamp = secondsSinceEpoch;
    }

    @Override
    public void writeEnd() throws IOException {
        // Finish any partially-completed block.
        endBlock();
        // Send a primitive block with no primitive group to the writer thread, signaling it to shut down and clean up.
        try {
            synchronousQueue.put(Osmformat.PrimitiveBlock.getDefaultInstance());
            writerThread.join();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        LOG.info("Finished writing PBF format.");
    }

    @Override
    public void writeNode(long id, Node node) throws IOException {

        checkBlockTransition(OSMEntity.Type.NODE);

        long idDelta = id - prevId;
        prevId = id;
        denseNodesBuilder.addId(idDelta);

        /* Fixed-precision latitude and longitude */
        long fixedLat = (long) (node.fixedLat); // TODO are longs necessary?
        long fixedLon = (long) (node.fixedLon);
        long fixedLatDelta = fixedLat - prevFixedLat;
        long fixedLonDelta = fixedLon - prevFixedLon;
        prevFixedLat = fixedLat;
        prevFixedLon = fixedLon;
        denseNodesBuilder.addLat(fixedLatDelta).addLon(fixedLonDelta);

        /* Tags for a whole dense node block are stored as: key, val, key, val, 0, key, val, key, val, 0 */
        if (node.tags != null) {
            for (OSMEntity.Tag tag : node.tags) {
                if (tag.value == null) tag.value = "";
                int keyCode = stringTable.getCode(tag.key);
                int valCode = stringTable.getCode(tag.value);
                denseNodesBuilder.addKeysVals(keyCode);
                denseNodesBuilder.addKeysVals(valCode);
            }
        }
        denseNodesBuilder.addKeysVals(0);

    }

    @Override
    public void writeWay(long id, Way way) throws IOException {

        checkBlockTransition(OSMEntity.Type.WAY);
        Osmformat.Way.Builder builder = Osmformat.Way.newBuilder().setId(id);

        /* Tags */
        if (way.tags != null) {
            for (OSMEntity.Tag tag : way.tags) {
                if (tag.value == null) tag.value = "";
                builder.addKeys(stringTable.getCode(tag.key));
                builder.addVals(stringTable.getCode(tag.value));
            }
        }

        /* Node References */
        long prevNodeRef = 0;
        for (long ref : way.nodes) {
            builder.addRefs(ref - prevNodeRef); // delta-coded node references
            prevNodeRef = ref;
        }

        /* TODO Should we trigger the build here or just call with the builder? */
        primitiveGroupBuilder.addWays(builder.build());

    }

    @Override
    public void writeRelation(long id, Relation relation) throws IOException {

        checkBlockTransition(OSMEntity.Type.RELATION);
        Osmformat.Relation.Builder builder = Osmformat.Relation.newBuilder().setId(id);

        /* Tags */
        if (relation.tags != null) {
            for (OSMEntity.Tag tag : relation.tags) {
                if (tag.value == null) tag.value = "";
                builder.addKeys(stringTable.getCode(tag.key));
                builder.addVals(stringTable.getCode(tag.value));
            }
        }

        /* Relation members */
        long lastMemberId = 0;
        for (Relation.Member member : relation.members) {
            builder.addMemids(member.id - lastMemberId); // delta-coded member references
            lastMemberId = member.id;
            builder.addRolesSid(stringTable.getCode(member.role));
            Osmformat.Relation.MemberType memberType;
            if (member.type == OSMEntity.Type.NODE) {
                memberType = Osmformat.Relation.MemberType.NODE;
            } else if (member.type == OSMEntity.Type.WAY) {
                memberType = Osmformat.Relation.MemberType.WAY;
            } else if (member.type == OSMEntity.Type.RELATION) {
                memberType = Osmformat.Relation.MemberType.RELATION;
            } else {
                throw new RuntimeException("Member type was not defined.");
            }
            builder.addTypes(memberType);
        }

        /* TODO Should we trigger the build here or just call with the builder? */
        primitiveGroupBuilder.addRelations(builder.build());

    }

    /** A zero-length BlockingQueue that hands tasks to the compression/writing thread. */
    private final SynchronousQueue<Osmformat.PrimitiveBlock> synchronousQueue = new SynchronousQueue<>();

    /** Runnable interface implementation that compresses and writes output blocks asynchronously. */
    @Override
    public void run() {
        while (true) {
            try {
                Osmformat.PrimitiveBlock block = synchronousQueue.take(); // block until work is available
                if (block.getPrimitivegroupCount() == 0) {
                    break; // a block with no primitive groups tells the writer thread to shut down.
                }
                writeOneBlob(block);
            } catch (InterruptedException ex) {
                LOG.error("Block writer thread was interrupted while waiting for work.");
                break;
            }
        }
        try {
            downstream.flush();
            downstream.close();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

}
