/**
 * Copyright (c) 2006 - 2010 Smaxe Ltd (www.smaxe.com).
 * All rights reserved.
 */

import com.smaxe.uv.ProtocolLayerInfo;
import com.smaxe.uv.client.ICamera;
import com.smaxe.uv.client.IMicrophone;
import com.smaxe.uv.client.camera.AbstractCamera;
import com.smaxe.uv.client.rtmp.INetConnection;
import com.smaxe.uv.client.rtmp.INetStream;
import com.smaxe.uv.client.rtmp.License;
import com.smaxe.uv.client.rtmp.NetConnection;
import com.smaxe.uv.client.rtmp.NetStream;
import com.smaxe.uv.stream.MediaDataFactory;

import java.awt.Rectangle;
import java.awt.Robot;
import java.awt.image.BufferedImage;
import java.awt.image.DataBuffer;
import java.io.ByteArrayOutputStream;
import java.io.OutputStream;
import java.util.Map;
import java.util.zip.Deflater;
import java.util.zip.DeflaterOutputStream;

/**
 * <code>ExRtmpDesktopPublisherX</code> - publishes part of Desktop screen to the RTMP server and
 * provides means to manage upload bandwidth.
 * <p> Note:
 * <br> - This example encodes desktop using ScreenVideo codec implementation.
 * 
 * @author Andrei Sochirca
 */
public final class ExRtmpDesktopPublisherX extends Object
{
    /**
     * Entry point.
     * 
     * @param args
     * @throws Exception if an exception occurred
     */
    public static void main(final String[] args) throws Exception
    {
        // NOTE:
        // you can get Evaluation Key at:
        // http://www.smaxe.com/order.jsf#request_evaluation_key
        // or buy at:
        // http://www.smaxe.com/order.jsf
        License.setKey("SET-YOUR-KEY");
        
        
        final String url = "rtmp://localhost:1935/live";
        
        final DesktopCamera camera = new DesktopCamera(0 /*x*/, 0 /*y*/, 320 /*width*/, 240 /*height*/);
        
        new Thread(new Runnable()
        {
            public void run()
            {
                Publisher publisher = new Publisher();
                
                publisher.publish(url, "desktop", null /*microphone*/, camera);
            }
        }).start();
    }
    
    /**
     * <code>DesktopCamera</code> - {@link ICamera} implementation that captures desktop.
     * 
     * @author Andrei Sochirca
     */
    public final static class DesktopCamera extends AbstractCamera
    {
        /**
         * <code>CaptureRunnable</code> - {@link Runnable} implementation
         * that captures desktop.
         */
        private final class CaptureRunnable extends Object implements Runnable
        {
            // fields
            private volatile int x = 0;
            private volatile int y = 0;
            private volatile int width = 320;
            private volatile int height = 240;
            
            private volatile boolean active = true;
            
            private Deflater deflater = new Deflater();
            
            /**
             * Constructor.
             * 
             * @param x
             * @param y
             * @param width
             * @param height
             */
            public CaptureRunnable(final int x, final int y, final int width, final int height)
            {
                this.x = x;
                this.y = y;
                this.width = width;
                this.height = height;
            }
            
            /**
             * Sets the origin.
             * 
             * @param x
             * @param y
             */
            public void setOrigin(final int x, final int y)
            {
                this.x = x;
                this.y = y;
            }
            
            /**
             * Starts the capture.
             */
            public void start()
            {
            }
            
            /**
             * Stops the capture.
             */
            public void stop()
            {
            }
            
            /**
             * Releases the capture resources.
             */
            public void release()
            {
                active = false;
            }
            
            // Runnable implementation
            
            public void run()
            {
                final int blockWidth = 32;
                final int blockHeight = 32;
                
                final int timeBetweenFrames = 100; // 1000 / frameRate
                
                int frameCounter = 0;
                
                try
                {
                    Robot robot = new Robot();
                    
                    byte[] previous = null;
                    
                    while (active)
                    {
                        final long ctime = System.currentTimeMillis();
                        
                        BufferedImage image = robot.createScreenCapture(new Rectangle(x, y, width, height));
                        
                        byte[] current = toBGR(image);
                        
                        try
                        {
                            final byte[] packet = encode(current, previous, blockWidth, blockHeight, width, height);
                            
                            fireOnVideoData(MediaDataFactory.create(timeBetweenFrames, packet));
                            
                            previous = current;
                            
                            if (++frameCounter % 10 == 0) previous = null;
                        }
                        catch (Exception e)
                        {
                            e.printStackTrace();
                        }
                        
                        final int spent = (int) (System.currentTimeMillis() - ctime);
                        
                        Thread.sleep(Math.max(0, timeBetweenFrames - spent));
                    }
                }
                catch (Exception e)
                {
                    e.printStackTrace();
                }
            }
            
