package backtype.storm.testing;

import backtype.storm.Config;
import backtype.storm.task.TopologyContext;
import backtype.storm.topology.OutputFieldsDeclarer;
import backtype.storm.transactional.TransactionAttempt;
import backtype.storm.coordination.BatchOutputCollector;
import backtype.storm.transactional.partitioned.IPartitionedTransactionalSpout;
import backtype.storm.transactional.partitioned.IPartitionedTransactionalSpout.Emitter;
import backtype.storm.tuple.Fields;
import backtype.storm.utils.RegisteredGlobalState;
import backtype.storm.utils.Utils;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * This spout only works in local mode.
 */
public class MemoryTransactionalSpout implements IPartitionedTransactionalSpout<MemoryTransactionalSpoutMeta> {
    public static String TX_FIELD = MemoryTransactionalSpout.class.getName() + "/id";
    
    private String _id;
    private String _finishedPartitionsId;
    private int _takeAmt;
    private Fields _outFields;
    
    public MemoryTransactionalSpout(Map<Integer, List<List<Object>>> partitions, Fields outFields, int takeAmt) {
        _id = RegisteredGlobalState.registerState(partitions);
        Map<Integer, Boolean> finished = Collections.synchronizedMap(new HashMap<Integer, Boolean>());
        _finishedPartitionsId = RegisteredGlobalState.registerState(finished);
        _takeAmt = takeAmt;
        _outFields = outFields;
    }
    
    public boolean isExhaustedTuples() {
        Map<Integer, Boolean> statuses = getFinishedStatuses();
        for(Integer partition: getQueues().keySet()) {
            if(!statuses.containsKey(partition) || !getFinishedStatuses().get(partition)) {
                return false;
            }
        }
        return true;
    }
    
    class Coordinator implements IPartitionedTransactionalSpout.Coordinator {

        @Override
        public int numPartitions() {
            return getQueues().size();
        }

        @Override
        public boolean isReady() {
            return true;
        }        
        
        @Override
        public void close() {
        }        
    }
    
    class Emitter implements IPartitionedTransactionalSpout.Emitter<MemoryTransactionalSpoutMeta> {
        
        Integer _maxSpoutPending;
        Map<Integer, Integer> _emptyPartitions = new HashMap<Integer, Integer>();
        
        public Emitter(Map conf) {
            Object c = conf.get(Config.TOPOLOGY_MAX_SPOUT_PENDING);
            if(c==null) _maxSpoutPending = 1;
            else _maxSpoutPending = Utils.getInt(c);
        }
        
        
        @Override
        public MemoryTransactionalSpoutMeta emitPartitionBatchNew(TransactionAttempt tx, BatchOutputCollector collector, int partition, MemoryTransactionalSpoutMeta lastPartitionMeta) {
            int index;
            if(lastPartitionMeta==null) {
                index = 0;
            } else {
                index = lastPartitionMeta.index + lastPartitionMeta.amt;
            }
            List<List<Object>> queue = getQueues().get(partition);
            int total = queue.size();
            int left = total - index;
            int toTake = Math.min(left, _takeAmt);
            
            MemoryTransactionalSpoutMeta ret = new MemoryTransactionalSpoutMeta(index, toTake);
            emitPartitionBatch(tx, collector, partition, ret);
            if(toTake==0) {
                // this is a pretty hacky way to determine when all the partitions have been committed
                // wait until we've emitted max-spout-pending empty partitions for the partition
                int curr = Utils.get(_emptyPartitions, partition, 0) + 1;
                _emptyPartitions.put(partition, curr);
                if(curr > _maxSpoutPending) {
                    getFinishedStatuses().put(partition, true);
                }
            }
            return ret;   
        }

        @Override
        public void emitPartitionBatch(TransactionAttempt tx, BatchOutputCollector collector, int partition, MemoryTransactionalSpoutMeta partitionMeta) {
            List<List<Object>> queue = getQueues().get(partition);
            for(int i=partitionMeta.index; i < partitionMeta.index + partitionMeta.amt; i++) {
                List<Object> toEmit = new ArrayList<Object>(queue.get(i));
                toEmit.add(0, tx);
                collector.emit(toEmit);                
            }
        }
                
        @Override
        public void close() {
        }        
    } 
    
    @Override
    public IPartitionedTransactionalSpout.Coordinator getCoordinator(Map conf, TopologyContext context) {
        return new Coordinator();
    }

    @Override
    public IPartitionedTransactionalSpout.Emitter<MemoryTransactionalSpoutMeta> getEmitter(Map conf, TopologyContext context) {
        return new Emitter(conf);
    }

    @Override
    public void declareOutputFields(OutputFieldsDeclarer declarer) {
        List<String> toDeclare = new ArrayList<String>(_outFields.toList());
        toDeclare.add(0, TX_FIELD);
        declarer.declare(new Fields(toDeclare));
    }

    @Override
    public Map<String, Object> getComponentConfiguration() {
        Config conf = new Config();
        conf.registerSerialization(MemoryTransactionalSpoutMeta.class);
        return conf;
    }
    
    public void startup() {
        getFinishedStatuses().clear();
    }
    
    public void cleanup() {
        RegisteredGlobalState.clearState(_id);
        RegisteredGlobalState.clearState(_finishedPartitionsId);
    }
    
    private Map<Integer, List<List<Object>>> getQueues() {
        return (Map<Integer, List<List<Object>>>) RegisteredGlobalState.getState(_id);
    }
    
    private Map<Integer, Boolean> getFinishedStatuses() {
        return (Map<Integer, Boolean>) RegisteredGlobalState.getState(_finishedPartitionsId);
    }
}

@@ -1056,24 +1056,89 @@
 _outFields;%0A
+    private Map%3CInteger, List%3CList%3CObject%3E%3E%3E _initialPartitions;%0A
     %0A    pub
@@ -1545,24 +1545,65 @@
  outFields;%0A
+        _initialPartitions = partitions;%0A
     %7D%0A    %0A 
@@ -5790,24 +5790,27 @@
 etQueues() %7B
+   
 %0A        ret
@@ -5798,38 +5798,70 @@
 () %7B   %0A        
-return
+Map%3CInteger, List%3CList%3CObject%3E%3E%3E ret =
  (Map%3CInteger, L
@@ -5913,24 +5913,98 @@
 State(_id);%0A
+        if(ret!=null) return ret;%0A        else return _initialPartitions;%0A
     %7D%0A    %0A 

make MemoryTransactionalSpout work on a cluster