#
# Copyright (c) 2013 Big Switch Networks, Inc.
#
# Licensed under the Eclipse Public License, Version 1.0 (the
# "License"); you may not use this file except in compliance with the
# License. You may obtain a copy of the License at
#
#      http://www.eclipse.org/legal/epl-v10.html
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied. See the License for the specific language governing
# permissions and limitations under the License.
#

from thrift.transport import TTransport
from thrift.transport import TSocket
from thrift.protocol import TBinaryProtocol
from cassandra import Cassandra
from cassandra.ttypes import *

# FIXME: This class is derived from the django_cassandra backend.
# Should refactor this code better.

class CassandraConnection(object):
    def __init__(self, host, port, keyspace, user, password):
        self.host = host
        self.port = port
        self.keyspace = keyspace
        self.user = user
        self.password = password
        self.transport = None
        self.client = None
        self.keyspace_set = False
        self.logged_in = False
        
    def set_keyspace(self):
        if not self.keyspace_set:
            self.client.set_keyspace(self.keyspace)
            self.keyspace_set = True
    
    def login(self):
        # TODO: This user/password auth code hasn't been tested
        if not self.logged_in:
            if self.user:
                credentials = {'username': self.user, 'password': self.password}
                self.client.login(AuthenticationRequest(credentials))
            self.logged_in = True
            
    def connect(self, set_keyspace=False, login=False):
        if self.transport == None:
            # Create the client connection to the Cassandra daemon
            socket = TSocket.TSocket(self.host, int(self.port))
            transport = TTransport.TFramedTransport(TTransport.TBufferedTransport(socket))
            protocol = TBinaryProtocol.TBinaryProtocolAccelerated(transport)
            transport.open()
            self.transport = transport
            self.client = Cassandra.Client(protocol)
            
        if login:
            self.login()
        
        if set_keyspace:
            self.set_keyspace()
                
    def disconnect(self):
        if self.transport != None:
            self.transport.close()
            self.transport = None
            self.client = None
            self.keyspace_set = False
            self.logged_in = False
            
    def is_connected(self):
        return self.transport != None
    
    def reconnect(self):
        self.disconnect()
        self.connect(True, True)
    
    def get_client(self):
        if self.client == None:
            self.connect(True,True)
        return self.client