/**
 * Copyright (c) 2006 - 2008 Smaxe Ltd (www.smaxe.com).
 * All rights reserved.
 */
package com.smaxe.uv.fformat;

import com.smaxe.io.ByteArray;

import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.List;

/**
 * <code>FlvTool</code> - FLV tool.
 * 
 * @author Andrei Sochirca
 */
public final class FlvTool extends Object
{
    /**
     * Returns video frame type by <code>tag</code>.
     * 
     * @param tag
     * @return frame encoded in the <code>tag</code>
     */
    public static int getVideoFrame(final int tag)
    {
        return (tag >> 4) & 0x0F;
    }
    
    /**
     * Returns video codec by <code>tag</code>.
     * 
     * @param tag
     * @return codec encoded in the <code>tag</code>
     */
    public static int getVideoCodec(final int tag)
    {
        return tag & 0x0F;
    }
    
    /**
     * <code>FileInfo</code> - FLV file info.
     */
    public final static class FileInfo extends Object
    {
        /**
         * Number of audio frames.
         */
        public int audioFrames = 0;
        /**
         * Size of audio frames (in bytes).
         */
        public int audioFramesSize = 0;
        /**
         * Audio stream duration (in milliseconds).
         */
        public int audioStreamDuration = 0;
        
        /**
         * Number of video frames.
         */
        public int videoFrames = 0;
        /**
         * Size of video frames (in bytes).
         */
        public int videoFramesSize = 0;
        /**
         * Video stream duration (in milliseconds).
         */
        public int videoStreamDuration = 0;
        
        /**
         * Number of data frames.
         */
        public int dataFrames = 0;
        /**
         * Size of data frames (in bytes).
         */
        public int dataFramesSize = 0;
        
        /**
         * Constructor.
         */
        public FileInfo()
        {
        }
        
        @Override
        public String toString()
        {
            return "FileInfo [" + "audio (frames/size/duration): " + audioFrames + "/" + audioFramesSize + "/" + audioStreamDuration + "; " + 
                "video (frames/size/duration): " + videoFrames + "/" + videoFramesSize + "/" + videoStreamDuration +  "; " +
                "data (frames/size): " + dataFrames + "/" + dataFramesSize + "]";
        }
    }
    
    /**
     * <code>GetDataFramesTagProcessor</code>
     */
    private final static class GetDataFramesTagProcessor extends Flv.TagProcessorAdapter
    {
        // fields
        private List<Flv.ScriptData[]> frames = null;
        
        /**
         * Constructor.
         * 
         * @param frames data frames
         */
        public GetDataFramesTagProcessor(List<Flv.ScriptData[]> frames)
        {
            this.frames = frames;
        }
        
        @Override
        public void onTag(final int type, final int dataSize, final long timestamp, final InputStream is) throws IOException
        {
            switch (type)
            {
                case Flv.Tag.TYPE_DATA:
                {
                    frames.add(Flv.readScriptData(is, dataSize));
                    
                } break;
            }
        }
    }
    
    /**
     * <code>GetInfoTagProcessor</code>
     */
    private final static class GetInfoTagProcessor extends Flv.TagProcessorAdapter
    {
        // fields
        private FileInfo info = null;
        
        private long lastAudioTimestamp = -1;
        private long lastVideoTimestamp = -1;
        
        /**
         * Constructor.
         * 
         * @param info file info 
         */
        public GetInfoTagProcessor(FileInfo info)
        {
            this.info = info;
        }
        
        @Override
        public void onTag(final int type, final int dataSize, final long timestamp, InputStream is) throws IOException
        {
            switch (type)
            {
                case Flv.Tag.TYPE_AUDIO:
                {
                    if (lastAudioTimestamp < 0) lastAudioTimestamp = timestamp;
                    
                    info.audioFrames++;
                    info.audioFramesSize += dataSize;
                    info.audioStreamDuration += timestamp - lastAudioTimestamp;
                    
                    lastAudioTimestamp = timestamp;
                    
                } break;
                case Flv.Tag.TYPE_VIDEO:
                {
                    if (lastVideoTimestamp < 0) lastVideoTimestamp = timestamp;
                    
                    info.videoFrames++;
                    info.videoFramesSize += dataSize;
                    info.videoStreamDuration += timestamp - lastVideoTimestamp;
                    
                    lastVideoTimestamp = timestamp;
                    
                } break;
                case Flv.Tag.TYPE_DATA:
                {
                    info.dataFrames++;
                    info.dataFramesSize += dataSize;
                    
                } break;
            }
        }
    }
    
