package backtype.storm.transactional;

import backtype.storm.Config;
import backtype.storm.coordination.CoordinatedBolt.FinishedCallback;
import backtype.storm.coordination.FinishedTuple;
import backtype.storm.task.OutputCollector;
import backtype.storm.task.TopologyContext;
import backtype.storm.topology.IRichBolt;
import backtype.storm.topology.OutputFieldsDeclarer;
import backtype.storm.tuple.Tuple;
import backtype.storm.utils.TimeCacheMap;
import backtype.storm.utils.Utils;
import java.util.Map;

public class TransactionalBoltExecutor implements IRichBolt, FinishedCallback {
    byte[] _boltSer;
    TimeCacheMap<TransactionAttempt, ITransactionalBolt> _openTransactions;
    Map _conf;
    TopologyContext _context;
    TransactionalOutputCollectorImpl _collector;
    
    public TransactionalBoltExecutor(ITransactionalBolt bolt) {
        _boltSer = Utils.serialize(bolt);
    }
    
    @Override
    public void prepare(Map conf, TopologyContext context, OutputCollector collector) {
        _conf = conf;
        _context = context;
        _collector = new TransactionalOutputCollectorImpl(collector);
        _openTransactions = new TimeCacheMap<TransactionAttempt, ITransactionalBolt>(Utils.getInt(conf.get(Config.TOPOLOGY_MESSAGE_TIMEOUT_SECS)));
    }

    @Override
    public void execute(Tuple input) {
        TransactionAttempt attempt = (TransactionAttempt) input.getValue(0);
        String stream = input.getSourceStreamId();
        ITransactionalBolt bolt = _openTransactions.get(attempt);
        // bolt != null && receiving commit message -> this task processed the whole batch for this attempt
        // this is because: commit message only sent when tuple tree acked
        // if it failed after, then bolt will equal null, if it failed before, then it doesn't get acked
        //
        // this task processed the whole batch for this attempt && receiving commit message-> bolt != null
        // because the batch tuple is sent, guaranteeing that it sees at least one tuple for the batch
        if(stream.equals(TransactionalSpoutCoordinator.TRANSACTION_COMMIT_STREAM_ID)) {
                if(bolt!=null) {
                    ((ICommittable)bolt).commit();
                    _collector.ack(input);
                    _openTransactions.remove(attempt);
                } else {
                    _collector.fail(input);
                }
        } else {
            if(bolt==null) {
                bolt = newTransactionalBolt();
                bolt.prepare(_conf, _context, _collector, attempt);
                _openTransactions.put(attempt, bolt);            
            }

            // it is sent the batch id to guarantee that it creates the TransactionalBolt for the attempt before commit
            if(!stream.equals(TransactionalSpoutCoordinator.TRANSACTION_BATCH_STREAM_ID)) {
                bolt.execute(input);
            }       
            _collector.ack(input);
        }
    }

    @Override
    public void cleanup() {
    }

    @Override
    public void finishedId(FinishedTuple tup) {
        TransactionAttempt attempt = (TransactionAttempt) tup.getId();
        ITransactionalBolt bolt = _openTransactions.get(attempt);
        // can equal null if the TimeCacheMap expired it
        if(bolt!=null) {
            if(!(bolt instanceof ICommittable)) {
                _openTransactions.remove(attempt);
            }
            bolt.finishBatch();
        }
    }

    @Override
    public void declareOutputFields(OutputFieldsDeclarer declarer) {
        newTransactionalBolt().declareOutputFields(declarer);
    }
    
    @Override
    public Map<String, Object> getComponentConfiguration() {
        return newTransactionalBolt().getComponentConfiguration();
    }

    private ITransactionalBolt newTransactionalBolt() {
        return (ITransactionalBolt) Utils.deserialize(_boltSer);
    }
}

@@ -629,33 +629,36 @@
 sactionAttempt, 
-I
+Open
 TransactionalBol
@@ -652,22 +652,16 @@
 nsaction
-alBolt
 %3E _openT
@@ -1176,17 +1176,20 @@
 ttempt, 
-I
+Open
 Transact
@@ -1191,22 +1191,16 @@
 nsaction
-alBolt
 %3E(Utils.
@@ -1444,33 +1444,36 @@
 amId();%0A        
-I
+Open
 TransactionalBol
@@ -1463,35 +1463,27 @@
 nTransaction
-alBolt bolt
+ tx
  = _openTran
@@ -2117,34 +2117,47 @@
          if(
-bol
 t
+x
 !=null
+ && tx.finished
 ) %7B%0A        
@@ -2183,16 +2183,19 @@
 ittable)
+tx.
 bolt).co
@@ -2419,20 +2419,18 @@
      if(
-bol
 t
+x
 ==null) 
@@ -2451,18 +2451,21 @@
     
-bol
 t
+x
  = new
+ Open
 Tran
@@ -2471,22 +2471,16 @@
 nsaction
-alBolt
 ();%0A    
@@ -2487,24 +2487,27 @@
             
+tx.
 bolt.prepare
@@ -2593,20 +2593,18 @@
 ttempt, 
-bol
 t
+x
 );      
@@ -2849,24 +2849,27 @@
             
+tx.
 bolt.execute
@@ -3140,17 +3140,20 @@
         
-I
+Open
 Transact
@@ -3155,27 +3155,19 @@
 nsaction
-alBolt bolt
+ tx
  = _open
@@ -3261,20 +3261,18 @@
      if(
-bol
 t
+x
 !=null) 
@@ -3290,20 +3290,18 @@
    if(!(
-bol
 t
+x
  instanc
@@ -3398,16 +3398,19 @@
         
+tx.
 bolt.fin
@@ -3417,24 +3417,56 @@
 ishBatch();%0A
+            tx.finished = true;%0A
         %7D%0A  
@@ -3769,24 +3769,177 @@
 ion();%0A    %7D
+%0A    %0A    private class OpenTransaction %7B%0A        public boolean finished = false;%0A        public ITransactionalBolt bolt = newTransactionalBolt();%0A    %7D
 %0A%0A    privat

make transactionalbolt enforce that it only commits when it has received a finishedBatch message