package org.apache.cassandra.jmeter.config;
/*
 * Copyright 2014 Steven Lowenthal
 *
 * Licensed under the Apache License, Version 2.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.apache.org/licenses/LICENSE-2.0
 *
 * 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.
 */

import com.datastax.driver.core.Session;
import com.datastax.driver.core.policies.*;
import org.apache.jmeter.config.ConfigElement;
import org.apache.jmeter.testbeans.TestBean;
import org.apache.jmeter.testbeans.TestBeanHelper;
import org.apache.jmeter.testelement.AbstractTestElement;
import org.apache.jmeter.testelement.TestStateListener;
import org.apache.jmeter.threads.JMeterContextService;
import org.apache.jmeter.threads.JMeterVariables;
import org.apache.jorphan.logging.LoggingManager;
import org.apache.log.Logger;

import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.UnknownHostException;
import java.util.*;

public class CassandraConnection extends AbstractTestElement
    implements ConfigElement, TestStateListener, TestBean
    {

    // Load Balancer constants
    public static final String ROUND_ROBIN = "RoundRobin";
    public static final String DC_AWARE_ROUND_ROBIN = "DCAwareRoundRobin";
    public static final String DC_TOKEN_AWARE = "TokenAware(DCAwareRoundRobin)";
    public static final String WHITELIST = "WhiteList";
    public static final String DEFAULTLOADBALANCER = "Default";

    private static final Logger log = LoggingManager.getLoggerForClass();

    private static final long serialVersionUID = 233L;

    private transient String contactPoints, keyspace, username, password, sessionName, loadBalancer, localDataCenter;

    private final transient Set<InetAddress> contactPointsI = new HashSet<InetAddress>();
    private final transient Set<InetSocketAddress> contactPointsIS = new HashSet<InetSocketAddress>();

    // TODO - Add Port Number

    /*
     *  The datasource is set up by testStarted and cleared by testEnded.
     *  These are called from different threads, so access must be synchronized.
     *  The same instance is called in each case.
    */

    // Keep a record of the pre-thread pools so that they can be disposed of at the end of a test

    public CassandraConnection() {
    }

    public void testEnded() {
          CassandraSessionFactory.destroyClusters();
    }

    public void testEnded(String host) {
        testEnded();
    }

    @SuppressWarnings("deprecation") // call to TestBeanHelper.prepare() is intentional
    public void testStarted() {
        this.setRunningVersion(true);
        TestBeanHelper.prepare(this);
        JMeterVariables variables = getThreadContext().getVariables();
        LoadBalancingPolicy loadBalancingPolicy = null;

        if (loadBalancer.contentEquals(DC_AWARE_ROUND_ROBIN)) {
            // in driver v2.0.2+, we can use the default constructor on
            // dcawareroundrobinpolicy
            if (localDataCenter.isEmpty()) {
                loadBalancingPolicy = new DCAwareRoundRobinPolicy();
            }   else {
                loadBalancingPolicy = new DCAwareRoundRobinPolicy(localDataCenter);
            }
        } else if (loadBalancer.contentEquals(WHITELIST)) {
            loadBalancingPolicy = new WhiteListPolicy(new RoundRobinPolicy(), contactPointsIS);
        } else if (loadBalancer.contentEquals(ROUND_ROBIN)) {
            loadBalancingPolicy = new RoundRobinPolicy();
        } else if (loadBalancer.contentEquals(DC_TOKEN_AWARE)) {
            loadBalancingPolicy = new TokenAwarePolicy(new DCAwareRoundRobinPolicy());
        } else if (loadBalancer.contentEquals(DEFAULTLOADBALANCER)) {
            loadBalancingPolicy = null;
        }

        Session session = CassandraSessionFactory.createSession(sessionName, contactPointsI, keyspace, username, password, loadBalancingPolicy);

        variables.putObject(sessionName, session);
    }

    public void testStarted(String host) {
        testStarted();
    }

    @Override
    public Object clone() {
        return (CassandraConnection) super.clone();
    }

    /*
     * Utility routine to get the connection from the pool.
     * Purpose:
     * - allows CassandraSampler to be entirely independent of the pooling classes
     * - allows the pool storage mechanism to be changed if necessary
     */
    public static Session getSession(String sessionName) {
                return (Session) JMeterContextService.getContext().getVariables().getObject(sessionName);
     }

    // used to hold per-thread singleton connection pools
    private static final ThreadLocal<Map<String, Session>> perThreadPoolMap =
        new ThreadLocal<Map<String, Session>>(){
        @Override
        protected Map<String, Session> initialValue() {
            return new HashMap<String, Session>();
        }
    };



    public void addConfigElement(ConfigElement config) {
    }

    public boolean expectsModification() {
        return false;
    }

    /**
     * @return Returns the poolname.
     */
    public String getContactPoints() {
        return contactPoints.toString();
    }

    /**
     * @param contactPoints
     *            The poolname to set.
     */
    public void setContactPoints(String contactPoints) throws UnknownHostException {
        this.contactPoints = contactPoints;
        for (String contactPt : contactPoints.split(",")) {
            this.contactPointsI.add(InetAddress.getByName(contactPt));
            // TODO - 9160 should not really be hard coded.
            this.contactPointsIS.add(new InetSocketAddress(contactPt, 9042));
        }
    }

    /**
     * @return Returns the keyspace.
     */
    public String getKeyspace() {
        return keyspace;
    }

    /**
     * @param keyspace
     *            The keyspace to set.
     */
    public void setKeyspace(String keyspace) {
        this.keyspace = keyspace;
    }

    /**
     * @return Returns the password.
     */
    public String getPassword() {
        return password;
    }

    /**
     * @param password
     *            The password to set.
     */
    public void setPassword(String password) {
        this.password = password;
    }

    /**
     * @return Returns the username.
     */
    public String getUsername() {
        return username;
    }

    /**
     * @param username
     *            The username to set.
     */
    public void setUsername(String username) {
        this.username = username;
    }


   public String getSessionName() {
       return sessionName;
   }

   public void setSessionName(String sessionName) {
       this.sessionName = sessionName;
   }

   public String getLoadBalancer() {
       return loadBalancer;
   }

   public void setLoadBalancer(String loadBalancer) {
       this.loadBalancer = loadBalancer;
   }

    public String getLocalDataCenter() {
        return localDataCenter;
    }

    public void setLocalDataCenter(String localDataCenter) {
        this.localDataCenter = localDataCenter;
    }
 }