    /**
     * <code>ConvertFlvToMp3TagProcessor</code>
     */
    private final static class ConvertFlvToMp3TagProcessor extends Flv.TagProcessorAdapter
    {
        /**
         * <code>FLV_AUDIO_TAG_SIZE</code> - audio tag size.
         */
        private final static int FLV_AUDIO_TAG_SIZE = 1;
        
        // fields
        private OutputStream aos = null;
        private byte[] buf = null;
        
        /**
         * Constructor.
         * 
         * @param aos audio output stream
         */
        public ConvertFlvToMp3TagProcessor(final OutputStream aos)
        {
            this.aos = aos;
            this.buf = new byte[1024];
        }
        
        @Override
        public void onTag(final int type, final int dataSize, final long timestamp, InputStream is) throws IOException
        {
            switch (type)
            {
                case Flv.Tag.TYPE_AUDIO:
                {
                    if (dataSize > 0)
                    {
                        is.read(buf, 0, dataSize);
                        
                        aos.write(buf, FLV_AUDIO_TAG_SIZE, dataSize - FLV_AUDIO_TAG_SIZE);
                    }
                } break;
            }
        }
    }
    
    /**
     * <code>CheckCorrectnessTagProcessor</code> - check file correctness, i.e.
     * if audio/video tags are valid (known codec id).
     */
    private final static class CheckCorrectnessTagProcessor extends Flv.TagProcessorAdapter
    {
        // fields
        
        /**
         * Constructor.
         * 
         * @param info file info 
         */
        public CheckCorrectnessTagProcessor()
        {
        }
        
        @Override
        public void onTag(final int type, final int dataSize, final long timestamp, InputStream is) throws IOException
        {
            switch (type)
            {
                case Flv.Tag.TYPE_AUDIO:
                {
                    final int tag = is.read();
                    
                } break;
                case Flv.Tag.TYPE_VIDEO:
                {
                    final int tag = is.read();
                    
                } break;
                case Flv.Tag.TYPE_DATA:
                {
                    is.read();
                } break;
            }
        }
    }
    
    /**
     * <code>SplitTagProcessor</code>
     */
    private final static class SplitTagProcessor extends Flv.TagProcessorAdapter
    {
        // fields
        private OutputStream aos = null;
        private OutputStream vos = null;
        
        private long audioTimestamp = -1;
        private long videoTimestamp = -1;
        
        private byte[] buf = null;
        
        /**
         * Constructor.
         * 
         * @param aos audio output stream
         * @param vos video output stream
         */
        public SplitTagProcessor(final OutputStream aos, final OutputStream vos)
        {
            this.aos = aos;
            this.vos = vos;
            this.buf = new byte[1024];
        }
        
        @Override
        public void onTag(final int type, final int dataSize, final long timestamp, InputStream is) throws IOException
        {
            switch (type)
            {
                case Flv.Tag.TYPE_AUDIO:
                {
                    if (audioTimestamp < 0)
                    {
                        audioTimestamp = timestamp;
                        
                        Flv.writeHeader(aos, new Flv.Header(true /*audio*/, false /*video*/));
                    }
                    
                    is.read(buf, 0, dataSize);
                    
                    Flv.writeTag(aos, type, timestamp - audioTimestamp, buf, 0, dataSize);
                    
                } break;
                case Flv.Tag.TYPE_VIDEO:
                {
                    if (videoTimestamp < 0)
                    {
                        videoTimestamp = timestamp;
                        
                        Flv.writeHeader(vos, new Flv.Header(false /*audio*/, true /*video*/));
                    }
                    
                    final int vtag = is.read();
                    
                    if (getVideoFrame(vtag) < 0x05 /*command*/)
                    {
                        Flv.writeTagHeader(vos, type, timestamp - videoTimestamp, dataSize);
                        
                        // copy video frame data
                        vos.write(vtag);
                        
                        int size = dataSize - 1;
                        do
                        {
                            final int len = Math.min(buf.length, size);
                            
                            is.read(buf, 0, len);
                            
                            vos.write(buf, 0, len);
                            
                            size -= len;
                        }
                        while (size > 0);
                        
                        Flv.writeTagTrailer(vos, dataSize);
                    }
                } break;
            }
        }
    }
    
