package org.mortbay.cometd.client;

import java.io.InputStreamReader;
import java.io.LineNumberReader;
import java.net.InetSocketAddress;
import java.security.SecureRandom;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

import org.mortbay.cometd.AbstractBayeux;
import org.mortbay.jetty.client.HttpClient;
import org.mortbay.jetty.client.HttpDestination;
import org.mortbay.thread.BoundedThreadPool;
import org.mortbay.util.ajax.JSON;

import dojox.cometd.Bayeux;
import dojox.cometd.Client;
import dojox.cometd.Listener;

public class BayeuxLoadGenerator
{
    SecureRandom _random= new SecureRandom();
    HttpClient http;
    InetSocketAddress address=new InetSocketAddress("localhost",8080);
    String uri="/cometd";
    int clients=10000;
    BayeuxClient[] client;
    int _usersPerRoom=10;
    int _burst=20;
    long _pauseConnect=10;
    long _pause=1000;
    int _chat=1000;
    long _minLatency;
    long _maxLatency;
    long _totalLatency;
    int _got;
    AtomicInteger _subscribed = new AtomicInteger();
    
    public BayeuxLoadGenerator() throws Exception
    {
        http=new HttpClient();
        
        http.setConnectorType(HttpClient.CONNECTOR_SELECT_CHANNEL);
        // http.setConnectorType(HttpClient.CONNECTOR_SOCKET);
        http.setMaxConnectionsPerAddress(10000);
        
        BoundedThreadPool pool = new BoundedThreadPool();
        pool.setMaxThreads(255);
        pool.setDaemon(true);
        http.setThreadPool(pool);
        http.start();
       
        
    }
    
    
    public void generateLoad() throws Exception
    {
        LineNumberReader in = new LineNumberReader(new InputStreamReader(System.in));

        System.err.print("chat size: ");
        String t = in.readLine().trim();
        int size=Integer.parseInt(t);
        String chat="";
        for (int i=0;i<size;i++)
            chat+="x";
        
        System.err.print("clients: ");
        t = in.readLine().trim();
        clients=Integer.parseInt(t);
        System.err.print("per room: ");
        t = in.readLine().trim();
        _usersPerRoom=Integer.parseInt(t);
        
        client = new BayeuxClient[clients];

        HttpDestination destination = http.getDestination(address,false);
        
        for (int i=0;i<clients;i++)
        {
            client[i] = new BayeuxClient(http,address,uri)
            {
                public void deliver(Client from, Map<String, Object> message)
                {
                    if (Bayeux.META_SUBSCRIBE.equals(message.get(Bayeux.CHANNEL_FIELD)) &&
                        ((Boolean)message.get(Bayeux.SUCCESSFUL_FIELD)).booleanValue())
                        _subscribed.incrementAndGet();
                    super.deliver(from,message);
                }
            };
            
            int room=i/_usersPerRoom;
            final String name = "Client"+(i>=1000?"":i>=100?"0":i>=10?"00":"000")+i;
            final String id=client[i].getClientId();
            if (i%_burst==0)
            {
                System.err.println(name);
                Thread.sleep(_pauseConnect*_burst);
            }
            
            Listener listener = new Listener()
            {
                public void deliver(Client fromClient, Client toClient, Map<String, Object> msg)
                {
                    Object data=(Object)msg.get(AbstractBayeux.DATA_FIELD);
                    if (data!=null)
                    {
                        String msgId=(String)msg.get(AbstractBayeux.ID_FIELD);
                        // System.err.println(name+": "+data);
                        if (msgId!=null)
                        {
                            long latency= System.currentTimeMillis()-Long.parseLong(msgId);
                            synchronized(BayeuxLoadGenerator.this)
                            {
                                _got++;
                                if (_maxLatency<latency)
                                    _maxLatency=latency;
                                if (_minLatency==0 || latency<_minLatency)
                                    _minLatency=latency;
                                _totalLatency+=latency;
                            }
                        }
                    }
                }

                public void removed(String clientId, boolean timeout)
                {
                }
                
                public String toString()
                {
                    return name+"/"+id;
                }
            };
            client[i].setListener(listener);
            client[i].subscribe(room>0?("/chat/demo/"+room):"/chat/demo");
            client[i].start();
            Thread.sleep(_pauseConnect);
        }

        Thread.sleep(10*_pause);
        
        while(_subscribed.get()<clients)
        {
            // System.err.println(destination.toDetailString());
            System.err.println("Subscribed:"+_subscribed.get()+" < "+clients+" ...");
            Thread.sleep(1000);
        }
        
        
        while (true)
        {
            synchronized(this)
            {
                _got=0;
                _minLatency=0;
                _maxLatency=0;
                _totalLatency=0;
            }
            
            System.err.print("chat: ");
            t = in.readLine().trim();
            _chat=Integer.parseInt(t);
            if (_chat<=0)
                break;
            
            System.err.print("pause: ");
            t = in.readLine().trim();
            _pause=Integer.parseInt(t);
            System.err.print("burst: ");
            t = in.readLine().trim();
            _burst=Integer.parseInt(t);
            

            long start=System.currentTimeMillis();
            for (int i=1;i<=_chat;i++)
            {
                // System.err.print(i);
                // System.err.print(',');
                int c = _random.nextInt(clients);
                int room=c/_usersPerRoom;
                final String name = "Client"+(i>=1000?"":i>=100?"0":i>=10?"00":"000")+i;
                client[c].publish(room>0?("/chat/demo/"+room):"/chat/demo",new JSON.Literal("{\"user\":\""+name+"\",\"chat\":\""+chat+" "+i+"\"}"),""+System.currentTimeMillis());

                if (i%10==0)
                {
                    System.err.print('.');
                    if (i%1000==0)
                        System.err.println();
                }
                
                if (_pause>0 && i%_burst==0)
                {
                    Thread.sleep(_pause);
                }
                
                    
            }

            Thread.sleep(_maxLatency);

            for (int i=0;i<clients;i++)
            {
                if (!client[i].isPolling())
                    System.err.println("PROBLEM WITH "+client[i]);
            }

            System.err.println();

            long last=0;
            if (_got<_chat*_usersPerRoom) 
                Thread.sleep(500);
            
            while (_got<_chat*_usersPerRoom) 
            {
                // System.err.println(destination.toDetailString());
                System.err.println("Got:"+_got+" < "+(_chat*_usersPerRoom)+" ...");
                Thread.sleep(2000);
                if (last!=0 && _got==last)
                    break;
                last=_got;
            }

            long end=System.currentTimeMillis();

            System.err.println("Got "+_got+" at "+(_got*1000/(end-start))+"/s, latency min/ave/max ="+_minLatency+"/"+(_totalLatency/_got)+"/"+_maxLatency+"ms");
        }
        

        for (int i=0;i<clients;i++)
        {
            client[i].remove(false);
        }

    }
    
    
    public static void main(String[] args) throws Exception
    {
        BayeuxLoadGenerator gen = new BayeuxLoadGenerator();
        
        if (args.length>0)
            gen.clients=Integer.parseInt(args[0]);
        if (args.length>1)
            gen._chat=Integer.parseInt(args[1]);
            
        
        gen.generateLoad();
    }
    
}
