package com.common.redis.serializer;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.redis.serializer.RedisSerializer;
import org.springframework.data.redis.serializer.SerializationException;

import java.io.ByteArrayOutputStream;

/**
 * @author xiaobzhou
 * date 2019-07-10 19:00
 */
public class KryoRedisSerializer<T> implements RedisSerializer<T> {

    private static final Logger LOG = LoggerFactory.getLogger(KryoRedisSerializer.class);

    public static final byte[] EMPTY_BYTE_ARRAY = new byte[0];

    /**
     * 初始化KRYO对象
     */
    private static final ThreadLocal<Kryo> KRYO = ThreadLocal.withInitial(Kryo::new);

    private Class<T> classType;

    public void cleanKryo() {
        KRYO.remove();
    }
    public KryoRedisSerializer(Class<T> classType) {
        super();
        this.classType = classType;
    }

    /**
     * Kryo序列化方法
     *
     * @param o
     * @return
     * @throws SerializationException
     */
    @Override
    public byte[] serialize(T o) throws SerializationException {

        if (o == null) {
            return EMPTY_BYTE_ARRAY;
        }

        Kryo kryo = KRYO.get();
        kryo.setReferences(false);
        kryo.register(classType);

        try {
            ByteArrayOutputStream byteOps = new ByteArrayOutputStream();
            Output output = new Output(byteOps);
            kryo.writeClassAndObject(output, o);
            output.flush();
            return byteOps.toByteArray();
        } catch (Exception ex) {
            LOG.error(ex.getMessage(), ex);
        }

        return EMPTY_BYTE_ARRAY;
    }

    /**
     * Kryo反序列化方法
     *
     * @param bytes
     * @return
     * @throws SerializationException
     */
    @Override
    public T deserialize(byte[] bytes) throws SerializationException {

        if (bytes != null) {
            Kryo kryo = KRYO.get();
            kryo.setReferences(false);
            kryo.register(classType);

            try {
                Input input = new Input(bytes);
                return (T) kryo.readClassAndObject(input);
            } catch (Exception ex) {
                LOG.error(ex.getMessage(), ex);
            }
        }

        return null;
    }
}