package test.udp.traffic;
   
   import java.net.SocketAddress;
   import java.util.Collections;
   import java.util.LinkedList;
   import java.util.Random;
   import java.util.concurrent.TimeUnit;
   import java.util.concurrent.locks.Condition;
   import java.util.concurrent.locks.Lock;
  import java.util.concurrent.locks.ReentrantLock;
  
  import org.jboss.netty.buffer.ChannelBuffer;
  import org.jboss.netty.buffer.ChannelBuffers;
  import org.jboss.netty.channel.Channel;
  import org.jboss.netty.channel.ChannelFuture;
  import org.jboss.netty.channel.ChannelHandlerContext;
  import org.jboss.netty.channel.DownstreamMessageEvent;
  import org.jboss.netty.channel.MessageEvent;
  import org.jboss.netty.util.HashedWheelTimer;
  import org.jboss.netty.util.Timeout;
  import org.jboss.netty.util.Timer;
  import org.jboss.netty.util.TimerTask;
  
  public class LeakyBucketHandler
  {
  	public enum OverflowStrategy
  	{
  		WAIT,
  		DISCARD,
  		IGNORE,
  		ERROR
  	}
  	
  	// TODO object to queue: ctx and message event, so that we can send up/down pipline using ctx
  	class QueueEvent
  	{
  		MessageEvent _e;
  		ChannelHandlerContext _ctx;
  	}
  	
  	// maximal bucket size before overflow
  	int _bucketSize = 2*1024;
  	// what do we do if bucket overflows ?
  	OverflowStrategy _overflowStrategy = OverflowStrategy.WAIT;
  	// maximal output flow rate in bytes/s
  	double _maxOutputRate = 1024;
  	 // limit tx buffer size (required for mobile udp)
  	int _maxBufferSize = 512;
  	// avoid output of small buffers per tick. if required wait some ticks before output
  	// output buffers smaller than minBufferSplit, only for complete buffers from the queue.
  	int _minBufferSplit = 512;
  	// bucket leak interval
  	long _timerTick = 100;
  
  	
  	Timer _timer = new HashedWheelTimer(); //TODO will be set in constructor
  	LinkedList<MessageEvent> _queue = new LinkedList<MessageEvent>();
  	volatile int _currentBucketSize;
  	Lock _dataLock = new ReentrantLock();
  	Condition _queueNotEmpty = _dataLock.newCondition();
  	Condition _bucketNotFull = _dataLock.newCondition();
  	volatile int _waitingSize;
  	volatile boolean _stoped = true;
  	volatile Timeout _timeout;
  	volatile int _waitingTicks = 1;
  	volatile long _maxBytesPerTick = (long)(_maxOutputRate * _timerTick) / 1000;
  
  	TimerTask _timerTickTask = new TimerTask()
  	{
  
  		@Override
  		public void run(Timeout timeout) throws Exception
  		{
  			if (_stoped)
  				return;
  			doTimeTick();
  		}
  		
  	};
  	
  	synchronized public void submit(MessageEvent e)
  	{
  		System.out.println("+ input "+((ChannelBuffer)e.getMessage()).readableBytes());
  		_dataLock.lock();
  		try
  		{
  		ChannelBuffer buffer = (ChannelBuffer) e.getMessage();
  		// do we have enough space in the buffer ?
  		if (_currentBucketSize == 0 || buffer.readableBytes() + _currentBucketSize < _bucketSize)
  		{
  			// yes ! queue the data
  			queue(e);
  		}
  		else
  			// no! handle overflow
  			doOverflow(e);
  		}
  		finally
  		{
 			_dataLock.unlock();
 		}
 		System.out.println("- input");
 	}
 
 	private void queue(MessageEvent e)
 	{
 		// note: we are safe, lock was set in caller
 		System.out.println("+ q");
 		ChannelBuffer buffer = (ChannelBuffer) e.getMessage();
 		_currentBucketSize += buffer.readableBytes();
 		_queue.addFirst(e);
 		_queueNotEmpty.signal();	
 		System.out.println("- q");
 	}
 
 	private void doOverflow(MessageEvent e)
 	{
 		// TODO log overflow
 		switch (_overflowStrategy)
 		{
 		case WAIT: doWait(e); queue(e); break;
 		case DISCARD: break;
 		case IGNORE: queue(e);
 		case ERROR: throw new RuntimeException("Bucket Overflow");
 		}
 	}
 
 	private void doWait(MessageEvent e)
 	{
 		System.out.println("+ w");
 		ChannelBuffer buffer = (ChannelBuffer) e.getMessage();
 		_waitingSize = buffer.readableBytes();
 		try
 		{
 			_bucketNotFull.await();
 		}
 		catch (InterruptedException e1)
 		{
 			e1.printStackTrace();
 		}
 		_waitingSize = 0;
 		System.out.println("- w");
 	}
 	
 	private void doTimeTick()
 	{
 		System.out.println("+ t");
 		_dataLock.lock();
 		try
 		{
 			while (_queue.isEmpty())
 				try
 				{
 					System.out.println("emtpy");
 					_queueNotEmpty.await();
 				}
 				catch (InterruptedException e1)
 				{
 					e1.printStackTrace();
 				}
 			System.out.println("working");
 			// max bytes we can send in this tick
 			long max = _maxBytesPerTick * _waitingTicks;
 			System.out.println("max  "+max);
 			// TODO check if channel writeable
 			while (max > 0 && !_queue.isEmpty())
 			{
 				MessageEvent e = _queue.getLast();
 				ChannelBuffer b = (ChannelBuffer) e.getMessage();
 				// how many bytes will we send in this tick ?
 				int toSendSize = (int) Math.min(max, b.readableBytes());
 				toSendSize = Math.min(toSendSize, _maxBufferSize);
 				// if we need to split the current buffer make sure the split is not too small
 				if (toSendSize < _minBufferSplit && toSendSize != b.readableBytes() && toSendSize != _maxBufferSize)
 				{
 					// wait for next tick and allow more bytes in next tick
 					_waitingTicks++;
 					break;
 				}
 				_waitingTicks = 1;
 				// do we need to split the current buffer ?
 				if (toSendSize == b.readableBytes())
 				{
 					// no! output the complete buffer
 					_queue.removeLast();
 					System.out.println("remove buffer");
 					output(e, b);
 				}
 				else
 				{
 					// yes, split the buffer
 					// output the tail of the current buffer
 					// TODO when spliting: handle future correctly: create new one
 					ChannelBuffer toSendBuffer = b.slice(b.writerIndex()-toSendSize, toSendSize);
 					// keep the rest of the buffer
 					b.writerIndex(b.writerIndex()-toSendSize);
 					output(e, toSendBuffer);
 				}
 				max -= toSendSize;
 				_currentBucketSize -= toSendSize;
 			}
 			// is someone waiting to add a buffer to the queue ?
 			if (_waitingSize != 0 && _waitingSize+_currentBucketSize<_bucketSize)
 				_bucketNotFull.signal();				
 			
 		}
 		finally
 		{
 			_dataLock.unlock();
 		}
 		if (!_stoped)
 			setTimer();
 		System.out.println("- t");
 	}
 
 	private void setTimer()
 	{
 		_timeout = _timer.newTimeout(_timerTickTask, _timerTick, TimeUnit.MILLISECONDS);
 	}
 	
 	public void start()
 	{
 		if (!_stoped)
 			return;
 		_stoped = false;
 		setTimer();
 	}
 	
 	public void stop()
 	{
 		_stoped = true;
 		if (_timeout != null)
 			_timeout.cancel();
 		_timeout = null;
 		
 	}
 
 	private void output(MessageEvent e, ChannelBuffer b)
 	{
 		// TODO output up/downstream using QueueEvent
 		System.out.println("output "+b.readableBytes());
 	}
 	
 	public void setMaxRate(double rate)
 	{
 		_maxOutputRate  = rate;
 		_maxBytesPerTick = (long)(_maxOutputRate * _timerTick) / 1000;
 	}
 	
 	public void setTimerTick(long timerTick)
 	{
 		_timerTick = timerTick;
 		_maxBytesPerTick = (long)(_maxOutputRate * _timerTick) / 1000;		
 	}
 
 	static class TestMessageEvent implements MessageEvent
 	{
 		Object _m;
 		TestMessageEvent(ChannelBuffer b)
 		{
 			_m = b;
 		}
 
 		@Override
 		public Object getMessage()
 		{
 			// TODO Auto-generated method stub
 			return _m;
 		}
 
 		@Override
 		public SocketAddress getRemoteAddress()
 		{
 			// TODO Auto-generated method stub
 			return null;
 		}
 
 		@Override
 		public Channel getChannel()
 		{
 			// TODO Auto-generated method stub
 			return null;
 		}
 
 		@Override
 		public ChannelFuture getFuture()
 		{
 			// TODO Auto-generated method stub
 			return null;
 		}
 		
 	}
 	
 	public static void main(String[] args)
 	{
 		ChannelBuffer cb = ChannelBuffers.dynamicBuffer();
 		for (int i = 0; i< 200; i++)
 		{
 			cb.writeInt(i);
 		}
 		LeakyBucketHandler h = new LeakyBucketHandler();
 		h.start();
 		Random r = new Random(System.currentTimeMillis());
 		while (true)
 		{
 			int size = r.nextInt(800);
 			if (size == 0)
 				size = 10;
 			MessageEvent e = new TestMessageEvent(ChannelBuffers.wrappedBuffer(cb).slice(0, size));
 			h.submit(e);
 			try
 			{
 				long s = (long) (500*r.nextFloat());
 				s = (s/10)*10;
 				System.out.println("sleep "+s);
 				Thread.sleep(s);
 			}
 			catch (InterruptedException e1)
 			{
 				// TODO Auto-generated catch block
 				e1.printStackTrace();
 			}
 			
 		}
 	}
 
 
 }