            // inner use methods
            /**
             * @param image
             * @return BGR image content
             */
            private byte[] toBGR(BufferedImage image)
            {
                final int width = image.getWidth();
                final int height = image.getHeight();
                
                byte[] buf = new byte[3 * width * height];
                
                final DataBuffer buffer = image.getData().getDataBuffer();
                
                for (int y = 0; y < height; y++)
                {
                    for (int x = 0; x < width; x++)
                    {
                        final int rgb = buffer.getElem(y * width + x);
                        final int offset = 3 * (y * width + x);
                        
                        buf[offset + 0] = (byte) (rgb & 0xFF);
                        buf[offset + 1] = (byte) ((rgb >> 8) & 0xFF);
                        buf[offset + 2] = (byte) ((rgb >> 16) & 0xFF);
                    }
                }
                
                return buf;
            }
            
            /**
             * Performs 'ScreenVideo' encode.
             * 
             * @param current
             * @param previous
             * @param blockWidth
             * @param blockHeight
             * @param width
             * @param height
             * @return buffer
             * @throws Exception
             */
            private byte[] encode(final byte[] current, final byte[] previous, final int blockWidth, final int blockHeight, final int width, final int height) throws Exception
            {
                ByteArrayOutputStream baos = new ByteArrayOutputStream(16 * 1024);
                
                if (previous == null)
                {
                    baos.write(getTag(0x01 /*key-frame*/, 0x03 /*ScreenVideo codec*/));
                }
                else
                {
                    baos.write(getTag(0x02 /*inter-frame*/, 0x03 /*ScreenVideo codec*/));
                }
                
                // write header
                final int wh = width + ((blockWidth / 16 - 1) << 12);
                final int hh = height + ((blockHeight / 16 - 1) << 12);
                
                writeShort(baos, wh);
                writeShort(baos, hh);
                
                // write content
                int y0 = height;
                int x0 = 0;
                int bwidth = blockWidth;
                int bheight = blockHeight;
                
                while (y0 > 0)
                {
                    bheight = Math.min(y0, blockHeight);
                    y0 -= bheight;
                    
                    bwidth = blockWidth;
                    x0 = 0;
                    
                    while (x0 < width)
                    {
                        bwidth = (x0 + blockWidth > width) ? width - x0 : blockWidth;
                        
                        final boolean changed = isChanged(current, previous, x0, y0, bwidth, bheight, width, height);
                        
                        if (changed)
                        {
                            ByteArrayOutputStream blaos = new ByteArrayOutputStream(4 * 1024);
                            
                            DeflaterOutputStream dos = new DeflaterOutputStream(blaos, deflater);
                            
                            for (int y = 0; y < bheight; y++)
                            {
                                dos.write(current, 3 * ((y0 + bheight - y - 1) * width + x0), 3 * bwidth);
                            }
                            
                            dos.finish();
                            deflater.reset();
                            
                            final byte[] bbuf = blaos.toByteArray();
                            final int written = bbuf.length;
                            
                            // write DataSize
                            writeShort(baos, written);
                            // write Data
                            baos.write(bbuf, 0, written);
                        }
                        else
                        {
                            // write DataSize
                            writeShort(baos, 0);
                        }
                        
                        x0 += bwidth;
                    }
                }
                
                return baos.toByteArray();
            }
            
            /**
             * Writes short value to the {@link OutputStream <tt>os</tt>}.
             * 
             * @param os
             * @param n
             * @throws Exception if an exception occurred
             */
            private void writeShort(OutputStream os, final int n) throws Exception
            {
                os.write((n >> 8) & 0xFF);
                os.write((n >> 0) & 0xFF);
            }
            
            /**
             * Checks if image block is changed.
             * 
             * @param current
             * @param previous
             * @param x0
             * @param y0
             * @param blockWidth
             * @param blockHeight
             * @param width
             * @param height
             * @return <code>true</code> if changed, otherwise <code>false</code>
             */
            public boolean isChanged(final byte[] current, final byte[] previous, final int x0, final int y0, final int blockWidth, final int blockHeight, final int width, final int height)
            {
                if (previous == null) return true;
                
                for (int y = y0, ny = y0 + blockHeight; y < ny; y++)
                {
                    final int foff = 3 * (x0 + width * y);
                    final int poff = 3 * (x0 + width * y);
                    
                    for (int i = 0, ni = 3 * blockWidth; i < ni; i++)
                    {
                        if (current[foff + i] != previous[poff + i]) return true;
                    }
                }
                
                return false;
            }
            
            /**
             * @param frame
             * @param codec
             * @return tag
             */
            public int getTag(final int frame, final int codec)
            {
                return ((frame & 0x0F) << 4) + ((codec & 0x0F) << 0);
            }
        }
        
        // fields
        private CaptureRunnable capture = null;
        private Thread t = null;
        
        /**
         * Constructor.
         */
        public DesktopCamera()
        {
            this(0 /*x*/, 0 /*y*/, 320 /*width*/, 240 /*height*/);
        }
        
        /**
         * Constructor.
         * 
         * @param x
         * @param y
         * @param width
         * @param height
         */
        public DesktopCamera(final int x, final int y, final int width, final int height)
        {
            capture = new CaptureRunnable(x, y, width, height);
        }
        
        /**
         * Starts desktop capture.
         */
        public void start()
        {
            if (t == null)
            {
                t = new Thread(capture);
                t.start();
            }
            
            capture.start();
        }
        
