package elephantdb.serialize;

import cascading.kryo.KryoFactory;
import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.apache.hadoop.conf.Configuration;
import org.apache.log4j.Logger;

import java.util.ArrayList;
import java.io.ByteArrayOutputStream;
import java.io.IOException;

/**
 * User: sritchie
 * Date: 1/2/12
 * Time: 12:22 PM
 */
public class KryoSerializer implements Serializer {
    public static final Logger LOG = Logger.getLogger(KryoSerializer.class);
    
    private static final ThreadLocal<Kryo> kryo = new ThreadLocal<Kryo>();
    private Iterable<KryoFactory.ClassPair> kryoPairs = new ArrayList<KryoFactory.ClassPair>();
    private ByteArrayOutputStream byteStream;

    public KryoSerializer() {
    }

    public KryoSerializer(Iterable<KryoFactory.ClassPair> kryoPairs) {
        setKryoPairs(kryoPairs);
    }

    private Kryo freshKryo() {
        Kryo k = new Kryo();
        byteStream = new ByteArrayOutputStream();
        KryoFactory factory = new KryoFactory(new Configuration());
        k.setInstantiatorStrategy(new StdInstantiatorStrategy());
        k.setRegistrationRequired(false);
        factory.registerBasic(k, getKryoPairs());
        return k;
    }

    public void setKryoPairs(Iterable<KryoFactory.ClassPair> pairs) {
        kryoPairs = pairs;
    }

    public Iterable<KryoFactory.ClassPair> getKryoPairs() {
        return kryoPairs;
    }

    public Kryo getKryo() {
        if (kryo.get() == null)
            kryo.set(freshKryo());

        return kryo.get();
    }

    public byte[] serialize(Object o) {
        byteStream.reset();
        Output ko = new Output(byteStream);
        getKryo().writeClassAndObject(ko, o);
        ko.flush();
        byte[] bytes = ko.toBytes();
        
        return bytes;
    }

    public Object deserialize(byte[] bytes) {
        return getKryo().readClassAndObject(new Input(bytes));
    }

    public <T> T deserialize(byte[] bytes, Class<T> klass) {
        return getKryo().readObject(new Input(bytes), klass);
    }
}

@@ -678,24 +678,139 @@
 al%3CKryo%3E();%0A
+    private static final ThreadLocal%3CByteArrayOutputStream%3E byteStream = new ThreadLocal%3CByteArrayOutputStream%3E();%0A
     private 
@@ -896,54 +896,8 @@
 %3E();
-%0A    private ByteArrayOutputStream byteStream;
 %0A%0A  
@@ -1720,32 +1720,223 @@
 ryo.get();%0A    %7D
+%0A    %0A    public ByteArrayOutputStream getByteStream() %7B%0A        if (byteStream.get() == null)%0A            byteStream.set(new ByteArrayOutputStream());%0A%0A        return byteStream.get();%0A    %7D
 %0A%0A    public byt
@@ -1961,33 +1961,36 @@
 ct o) %7B%0A        
-b
+getB
 yteStream.reset(
@@ -1982,16 +1982,18 @@
 teStream
+()
 .reset()
@@ -2025,17 +2025,20 @@
  Output(
-b
+getB
 yteStrea
@@ -2038,16 +2038,18 @@
 teStream
+()
 );%0A     

make byteStream a threadlocal