package org.forkalsrud.album.video;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.Arrays;

/**
 * Separates the FLV header boxes from the body boxes.
 * After streaming a file through the header can be rewritten with metadata appropriate for streaming.
 */
public class FlvFilter extends OutputStream {

    private static org.slf4j.Logger log = org.slf4j.LoggerFactory.getLogger(FlvFilter.class);

    public interface FlvReceiver {
        public void writeHeader(byte[] data);
        public void writeBody(byte[] data);
    }

    private final static int FLV_SIZE_TAGHEADER = 11;
    private final static int FLV_SIZE_TAGFOOTER = 4;
    private final static int FLV_SIZE_FILEHEADER = 13;

    private final static int FLV_TAG_AUDIO = 8;
    private final static int FLV_TAG_VIDEO = 9;
    private final static int FLV_TAG_SCRIPTDATA = 18;

    private boolean headerWritten = false;
    private FlvReceiver receiver;

    private byte[] fileHeader = new byte[FLV_SIZE_FILEHEADER];
    private int byteCounter = 0;
    private int incomingMetadataLength = 0;

    private byte[] currentBoxHeader = new byte[FLV_SIZE_TAGHEADER];
    private int currentBoxWriteIdx = 0;

    private int currentTagPos;
    private int currentTagType;
    private int currentDataSize;
    private int currentTagTimestamp;

    private int currentTagSize;
    private byte[] currentBox;

    private FlvMetadata metadata;
    private FlvMetadata extraMetadata;
    private int bodyLen = 0;

    public FlvFilter(FlvReceiver receiver, FlvMetadata metadata) {
        this.receiver = receiver;
        this.metadata = new FlvMetadata();
        this.extraMetadata = metadata;
    }

    /**
     * Receive some bytes of FLV output
     * @param b
     * @param off
     * @param len
     * @throws IOException
     */
    @Override
    public void write(byte[] b, int off, int len) throws IOException {

        // stateful
        // initially we receive the header 'FLV.....'
        // then we receive one box after another
        // if a box can be classified as a header we write it to header dst
        // else it is for body dst
        int remainingInputLength = len;
        int readOffset = off;
        if (byteCounter < FLV_SIZE_FILEHEADER) {
            int headerBytesToRead = Math.min(FLV_SIZE_FILEHEADER - byteCounter, remainingInputLength);
            for (int i = 0; i < headerBytesToRead; i++) {
                fileHeader[byteCounter++] = b[readOffset++];
            }
            remainingInputLength -= headerBytesToRead;
            currentTagPos = byteCounter;
        }
        while (remainingInputLength > 0) {
            // read the header first
            if (currentBoxWriteIdx < FLV_SIZE_TAGHEADER) {
                int headerBytesToRead = Math.min(FLV_SIZE_TAGHEADER - currentBoxWriteIdx, remainingInputLength);
                for (int i = 0; i < headerBytesToRead; i++) {
                    currentBoxHeader[currentBoxWriteIdx++] = b[readOffset++];
                    byteCounter++;
                }
                remainingInputLength -= headerBytesToRead;
                if (currentBoxWriteIdx < FLV_SIZE_TAGHEADER) {
                    // Don't have a full header yet, wait for next
                    if (remainingInputLength > 0) {
                        throw new RuntimeException("something is off in the lengths we're trying to read");
                    }
                    return;
                }
                // Got the header, prepare the buffer for the entire box
                currentTagType = decodeUint8(currentBoxHeader, 0);
                currentDataSize = decodeUint24(currentBoxHeader, 1);
                currentTagTimestamp = decodeTimestamp(currentBoxHeader, 4);
                currentTagSize = currentDataSize + FLV_SIZE_TAGHEADER + FLV_SIZE_TAGFOOTER;
                if (currentTagType == 0) {
                    log.info("Encountered tag type zero, exiting.");
                    currentBox = null;
                    return;
                }
                currentBox = new byte[currentTagSize];
                for (int i = 0; i < FLV_SIZE_TAGHEADER; i++) {
                    currentBox[i] = currentBoxHeader[i];
                }
            }
            int bodyBytesToRead = Math.min(currentTagSize - currentBoxWriteIdx, remainingInputLength);
            for (int i = 0; i < bodyBytesToRead; i++) {
                currentBox[currentBoxWriteIdx++] = b[readOffset++];
                byteCounter++;
            }
            remainingInputLength -= bodyBytesToRead;
            if (currentBoxWriteIdx < currentTagSize) {
                if (remainingInputLength > 0) {
                    throw new RuntimeException("something is off in the lengths we're trying to read");
                }
                return;
            }
            processBox();
            currentTagType = 0;
            currentDataSize = 0;
            currentTagTimestamp = 0;
            currentTagSize = 0;
            currentBox = null;
            currentBoxWriteIdx = 0;
            Arrays.fill(currentBoxHeader, (byte)0);
            currentTagPos = byteCounter;
        }
    }