    /**
     * <code>TagKeeper</code> - tag keeper.
     */
    private final static class TagKeeper extends Flv.TagProcessorAdapter
    {
        // fields
        private final int type;
        private final byte[] buf;
        
        private long timestamp = -1;
        private ByteArray data = null;
        
        /**
         * Constructor.
         * 
         * @param type frame type to keep
         */
        public TagKeeper(final int type)
        {
            this.type = type;
            this.buf = new byte[8 * 1024];
        }
        
        /**
         * Checks if keeper is valid, i.e. has data.
         * 
         * @return <code>true</code> if valid; otherwise <code>false</code>
         */
        public final boolean isValid()
        {
            return timestamp >= 0;
        }
        
        /**
         * Returns kept frame timestamp.
         * 
         * @return frame timestamp
         */
        public final long getTimestamp()
        {
            return timestamp;
        }
        
        /**
         * Writes the tag to the {@link OutputStream <tt>os</tt>} and invalidates the keeper.
         * 
         * @param os
         * @throws IOException if an I/O exception occurred
         */
        public void write(OutputStream os) throws IOException
        {
            Flv.writeTag(os, type, timestamp, data.array, data.offset, data.length);
            
            timestamp = -1;
        }
        
        @Override
        public void onTag(final int type, final int dataSize, final long timestamp, InputStream is) throws IOException
        {
            if (this.type != type) return;
            
            this.timestamp = timestamp;
            
            if (dataSize > buf.length)
            {
                byte[] buf = new byte[dataSize];
                
                is.read(buf, 0, dataSize);
                
                data = new ByteArray(buf, 0, dataSize);
            }
            else
            {
                is.read(buf, 0, dataSize);
                
                data = new ByteArray(buf, 0, dataSize);
            }
            
            // post-process... a bit of hack
            if (type == Flv.Tag.TYPE_VIDEO)
            {
                if (getVideoFrame(data.array[data.offset]) >= 0x05 /*command*/)
                {
                    this.timestamp = -1;
                }
            }
        }
    }
    
    /**
     * Returns <code>file</code> data frames: available data frames.
     * 
     * @param file flv file
     * @return data frames: list of <code>Flv.ScriptData</code>
     * @throws Exception if an exception occurred
     */
    public static List<Flv.ScriptData[]> getDataFrames(final String file) throws Exception
    {
        List<Flv.ScriptData[]> frames = new ArrayList<Flv.ScriptData[]>();
        
        processFile(file, new GetDataFramesTagProcessor(frames));
        
        return frames;
    }
    
    /**
     * Returns <code>file</code> info: number of audio/video/data frames/bytes.
     * 
     * @param file flv file
     * @return file info
     * @throws Exception if an exception occurred
     */
    public static FileInfo getInfo(final String file) throws Exception
    {
        FileInfo info = new FileInfo();
        
        processFile(file, new GetInfoTagProcessor(info));
        
        return info;
    }
    
    /**
     * @param file
     * @throws Exception if an exception occurred
     */
    public static void check(final String file) throws Exception
    {
        processFile(file, new CheckCorrectnessTagProcessor());
    }
    
