package com.continuuity.test.internal;

import com.continuuity.api.data.OperationException;
import com.continuuity.api.flow.flowlet.StreamEvent;
import com.continuuity.common.queue.QueueName;
import com.continuuity.data.operation.OperationContext;
import com.continuuity.data.operation.executor.OperationExecutor;
import com.continuuity.data.operation.ttqueue.QueueEnqueue;
import com.continuuity.data.operation.ttqueue.QueueEntry;
import com.continuuity.data.operation.ttqueue.QueueProducer;
import com.continuuity.streamevent.DefaultStreamEvent;
import com.continuuity.streamevent.StreamEventCodec;
import com.continuuity.test.StreamWriter;
import com.google.common.base.Charsets;
import com.google.common.collect.ImmutableMap;
import com.google.inject.Inject;
import com.google.inject.assistedinject.Assisted;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Map;

/**
 *
 */
public final class DefaultStreamWriter implements StreamWriter {

  private final OperationExecutor opex;
  private final OperationContext opCtx;
  private final QueueProducer queueProducer;
  private final QueueName queueName;
  private final StreamEventCodec codec;

  @Inject
  public DefaultStreamWriter(OperationExecutor opex,
                             @Assisted QueueName queueName,
                             @Assisted("accountId") String accountId,
                             @Assisted("applicationId") String applicationId) {
    this.opex = opex;
    this.opCtx = new OperationContext(accountId, applicationId);
    this.queueProducer = new QueueProducer("Testing");

    this.queueName = queueName;
    this.codec = new StreamEventCodec();
  }

  @Override
  public void send(String content) throws IOException {
    send(Charsets.UTF_8.encode(content));
  }

  @Override
  public void send(byte[] content) throws IOException {
    send(content, 0, content.length);
  }

  @Override
  public void send(byte[] content, int off, int len) throws IOException {
    send(ByteBuffer.wrap(content, off, len));
  }

  @Override
  public void send(ByteBuffer buffer) throws IOException {
    send(ImmutableMap.<String, String>of(), buffer);
  }

  @Override
  public void send(Map<String, String> headers, String content) throws IOException {
    send(headers, Charsets.UTF_8.encode(content));
  }

  @Override
  public void send(Map<String, String> headers, byte[] content) throws IOException {
    send(headers, content, 0, content.length);
  }

  @Override
  public void send(Map<String, String> headers, byte[] content, int off, int len) throws IOException {
    send(headers, ByteBuffer.wrap(content, off, len));
  }

  @Override
  public void send(Map<String, String> headers, ByteBuffer buffer) throws IOException {
    StreamEvent event = new DefaultStreamEvent(ImmutableMap.copyOf(headers), buffer);
    QueueEnqueue enqueue = new QueueEnqueue(queueProducer, queueName.toBytes(),
                                            new QueueEntry(codec.encodePayload(event)));
    try {
      opex.commit(opCtx, enqueue);
    } catch (OperationException e) {
      throw new IOException(e);
    }
  }
}

@@ -223,32 +223,26 @@
 eration.
-OperationContext
+StatusCode
 ;%0Aimport
@@ -277,34 +277,26 @@
 ion.
-executor.OperationExecutor
+ttqueue.QueueEntry
 ;%0Aim
@@ -316,37 +316,26 @@
 inuuity.data
+2
 .
-operation.tt
 queue.QueueE
@@ -333,23 +333,25 @@
 ue.Queue
-Enqueue
+2Producer
 ;%0Aimport
@@ -375,35 +375,32 @@
 data
-.operation.ttqueue.QueueEnt
+2.queue.QueueClientFacto
 ry;%0A
@@ -430,40 +430,158 @@
 data
-.operation.ttqueue.QueueProducer
+2.transaction.Transaction;%0Aimport com.continuuity.data2.transaction.TransactionAware;%0Aimport com.continuuity.data2.transaction.TransactionSystemClient
 ;%0Aim
@@ -1075,115 +1075,87 @@
 nal 
-OperationExecutor opex;%0A  private final OperationContext opCtx;%0A  private final QueueProducer queueProducer
+Queue2Producer producer;%0A  private final TransactionSystemClient txSystemClient
 ;%0A  
@@ -1273,30 +1273,114 @@
 ter(
-OperationExecutor opex
+QueueClientFactory queueClientFactory,%0A                             TransactionSystemClient txSystemClient
 ,%0A  
@@ -1589,152 +1589,137 @@
 Id) 
-%7B%0A    this.opex = opex;%0A    this.opCtx = new OperationContext(accountId, applicationId);%0A    this.queueProducer = new QueueProducer(%22Testing%22);%0A
+throws IOException %7B%0A%0A    this.producer = queueClientFactory.createProducer(queueName);%0A    this.txSystemClient = txSystemClient;
 %0A   
@@ -2956,218 +2956,572 @@
     
-QueueEnqueue enqueue = new QueueEnqueue(queueProducer, queueName.toBytes(),%0A                                            new QueueEntry(codec.encodePayload(event)));%0A    try %7B%0A      opex.commit(opCtx, enqueue
+TransactionAware txAware = (TransactionAware) producer;%0A%0A    // start tx to write in queue in tx%0A    Transaction tx = txSystemClient.start();%0A    txAware.startTx(tx);%0A    try %7B%0A      producer.enqueue(new QueueEntry(codec.encodePayload(event)));%0A      if (!txSystemClient.canCommit(tx, txAware.getTxChanges()) %7C%7C !txAware.commitTx() %7C%7C !txSystemClient.commit(tx)) %7B%0A        throw new OperationException(StatusCode.TRANSACTION_CONFLICT, %22Fail to commit%22);%0A      %7D%0A      txAware.postTxCommit();%0A    %7D catch (Exception e) %7B%0A      try %7B%0A        txAware.rollbackTx(
 );%0A
+  
     
@@ -3533,31 +3533,118 @@
 ch (
-OperationException e) %7B
+Exception ex) %7B%0A        throw new IOException(ex);%0A      %7D finally %7B%0A        txSystemClient.abort(tx);%0A      %7D
 %0A   

Use new stream in StreamWriter in test framework