        /**
         * Stops video capture.
         */
        public void stop()
        {
            capture.stop();
        }
        
        /**
         * Releases the resources.
         */
        public void release()
        {
            capture.release();
            t = null;
        }
    }
    
    /**
     * <code>Publisher</code> - publisher.
     */
    public static final class Publisher extends Object
    {
        /**
         * <code>NetConnectionListener</code> - {@link NetConnection} listener implementation.
         */
        private final class NetConnectionListener extends NetConnection.ListenerAdapter
        {
            /**
             * Constructor.
             */
            public NetConnectionListener()
            {
            }
            
            @Override
            public void onAsyncError(final INetConnection source, final String message, final Exception e)
            {
                System.out.println("Publisher#NetConnection#onAsyncError: " + message + " " + e);
            }
            
            @Override
            public void onIOError(final INetConnection source, final String message)
            {
                System.out.println("Publisher#NetConnection#onIOError: " + message);
            }
            
            @Override
            public void onNetStatus(final INetConnection source, final Map<String, Object> info)
            {
                System.out.println("Publisher#NetConnection#onNetStatus: " + info);
                
                final Object code = info.get("code");
                
                if (NetConnection.CONNECT_SUCCESS.equals(code))
                {
                }
                else
                if (NetConnection.CONNECT_BANDWIDTH.equals(code))
                {
                    final ProtocolLayerInfo connectionInfo = (ProtocolLayerInfo) info.get("info");
                    
                    final long serverReadBytes = (Long) info.get("acknowledgement");
                    final long uploadBufferSize = (Long) info.get("uploadBufferSize");
                    
                    final long diff = connectionInfo.writtenBytes - serverReadBytes;
                    
                    // TO DO
                    // keep track of diff... increasing over time means
                    // that outgoing bandwidth exceeds connection upload bandwidth
                    
                    // keep track of uploadBufferSize.. increasing over time means
                    // that outgoing bandwidth exceeds NetConnection.Configuration.MAX_UPLOAD_BANDWIDTH property
                }
                else
                {
                    disconnected = true;
                }
            }
        }
        
        
        // fields
        private volatile boolean disconnected = false;
        
        
        /**
         * Publishes the stream.
         * 
         * @param url
         * @param streamName
         * @param microphone microphone
         * @param camera camera
         */
        public void publish(final String url, final String streamName, final IMicrophone microphone, final DesktopCamera camera)
        {
            final NetConnection connection = new NetConnection();
            
            connection.configuration().put(NetConnection.Configuration.INACTIVITY_TIMEOUT, -1);
            connection.configuration().put(NetConnection.Configuration.RECEIVE_BUFFER_SIZE, 256 * 1024);
            connection.configuration().put(NetConnection.Configuration.SEND_BUFFER_SIZE, 256 * 1024);
            
            // BANDWIDTH MANAGEMENT
            // 'Acknowledgement' event provides you with number of bytes read by the server
            connection.configuration().put(NetConnection.Configuration.ENABLE_ACKNOWLEDGEMENT_EVENT_NOTIFICATION, true);
            // Server notifies about number of read bytes each 128Kib (defined by property)
            // Note: Wowza server doesn't support this feature and notifies client every 640Kib
            connection.configuration().put(NetConnection.Configuration.WINDOW_ACKNOWLEDGEMENT_SIZE, 128 * 1024);
            // if max upload bandwidth is defined then library tries to do not exceed this value
            // by sending video frames in chunks (you can change it after connection
            // is established using NetConnection#setMaxUploadBandwidth(bandwidth) method)
            // Note: this feature is available in 1.5.7 or later
            connection.configuration().put(NetConnection.Configuration.MAX_UPLOAD_BANDWIDTH, 32 * 1024 /*bytes per second*/);
            
            connection.addEventListener(new NetConnectionListener());
            
            connection.connect(url);
            
            // wait till connected
            while (!connection.connected() && !disconnected)
            {
                try
                {
                    Thread.sleep(100);
                }
                catch (Exception e) {/*ignore*/}
            }
            
            if (!disconnected)
            {
                final NetStream stream = new NetStream(connection);
                
                stream.addEventListener(new NetStream.ListenerAdapter()
                {
                    @Override
                    public void onNetStatus(final INetStream source, final Map<String, Object> info)
                    {
                        System.out.println("Publisher#NetStream#onNetStatus: " + info);
                        
                        final Object code = info.get("code");
                        
                        if (NetStream.PUBLISH_START.equals(code))
                        {
                            if (microphone != null)
                            {
                                stream.attachAudio(microphone);
                            }
                            
                            if (camera != null)
                            {
                                stream.attachCamera(camera, -1 /*snapshotMilliseconds*/);
                                
                                camera.start();
                            }
                        }
                    }
                });
                
                
                stream.publish(streamName, NetStream.LIVE);
            }
            
            while (!disconnected)
            {
                try
                {
                    Thread.sleep(100);
                }
                catch (Exception e) {/*ignore*/}
            }
            
            connection.close();
        }
    }
}