package org.pingles.cascading.cassandra;

import cascading.scheme.Scheme;
import cascading.tap.Tap;
import cascading.tuple.Fields;
import cascading.tuple.Tuple;
import cascading.tuple.TupleEntry;
import com.sun.source.tree.Tree;
import me.prettyprint.cassandra.serializers.TypeInferringSerializer;
import org.apache.cassandra.db.IColumn;
import org.apache.cassandra.hadoop.ColumnFamilyOutputFormat;
import org.apache.cassandra.hadoop.ConfigHelper;
import org.apache.cassandra.thrift.Column;
import org.apache.cassandra.thrift.ColumnOrSuperColumn;
import org.apache.cassandra.thrift.Mutation;
import org.apache.cassandra.thrift.SlicePredicate;
import org.apache.cassandra.utils.ByteBufferUtil;
import org.apache.commons.lang.NotImplementedException;
import org.apache.hadoop.mapred.JobConf;
import org.apache.hadoop.mapred.OutputCollector;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.charset.CharacterCodingException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.SortedMap;
import java.util.TreeMap;

public class CassandraScheme extends Scheme {
    private static final org.slf4j.Logger LOGGER = LoggerFactory.getLogger(CassandraScheme.class);
    private Fields keyField;
    private Fields[] columnFields;
    private Fields nameFields;

    /**
     * Creates a {@link Scheme} for dealing with a regular Column Family.
     * @param keyField      the field to use for the row key
     * @param columnFields  column names
     */
    public CassandraScheme(Fields keyField, Fields[] columnFields) {
        this.keyField = keyField;
        this.columnFields = columnFields;
    }

    public CassandraScheme(Fields fields) {
        this.nameFields = fields;
    }

    @Override
    public void sourceInit(Tap tap, JobConf jobConf) throws IOException {
        List<ByteBuffer> columnNames = new ArrayList<ByteBuffer>();

        for (int i = 0; i < nameFields.size(); i++) {
            Object columnName = nameFields.get(i);
            LOGGER.info("Adding input column name: {}", columnName);
            columnNames.add(TypeInferringSerializer.get().toByteBuffer(columnName));
        }
        SlicePredicate predicate = new SlicePredicate();
        predicate.setColumn_names(columnNames);
        ConfigHelper.setInputSlicePredicate(jobConf, predicate);
    }

    @Override
    public void sinkInit(Tap tap, JobConf jobConf) throws IOException {
        jobConf.setOutputKeyClass(ByteBuffer.class);
        jobConf.setOutputValueClass(Mutation.class);
        jobConf.setOutputFormat(ColumnFamilyOutputFormat.class);
    }

    @Override
    public Tuple source(Object key, Object value) {
        Tuple t = new Tuple();
        SortedMap<ByteBuffer, IColumn> values = (SortedMap<ByteBuffer, IColumn>) value;

        for (IColumn col : values.values()) {
            try {
                LOGGER.info("n: {}", ByteBufferUtil.string(col.name()));
            } catch (CharacterCodingException e) {
                throw new RuntimeException(e);
            }
            t.add(col.value());
        }

        // should pull column name as byte buffer selection from name fields?
//        int i = 0;
//        for (ByteBuffer n : values.keySet()) {
//            try {
//                LOGGER.info("n: {}", ByteBufferUtil.string(n));
//            } catch (CharacterCodingException e) {
//                throw new RuntimeException(e);
//            }
//
//            IColumn column = values.get(n);
//            tupleEntry.set(i, "test");
//            i++;
//        }

        return t;
    }

    @Override
    public void sink(TupleEntry tupleEntry, OutputCollector outputCollector) throws IOException {
        Tuple key = tupleEntry.selectTuple(keyField);
        TypeInferringSerializer<Object> serializer = TypeInferringSerializer.get();
        byte[] keyBytes = serializer.toBytes(key.get(0));

        for (Fields selector : columnFields) {
            TupleEntry values = tupleEntry.selectEntry(selector);

            for (int j = 0; j < values.getFields().size(); j++) {
                Fields fields = values.getFields();
                Tuple tuple = values.getTuple();

                Object name = fields.get(j);
                Object value = tuple.get(j);

                Mutation mutation = createColumnPutMutation(ByteBuffer.wrap(serializer.toBytes(name)), ByteBuffer.wrap(serializer.toBytes(value)));
                outputCollector.collect(ByteBuffer.wrap(keyBytes), Collections.singletonList(mutation));
            }
        }
    }

    private Mutation createColumnPutMutation(ByteBuffer name, ByteBuffer value) {
        Column column = new Column(name);
        column.setName(name);
        column.setValue(value);
        column.setTimestamp(System.currentTimeMillis());

        Mutation m = new Mutation();
        ColumnOrSuperColumn columnOrSuperColumn = new ColumnOrSuperColumn();
        columnOrSuperColumn.setColumn(column);
        m.setColumn_or_supercolumn(columnOrSuperColumn);

        return m;
    }
}

@@ -1392,48 +1392,38 @@
 me%7D 
-for dealing with a regular Column Family
+suitable for using with a Sink
 .%0A  
@@ -1678,24 +1678,127 @@
 lds;%0A    %7D%0A%0A
+    /**%0A     * Creates a %7B@link Scheme%7D suitable for using with a Source.%0A     * @param fields%0A     */%0A
     public C
@@ -3219,482 +3219,8 @@
  %7D%0A%0A
-        // should pull column name as byte buffer selection from name fields?%0A//        int i = 0;%0A//        for (ByteBuffer n : values.keySet()) %7B%0A//            try %7B%0A//                LOGGER.info(%22n: %7B%7D%22, ByteBufferUtil.string(n));%0A//            %7D catch (CharacterCodingException e) %7B%0A//                throw new RuntimeException(e);%0A//            %7D%0A//%0A//            IColumn column = values.get(n);%0A//            tupleEntry.set(i, %22test%22);%0A//            i++;%0A//        %7D%0A%0A
     

remove some commented code