    /**
     * Merges audio stream from the <code>audioFile</code> and video stream from
     * the <code>videoFile</code> and writes the result to the <code>file</code>.
     * 
     * @param audioFile audio flv file
     * @param videoFile video flv file
     * @param file result flv file
     * @throws Exception if an exception occurred
     */
    public static void merge(final String audioFile, final String videoFile, final String file) throws Exception
    {
        if (!new File(audioFile).exists()) throw new IllegalArgumentException("Audio file doesn't exists!");
        if (!new File(videoFile).exists()) throw new IllegalArgumentException("Video file doesn't exists!");
        
        InputStream ais = null;
        InputStream vis = null;
        OutputStream os = null;
        
        try
        {
            ais = new FileInputStream(audioFile);
            vis = new FileInputStream(videoFile);
            
            os = new BufferedOutputStream(new FileOutputStream(file), 32 * 1024);
            
            Flv.readHeader(ais);
            Flv.readHeader(vis);
            Flv.writeHeader(os, new Flv.Header(true, true));
            
            TagKeeper audioTag = new TagKeeper(Flv.Tag.TYPE_AUDIO);
            TagKeeper videoTag = new TagKeeper(Flv.Tag.TYPE_VIDEO);
            
            while (true)
            {
                while (!videoTag.isValid() && Flv.readTag(vis, videoTag)) {}
                
                while (!audioTag.isValid() && Flv.readTag(ais, audioTag)) {}
                
                if (videoTag.isValid())
                {
                    if (audioTag.isValid())
                    {
                        if (videoTag.getTimestamp() <= audioTag.getTimestamp())
                        {
                            videoTag.write(os);
                        }
                        else
                        {
                            audioTag.write(os);
                        }
                    }
                    else
                    {
                        videoTag.write(os);
                    }
                }
                else
                {
                    if (audioTag.isValid())
                    {
                        audioTag.write(os);
                    }
                    else
                    {
                        break;
                    }
                }
            }
        }
        catch (Exception e)
        {
            throw new Exception(e);
        }
        finally
        {
            if (ais != null)
            {
                try
                {
                    ais.close();
                }
                catch (Exception e) {/*ignore*/}
            }
            
            if (vis != null)
            {
                try
                {
                    vis.close();
                }
                catch (Exception e) {/*ignore*/}
            }
            
            if (os != null)
            {
                try
                {
                    os.close();
                }
                catch (Exception e) {/*ignore*/}
            }
        }
    }
    
    /**
     * Converts mp3 audio stream from the <code>flvFile</code> to a new <code>mp3File</code>.
     * 
     * @param flvFile flv file with mp3 audio stream
     * @param mp3File mp3 audio file
     * @throws Exception if an exception occurred
     */
    public static void convertFlvToMp3(final String flvFile, final String mp3File) throws Exception
    {
        if (new File(mp3File).exists()) throw new IllegalArgumentException("Audio file exists!");
        
        OutputStream aos = null;
        
        try
        {
            ConvertFlvToMp3TagProcessor processor = new ConvertFlvToMp3TagProcessor(
                    aos = new BufferedOutputStream(new FileOutputStream(mp3File), 32 * 1024));
            
            processFile(flvFile, processor);
        }
        catch (Exception e)
        {
            throw new Exception(e);
        }
        finally
        {
            if (aos != null)
            {
                try
                {
                    aos.close();
                }
                catch (Exception e) {/*ignore*/}
            }
        }
    }
    
    /**
     * Splits flv <code>file</code> into <code>audioFile</code> and <code>videoFile</code>
     * flv files.
     * 
     * @param file flv file to split
     * @param audioFile flv file with audio stream
     * @param videoFile flv file with video stream
     * @throws Exception if an exception occurred
     */
    public static void split(final String file, final String audioFile, final String videoFile) throws Exception
    {
        if (new File(audioFile).exists()) throw new IllegalArgumentException("Audio file exists!");
        if (new File(videoFile).exists()) throw new IllegalArgumentException("Video file exists!");
        
        OutputStream aos = null;
        OutputStream vos = null;
        
        try
        {
            SplitTagProcessor processor = new SplitTagProcessor(
                    aos = new BufferedOutputStream(new FileOutputStream(audioFile), 32 * 1024),
                    vos = new BufferedOutputStream(new FileOutputStream(videoFile), 32 * 1024));
            
            processFile(file, processor);
        }
        catch (Exception e)
        {
            throw new Exception(e);
        }
        finally
        {
            if (aos != null)
            {
                try
                {
                    aos.close();
                }
                catch (Exception e) {/*ignore*/}
            }
            
            if (vos != null)
            {
                try
                {
                    vos.close();
                }
                catch (Exception e) {/*ignore*/}
            }
        }
    }
    
    
    /**
     * Processes the <code>file</code> using {@link Flv.ITagProcessor <tt>processor</tt>}.
     * 
     * @param file flv file
     * @param processor tag processor
     * @throws Exception if an exception occurred
     */
    public static void processFile(final String file, final Flv.ITagProcessor processor) throws Exception
    {
        FileInputStream fis = null;
        
        try
        {
            fis = new FileInputStream(new File(file));
            
            Flv.readHeader(fis);
            
            while (Flv.readTag(fis, processor)) {}
        }
        catch (Exception e)
        {
            throw new Exception(e);
        }
        finally
        {
            if (fis != null)
            {
                try
                {
                    fis.close();
                }
                catch (Exception e) {/*ignore*/}
            }
        }
    }
    
    /**
     * Constructor.
     */
    private FlvTool()
    {
    }
}