package org.apache.hadoop.hive.cassandra;

import org.apache.cassandra.thrift.Cassandra;
import org.apache.cassandra.thrift.InvalidRequestException;
import org.apache.commons.lang.StringUtils;
import org.apache.thrift.TException;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.TTransportException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class CassandraClientHolder {

    private static final Logger log = LoggerFactory.getLogger(CassandraClientHolder.class);

    private Cassandra.Client client;
    private final TTransport transport;
    private String keyspace;

    public CassandraClientHolder(TTransport transport) throws CassandraException {
        this(transport, null);
    }

    public CassandraClientHolder(TTransport transport, String keyspace) throws CassandraException {
        this.transport = transport;
        this.keyspace = keyspace;
        initClient();
    }

    public boolean isOpen() {
        return client != null && transport != null && transport.isOpen();
    }

    public String getKeyspace() {
        return keyspace;
    }

    private void initClient() throws CassandraException {
        try {
            transport.open();
        } catch (TTransportException e) {
            throw new CassandraException("unable to connect to server", e);
        }

        client = new Cassandra.Client(new TBinaryProtocol(transport));

        // connect to last known keyspace
        setKeyspace(keyspace);
    }

    /**
     * Set the client with the (potentially) new keyspace. Safe to call this
     * repeatedly with the same keyspace.
     *
     * @param keyspace
     * @return
     * @throws CassandraException
     */
    public void setKeyspace(String ks) throws CassandraException {
        if (ks == null) {
            return;
        }

        if (keyspace == null || !StringUtils.equals(keyspace, ks)) {
            try {
                this.keyspace = ks;
                client.set_keyspace(keyspace);
            } catch (InvalidRequestException e) {
                throw new CassandraException(e);
            } catch (TException e) {
                throw new CassandraException(e);
            }
        }
    }

    public Cassandra.Client getClient() {
        return client;
    }

    public void close() {
        if (transport == null || !transport.isOpen()) {
            return;
        }
        try {
            transport.flush();
        } catch (Exception e) {
            log.error("Could not flush transport for client holder: " + toString(), e);
        } finally {
            try {
                if (transport.isOpen()) {
                    transport.close();
                }
            } catch (Exception e) {
                log.error("Error on transport close for client: " + toString(), e);
            }
        }
    }

}