    void processBox() throws IOException {

        // We expect one SCRIPTDATA box immediately following the file header
        // Then only audio and video
        if (currentTagType == FLV_TAG_SCRIPTDATA) {
            if (!headerWritten) {
                incomingMetadataLength = currentTagSize - (FLV_SIZE_TAGHEADER + FLV_SIZE_TAGFOOTER);
                ByteArrayInputStream scriptData = new ByteArrayInputStream(currentBox, FLV_SIZE_TAGHEADER, incomingMetadataLength);
                metadata.read(scriptData);
                metadata.merge(extraMetadata);
                generateHeaderChunk();
//                receiver.writeHeader(fileHeader, currentBox);  // copy the input directly
                headerWritten = true;
                metadata.clearKeyframes();
            } else {
                log.warn("SCRIPTDATA out of order");
                writeBody();
            }
        } else if (currentTagType == FLV_TAG_VIDEO || currentTagType == FLV_TAG_AUDIO) {
            if (!headerWritten) {
                log.error("Missing SCRIPTDATA");
                receiver.writeHeader(fileHeader);
                headerWritten = true;
            }
            if (currentTagType == FLV_TAG_VIDEO) {
                int flags = decodeUint8(currentBox, FLV_SIZE_TAGHEADER);
                boolean isKeyFrame = ((flags >> 4) & 0xf) == 1;
                metadata.addVideoFrame(currentTagPos, currentTagTimestamp, isKeyFrame);
            } else { // currentTagType == FLV_TAG_AUDIO
                metadata.addAudioFrame(currentTagPos, currentTagTimestamp);
            }
            writeBody();
        } else {
            log.error("Unknown box type: " + currentTagType);
            writeBody();
        }        
    }

    /**
     * 
     */
    private void writeBody() {
        receiver.writeBody(currentBox);
        bodyLen += currentBox.length;
    }


    @Override
    public void flush() throws IOException {
        if (currentBox != null) {
            throw new IOException("Premature end of file");
        }
        metadata.setDurationFromLastTimestamp();
    }

    /**
     * Writes out the header chunk which consists of a file header and a script data tag.
     * The tag body comes from the metadata object.  The file header is a copy of what we read in at the start.
     * @throws IOException
     */
    public void generateHeaderChunk() throws IOException {

        metadata.setFileSize(fileHeader.length + FLV_SIZE_TAGHEADER + incomingMetadataLength + FLV_SIZE_TAGFOOTER + bodyLen);
        int metadataLength = metadata.calculateLength();
        ByteArrayOutputStream out = new ByteArrayOutputStream(fileHeader.length + FLV_SIZE_TAGHEADER + metadataLength + FLV_SIZE_TAGFOOTER);
        // file header
        out.write(fileHeader);
        metadata.setFileSize(fileHeader.length + FLV_SIZE_TAGHEADER + metadataLength + FLV_SIZE_TAGFOOTER + bodyLen);
        metadata.setFileOffsetDelta(metadataLength - incomingMetadataLength);
        // tag header
        out.write(FLV_TAG_SCRIPTDATA); // Script data tag        1 byte
        metadata.writeUint24(out, metadataLength);  //data size  3 bytes
        metadata.writeTimestamp(out, 0); // timestamp            4 bytes
        metadata.writeUint24(out, 0); // Stream ID               3 bytes
        // body
        metadata.writeScriptDataObject(out);
        // tag footer
        metadata.writeUint32(out, FLV_SIZE_TAGHEADER + metadataLength + FLV_SIZE_TAGFOOTER); // size of tag
        receiver.writeHeader(out.toByteArray());
    }



    private int decodeUint24(byte[] buf, int offset) {
        int ch1 = buf[offset + 0] & 0xff;
        int ch2 = buf[offset + 1] & 0xff;
        int ch3 = buf[offset + 2] & 0xff;
        return ((ch1 << 16) + (ch2 << 8) + (ch3 << 0));
    }
    

    private int decodeUint8(byte[] buf, int offset) {
        return buf[offset] & 0xff;
    }

    private int decodeTimestamp(byte[] buf, int offset) {
        int ch1 = buf[offset + 0] & 0xff;
        int ch2 = buf[offset + 1] & 0xff;
        int ch3 = buf[offset + 2] & 0xff;
        int ch4 = buf[offset + 3] & 0xff;
        return ((ch4 << 24) + (ch1 << 16) + (ch2 << 8) + (ch3 << 0));
    }



    @Override
    public void write(int b) throws IOException {
        write(new byte[] { (byte)b });
    }


}
