/*
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 *
 */
package org.apache.cassandra.stress;


import java.io.IOError;
import java.io.IOException;
import java.io.InputStream;
import java.io.Serializable;
import java.net.URI;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

import com.google.common.base.Function;
import com.google.common.collect.Iterables;
import com.google.common.util.concurrent.Uninterruptibles;

import com.datastax.driver.core.*;
import com.datastax.driver.core.exceptions.AlreadyExistsException;
import org.apache.cassandra.config.CFMetaData;
import org.apache.cassandra.config.ColumnDefinition;
import org.apache.cassandra.cql3.QueryProcessor;
import org.apache.cassandra.cql3.statements.CreateKeyspaceStatement;
import org.apache.cassandra.cql3.statements.CreateTableStatement;
import org.apache.cassandra.exceptions.RequestValidationException;
import org.apache.cassandra.exceptions.SyntaxException;
import org.apache.cassandra.stress.generate.*;
import org.apache.cassandra.stress.generate.values.*;
import org.apache.cassandra.stress.operations.userdefined.TokenRangeQuery;
import org.apache.cassandra.stress.operations.userdefined.SchemaInsert;
import org.apache.cassandra.stress.operations.userdefined.SchemaQuery;
import org.apache.cassandra.stress.operations.userdefined.ValidatingSchemaQuery;
import org.apache.cassandra.stress.settings.*;
import org.apache.cassandra.stress.util.JavaDriverClient;
import org.apache.cassandra.stress.util.ThriftClient;
import org.apache.cassandra.stress.util.Timer;
import org.apache.cassandra.thrift.Compression;
import org.apache.cassandra.thrift.ThriftConversion;
import org.apache.thrift.TException;
import org.yaml.snakeyaml.Yaml;
import org.yaml.snakeyaml.constructor.Constructor;
import org.yaml.snakeyaml.error.YAMLException;

import static org.apache.cassandra.io.sstable.CQLSSTableWriter.parseStatement;

public class StressProfile implements Serializable
{
    private String keyspaceCql;
    private String tableCql;
    private List<String> extraSchemaDefinitions;
    public final String seedStr = "seed for stress";

    public String keyspaceName;
    public String tableName;
    private Map<String, GeneratorConfig> columnConfigs;
    private Map<String, StressYaml.QueryDef> queries;
    public Map<String, StressYaml.TokenRangeQueryDef> tokenRangeQueries;
    private Map<String, String> insert;

    transient volatile TableMetadata tableMetaData;
    transient volatile Set<TokenRange> tokenRanges;

    transient volatile GeneratorFactory generatorFactory;

    transient volatile BatchStatement.Type batchType;
    transient volatile DistributionFactory partitions;
    transient volatile RatioDistributionFactory selectchance;
    transient volatile RatioDistributionFactory rowPopulation;
    transient volatile PreparedStatement insertStatement;
    transient volatile Integer thriftInsertId;
    transient volatile List<ValidatingSchemaQuery.Factory> validationFactories;

    transient volatile Map<String, SchemaQuery.ArgSelect> argSelects;
    transient volatile Map<String, PreparedStatement> queryStatements;
    transient volatile Map<String, Integer> thriftQueryIds;

    private static final Pattern lowercaseAlphanumeric = Pattern.compile("[a-z0-9_]+");

    private void init(StressYaml yaml) throws RequestValidationException
    {
        keyspaceName = yaml.keyspace;
        keyspaceCql = yaml.keyspace_definition;
        tableName = yaml.table;
        tableCql = yaml.table_definition;
        queries = yaml.queries;
        tokenRangeQueries = yaml.token_range_queries;
        insert = yaml.insert;

        extraSchemaDefinitions = yaml.extra_definitions;

        assert keyspaceName != null : "keyspace name is required in yaml file";
        assert tableName != null : "table name is required in yaml file";
        assert queries != null : "queries map is required in yaml file";

        for (String query : queries.keySet())
            assert !tokenRangeQueries.containsKey(query) :
                String.format("Found %s in both queries and token_range_queries, please use different names", query);

        if (keyspaceCql != null && keyspaceCql.length() > 0)
        {
            try
            {
                String name = ((CreateKeyspaceStatement) QueryProcessor.parseStatement(keyspaceCql)).keyspace();
                assert name.equalsIgnoreCase(keyspaceName) : "Name in keyspace_definition doesn't match keyspace property: '" + name + "' != '" + keyspaceName + "'";
            }
            catch (SyntaxException e)
            {
                throw new IllegalArgumentException("There was a problem parsing the keyspace cql: " + e.getMessage());
            }
        }
        else
        {
            keyspaceCql = null;
        }

        if (tableCql != null && tableCql.length() > 0)
        {
            try
            {
                String name = CFMetaData.compile(tableCql, keyspaceName).cfName;
                assert name.equalsIgnoreCase(tableName) : "Name in table_definition doesn't match table property: '" + name + "' != '" + tableName + "'";
            }
            catch (RuntimeException e)
            {
                throw new IllegalArgumentException("There was a problem parsing the table cql: " + e.getMessage());
            }
        }
        else
        {
            tableCql = null;
        }

        columnConfigs = new HashMap<>();

        if (yaml.columnspec != null)
        {
            for (Map<String, Object> spec : yaml.columnspec)
            {
                lowerCase(spec);
                String name = (String) spec.remove("name");
                DistributionFactory population = !spec.containsKey("population") ? null : OptionDistribution.get((String) spec.remove("population"));
                DistributionFactory size = !spec.containsKey("size") ? null : OptionDistribution.get((String) spec.remove("size"));
                DistributionFactory clustering = !spec.containsKey("cluster") ? null : OptionDistribution.get((String) spec.remove("cluster"));

                if (!spec.isEmpty())
                    throw new IllegalArgumentException("Unrecognised option(s) in column spec: " + spec);
                if (name == null)
                    throw new IllegalArgumentException("Missing name argument in column spec");

                GeneratorConfig config = new GeneratorConfig(seedStr + name, clustering, size, population);
                columnConfigs.put(name, config);
            }
        }
    }

    public void maybeCreateSchema(StressSettings settings)
    {
        JavaDriverClient client = settings.getJavaDriverClient(false);

        if (keyspaceCql != null)
        {
            try
            {
                client.execute(keyspaceCql, org.apache.cassandra.db.ConsistencyLevel.ONE);
            }
            catch (AlreadyExistsException e)
            {
            }
        }

        client.execute("use " + keyspaceName, org.apache.cassandra.db.ConsistencyLevel.ONE);

        if (tableCql != null)
        {
            try
            {
                client.execute(tableCql, org.apache.cassandra.db.ConsistencyLevel.ONE);
            }
            catch (AlreadyExistsException e)
            {
            }

            System.out.println(String.format("Created schema. Sleeping %ss for propagation.", settings.node.nodes.size()));
            Uninterruptibles.sleepUninterruptibly(settings.node.nodes.size(), TimeUnit.SECONDS);
        }

        if (extraSchemaDefinitions != null)
        {
            for (String extraCql : extraSchemaDefinitions)
            {

                try
                {
                    client.execute(extraCql, org.apache.cassandra.db.ConsistencyLevel.ONE);
                }
                catch (AlreadyExistsException e)
                {
                }
            }

            System.out.println(String.format("Created extra schema. Sleeping %ss for propagation.", settings.node.nodes.size()));
            Uninterruptibles.sleepUninterruptibly(settings.node.nodes.size(), TimeUnit.SECONDS);
        }

        maybeLoadSchemaInfo(settings);
    }

    public void truncateTable(StressSettings settings)
    {
        JavaDriverClient client = settings.getJavaDriverClient(false);
        assert settings.command.truncate != SettingsCommand.TruncateWhen.NEVER;
        String cql = String.format("TRUNCATE %s.%s", keyspaceName, tableName);
        client.execute(cql, org.apache.cassandra.db.ConsistencyLevel.ONE);
        System.out.println(String.format("Truncated %s.%s. Sleeping %ss for propagation.",
                                         keyspaceName, tableName, settings.node.nodes.size()));
        Uninterruptibles.sleepUninterruptibly(settings.node.nodes.size(), TimeUnit.SECONDS);
    }

    private void maybeLoadSchemaInfo(StressSettings settings)
    {
        if (tableMetaData == null)
        {
            JavaDriverClient client = settings.getJavaDriverClient();

            synchronized (client)
            {

                if (tableMetaData != null)
                    return;

                TableMetadata metadata = client.getCluster()
                                               .getMetadata()
                                               .getKeyspace(keyspaceName)
                                               .getTable(quoteIdentifier(tableName));

                if (metadata == null)
                    throw new RuntimeException("Unable to find table " + keyspaceName + "." + tableName);

                //Fill in missing column configs
                for (ColumnMetadata col : metadata.getColumns())
                {
                    if (columnConfigs.containsKey(col.getName()))
                        continue;

                    columnConfigs.put(col.getName(), new GeneratorConfig(seedStr + col.getName(), null, null, null));
                }

                tableMetaData = metadata;
            }
        }
    }

    public Set<TokenRange> maybeLoadTokenRanges(StressSettings settings)
    {
        maybeLoadSchemaInfo(settings); // ensure table metadata is available

        JavaDriverClient client = settings.getJavaDriverClient();
        synchronized (client)
        {
            if (tokenRanges != null)
                return tokenRanges;

            Cluster cluster = client.getCluster();
            Metadata metadata = cluster.getMetadata();
            if (metadata == null)
                throw new RuntimeException("Unable to get metadata");

            List<TokenRange> sortedRanges = new ArrayList<>(metadata.getTokenRanges().size() + 1);
            for (TokenRange range : metadata.getTokenRanges())
            {
                //if we don't unwrap we miss the partitions between ring min and smallest range start value
                if (range.isWrappedAround())
                    sortedRanges.addAll(range.unwrap());
                else
                    sortedRanges.add(range);
            }

            Collections.sort(sortedRanges);
            tokenRanges = new LinkedHashSet<>(sortedRanges);
            return tokenRanges;
        }
    }

    public Operation getQuery(String name,
                              Timer timer,
                              PartitionGenerator generator,
                              SeedManager seeds,
                              StressSettings settings,
                              boolean isWarmup)
    {
        name = name.toLowerCase();
        if (!queries.containsKey(name))
            throw new IllegalArgumentException("No query defined with name " + name);

        if (queryStatements == null)
        {
            synchronized (this)
            {
                if (queryStatements == null)
                {
                    try
                    {
                        JavaDriverClient jclient = settings.getJavaDriverClient();
                        ThriftClient tclient = null;

                        if (settings.mode.api != ConnectionAPI.JAVA_DRIVER_NATIVE)
                            tclient = settings.getThriftClient();

                        Map<String, PreparedStatement> stmts = new HashMap<>();
                        Map<String, Integer> tids = new HashMap<>();
                        Map<String, SchemaQuery.ArgSelect> args = new HashMap<>();
                        for (Map.Entry<String, StressYaml.QueryDef> e : queries.entrySet())
                        {
                            stmts.put(e.getKey().toLowerCase(), jclient.prepare(e.getValue().cql));

                            if (tclient != null)
                                tids.put(e.getKey().toLowerCase(), tclient.prepare_cql3_query(e.getValue().cql, Compression.NONE));

                            args.put(e.getKey().toLowerCase(), e.getValue().fields == null
                                                                     ? SchemaQuery.ArgSelect.MULTIROW
                                                                     : SchemaQuery.ArgSelect.valueOf(e.getValue().fields.toUpperCase()));
                        }
                        thriftQueryIds = tids;
                        queryStatements = stmts;
                        argSelects = args;
                    }
                    catch (TException e)
                    {
                        throw new RuntimeException(e);
                    }
                }
            }
        }

        return new SchemaQuery(timer, settings, generator, seeds, thriftQueryIds.get(name), queryStatements.get(name),
                               ThriftConversion.fromThrift(settings.command.consistencyLevel), argSelects.get(name));
    }

    public Operation getBulkReadQueries(String name, Timer timer, StressSettings settings, TokenRangeIterator tokenRangeIterator, boolean isWarmup)
    {
        StressYaml.TokenRangeQueryDef def = tokenRangeQueries.get(name);
        if (def == null)
            throw new IllegalArgumentException("No bulk read query defined with name " + name);

        return new TokenRangeQuery(timer, settings, tableMetaData, tokenRangeIterator, def, isWarmup);
    }


    public PartitionGenerator getOfflineGenerator()
    {
        CFMetaData cfMetaData = CFMetaData.compile(tableCql, keyspaceName);

        //Add missing column configs
        Iterator<ColumnDefinition> it = cfMetaData.allColumnsInSelectOrder();
        while (it.hasNext())
        {
            ColumnDefinition c = it.next();
            if (!columnConfigs.containsKey(c.name.toString()))
                columnConfigs.put(c.name.toString(), new GeneratorConfig(seedStr + c.name.toString(), null, null, null));
        }

        List<Generator> partitionColumns = cfMetaData.partitionKeyColumns().stream()
                                                     .map(c -> new ColumnInfo(c.name.toString(), c.type.asCQL3Type().toString(), "", columnConfigs.get(c.name.toString())))
                                                     .map(c -> c.getGenerator())
                                                     .collect(Collectors.toList());

        List<Generator> clusteringColumns = cfMetaData.clusteringColumns().stream()
                                                             .map(c -> new ColumnInfo(c.name.toString(), c.type.asCQL3Type().toString(), "", columnConfigs.get(c.name.toString())))
                                                             .map(c -> c.getGenerator())
                                                             .collect(Collectors.toList());

        List<Generator> regularColumns = com.google.common.collect.Lists.newArrayList(cfMetaData.partitionColumns().selectOrderIterator()).stream()
                                                                                                             .map(c -> new ColumnInfo(c.name.toString(), c.type.asCQL3Type().toString(), "", columnConfigs.get(c.name.toString())))
                                                                                                             .map(c -> c.getGenerator())
                                                                                                             .collect(Collectors.toList());

        return new PartitionGenerator(partitionColumns, clusteringColumns, regularColumns, PartitionGenerator.Order.ARBITRARY);
    }

    public CreateTableStatement.RawStatement getCreateStatement()
    {
        CreateTableStatement.RawStatement createStatement = parseStatement(tableCql, CreateTableStatement.RawStatement.class, "CREATE TABLE");
        createStatement.prepareKeyspace(keyspaceName);

        return createStatement;
    }

    public SchemaInsert getOfflineInsert(Timer timer, PartitionGenerator generator, SeedManager seedManager, StressSettings settings)
    {
        assert tableCql != null;

        CFMetaData cfMetaData = CFMetaData.compile(tableCql, keyspaceName);

        List<ColumnDefinition> allColumns = com.google.common.collect.Lists.newArrayList(cfMetaData.allColumnsInSelectOrder());

        StringBuilder sb = new StringBuilder();
        sb.append("INSERT INTO ").append(quoteIdentifier(keyspaceName) + "." + quoteIdentifier(tableName)).append(" (");
        StringBuilder value = new StringBuilder();
        for (ColumnDefinition c : allColumns)
        {
            sb.append(quoteIdentifier(c.name.toString())).append(", ");
            value.append("?, ");
        }
        sb.delete(sb.lastIndexOf(","), sb.length());
        value.delete(value.lastIndexOf(","), value.length());
        sb.append(") ").append("values(").append(value).append(')');


        if (insert == null)
            insert = new HashMap<>();
        lowerCase(insert);

        partitions = select(settings.insert.batchsize, "partitions", "fixed(1)", insert, OptionDistribution.BUILDER);
        selectchance = select(settings.insert.selectRatio, "select", "fixed(1)/1", insert, OptionRatioDistribution.BUILDER);
        rowPopulation = select(settings.insert.rowPopulationRatio, "row-population", "fixed(1)/1", insert, OptionRatioDistribution.BUILDER);

        if (generator.maxRowCount > 100 * 1000 * 1000)
            System.err.printf("WARNING: You have defined a schema that permits very large partitions (%.0f max rows (>100M))%n", generator.maxRowCount);

        String statement = sb.toString();

        //CQLTableWriter requires the keyspace name be in the create statement
        String tableCreate = tableCql.replaceFirst("\\s+\"?"+tableName+"\"?\\s+", " \""+keyspaceName+"\".\""+tableName+"\" ");


        return new SchemaInsert(timer, settings, generator, seedManager, selectchance.get(), rowPopulation.get(), thriftInsertId, statement, tableCreate);
    }

    public SchemaInsert getInsert(Timer timer, PartitionGenerator generator, SeedManager seedManager, StressSettings settings)
    {
        if (insertStatement == null)
        {
            synchronized (this)
            {
                if (insertStatement == null)
                {
                    maybeLoadSchemaInfo(settings);

                    Set<ColumnMetadata> keyColumns = com.google.common.collect.Sets.newHashSet(tableMetaData.getPrimaryKey());
                    Set<ColumnMetadata> allColumns = com.google.common.collect.Sets.newHashSet(tableMetaData.getColumns());
                    boolean isKeyOnlyTable = (keyColumns.size() == allColumns.size());
                    //With compact storage
                    if (!isKeyOnlyTable && (keyColumns.size() == (allColumns.size() - 1)))
                    {
                        com.google.common.collect.Sets.SetView diff = com.google.common.collect.Sets.difference(allColumns, keyColumns);
                        for (Object obj : diff)
                        {
                            ColumnMetadata col = (ColumnMetadata)obj;
                            isKeyOnlyTable = col.getName().isEmpty();
                            break;
                        }
                    }

                    //Non PK Columns
                    StringBuilder sb = new StringBuilder();
                    if (!isKeyOnlyTable)
                    {
                        sb.append("UPDATE ").append(quoteIdentifier(tableName)).append(" SET ");
                        //PK Columns
                        StringBuilder pred = new StringBuilder();
                        pred.append(" WHERE ");

                        boolean firstCol = true;
                        boolean firstPred = true;
                        for (ColumnMetadata c : tableMetaData.getColumns()) {

                            if (keyColumns.contains(c)) {
                                if (firstPred)
                                    firstPred = false;
                                else
                                    pred.append(" AND ");

                                pred.append(quoteIdentifier(c.getName())).append(" = ?");
                            } else {
                                if (firstCol)
                                    firstCol = false;
                                else
                                    sb.append(',');

                                sb.append(quoteIdentifier(c.getName())).append(" = ");

                                switch (c.getType().getName())
                                {
                                case SET:
                                case LIST:
                                case COUNTER:
                                    sb.append(quoteIdentifier(c.getName())).append(" + ?");
                                    break;
                                default:
                                    sb.append("?");
                                    break;
                                }
                            }
                        }

                        //Put PK predicates at the end
                        sb.append(pred);
                    }
                    else
                    {
                        sb.append("INSERT INTO ").append(quoteIdentifier(tableName)).append(" (");
                        StringBuilder value = new StringBuilder();
                        for (ColumnMetadata c : tableMetaData.getPrimaryKey())
                        {
                            sb.append(quoteIdentifier(c.getName())).append(", ");
                            value.append("?, ");
                        }
                        sb.delete(sb.lastIndexOf(","), sb.length());
                        value.delete(value.lastIndexOf(","), value.length());
                        sb.append(") ").append("values(").append(value).append(')');
                    }

                    if (insert == null)
                        insert = new HashMap<>();
                    lowerCase(insert);

                    partitions = select(settings.insert.batchsize, "partitions", "fixed(1)", insert, OptionDistribution.BUILDER);
                    selectchance = select(settings.insert.selectRatio, "select", "fixed(1)/1", insert, OptionRatioDistribution.BUILDER);
                    rowPopulation = select(settings.insert.rowPopulationRatio, "row-population", "fixed(1)/1", insert, OptionRatioDistribution.BUILDER);
                    batchType = settings.insert.batchType != null
                                ? settings.insert.batchType
                                : !insert.containsKey("batchtype")
                                  ? BatchStatement.Type.LOGGED
                                  : BatchStatement.Type.valueOf(insert.remove("batchtype"));
                    if (!insert.isEmpty())
                        throw new IllegalArgumentException("Unrecognised insert option(s): " + insert);

                    Distribution visits = settings.insert.visits.get();
                    // these min/max are not absolutely accurate if selectchance < 1, but they're close enough to
                    // guarantee the vast majority of actions occur in these bounds
                    double minBatchSize = selectchance.get().min() * partitions.get().minValue() * generator.minRowCount * (1d / visits.maxValue());
                    double maxBatchSize = selectchance.get().max() * partitions.get().maxValue() * generator.maxRowCount * (1d / visits.minValue());
                    System.out.printf("Generating batches with [%d..%d] partitions and [%.0f..%.0f] rows (of [%.0f..%.0f] total rows in the partitions)%n",
                                      partitions.get().minValue(), partitions.get().maxValue(),
                                      minBatchSize, maxBatchSize,
                                      partitions.get().minValue() * generator.minRowCount,
                                      partitions.get().maxValue() * generator.maxRowCount);
                    if (generator.maxRowCount > 100 * 1000 * 1000)
                        System.err.printf("WARNING: You have defined a schema that permits very large partitions (%.0f max rows (>100M))%n", generator.maxRowCount);
                    if (batchType == BatchStatement.Type.LOGGED && maxBatchSize > 65535)
                    {
                        System.err.printf("ERROR: You have defined a workload that generates batches with more than 65k rows (%.0f), but have required the use of LOGGED batches. There is a 65k row limit on a single batch.%n",
                                          selectchance.get().max() * partitions.get().maxValue() * generator.maxRowCount);
                        System.exit(1);
                    }
                    if (maxBatchSize > 100000)
                        System.err.printf("WARNING: You have defined a schema that permits very large batches (%.0f max rows (>100K)). This may OOM this stress client, or the server.%n",
                                          selectchance.get().max() * partitions.get().maxValue() * generator.maxRowCount);

                    JavaDriverClient client = settings.getJavaDriverClient();
                    String query = sb.toString();

                    if (settings.mode.api != ConnectionAPI.JAVA_DRIVER_NATIVE)
                    {
                        try
                        {
                            thriftInsertId = settings.getThriftClient().prepare_cql3_query(query, Compression.NONE);
                        }
                        catch (TException e)
                        {
                            throw new RuntimeException(e);
                        }
                    }

                    insertStatement = client.prepare(query);
                }
            }
        }

        return new SchemaInsert(timer, settings, generator, seedManager, partitions.get(), selectchance.get(), rowPopulation.get(), thriftInsertId, insertStatement, ThriftConversion.fromThrift(settings.command.consistencyLevel), batchType);
    }

    public List<ValidatingSchemaQuery> getValidate(Timer timer, PartitionGenerator generator, SeedManager seedManager, StressSettings settings)
    {
        if (validationFactories == null)
        {
            synchronized (this)
            {
                if (validationFactories == null)
                {
                    maybeLoadSchemaInfo(settings);
                    validationFactories = ValidatingSchemaQuery.create(tableMetaData, settings);
                }
            }
        }

        List<ValidatingSchemaQuery> queries = new ArrayList<>();
        for (ValidatingSchemaQuery.Factory factory : validationFactories)
            queries.add(factory.create(timer, settings, generator, seedManager, ThriftConversion.fromThrift(settings.command.consistencyLevel)));
        return queries;
    }

    private static <E> E select(E first, String key, String defValue, Map<String, String> map, Function<String, E> builder)
    {
        String val = map.remove(key);

        if (first != null)
            return first;
        if (val != null && val.trim().length() > 0)
            return builder.apply(val);
        
        return builder.apply(defValue);
    }

    public PartitionGenerator newGenerator(StressSettings settings)
    {
        if (generatorFactory == null)
        {
            synchronized (this)
            {
                maybeLoadSchemaInfo(settings);
                if (generatorFactory == null)
                    generatorFactory = new GeneratorFactory();
            }
        }

        return generatorFactory.newGenerator(settings);
    }

    private class GeneratorFactory
    {
        final List<ColumnInfo> partitionKeys = new ArrayList<>();
        final List<ColumnInfo> clusteringColumns = new ArrayList<>();
        final List<ColumnInfo> valueColumns = new ArrayList<>();

        private GeneratorFactory()
        {
            Set<ColumnMetadata> keyColumns = com.google.common.collect.Sets.newHashSet(tableMetaData.getPrimaryKey());

            for (ColumnMetadata metadata : tableMetaData.getPartitionKey())
                partitionKeys.add(new ColumnInfo(metadata.getName(), metadata.getType().getName().toString(),
                                                 metadata.getType().isCollection() ? metadata.getType().getTypeArguments().get(0).getName().toString() : "",
                                                 columnConfigs.get(metadata.getName())));
            for (ColumnMetadata metadata : tableMetaData.getClusteringColumns())
                clusteringColumns.add(new ColumnInfo(metadata.getName(), metadata.getType().getName().toString(),
                                                     metadata.getType().isCollection() ? metadata.getType().getTypeArguments().get(0).getName().toString() : "",
                                                     columnConfigs.get(metadata.getName())));
            for (ColumnMetadata metadata : tableMetaData.getColumns())
                if (!keyColumns.contains(metadata))
                    valueColumns.add(new ColumnInfo(metadata.getName(), metadata.getType().getName().toString(),
                                                    metadata.getType().isCollection() ? metadata.getType().getTypeArguments().get(0).getName().toString() : "",
                                                    columnConfigs.get(metadata.getName())));
        }

        PartitionGenerator newGenerator(StressSettings settings)
        {
            return new PartitionGenerator(get(partitionKeys), get(clusteringColumns), get(valueColumns), settings.generate.order);
        }

        List<Generator> get(List<ColumnInfo> columnInfos)
        {
            List<Generator> result = new ArrayList<>();
            for (ColumnInfo columnInfo : columnInfos)
                result.add(columnInfo.getGenerator());
            return result;
        }
    }

    static class ColumnInfo
    {
        final String name;
        final String type;
        final String collectionType;
        final GeneratorConfig config;

        ColumnInfo(String name, String type, String collectionType, GeneratorConfig config)
        {
            this.name = name;
            this.type = type;
            this.collectionType = collectionType;
            this.config = config;
        }

        Generator getGenerator()
        {
            return getGenerator(name, type, collectionType, config);
        }

        static Generator getGenerator(final String name, final String type, final String collectionType, GeneratorConfig config)
        {
            switch (type.toUpperCase())
            {
                case "ASCII":
                case "TEXT":
                case "VARCHAR":
                    return new Strings(name, config);
                case "BIGINT":
                case "COUNTER":
                    return new Longs(name, config);
                case "BLOB":
                    return new Bytes(name, config);
                case "BOOLEAN":
                    return new Booleans(name, config);
                case "DECIMAL":
                    return new BigDecimals(name, config);
                case "DOUBLE":
                    return new Doubles(name, config);
                case "FLOAT":
                    return new Floats(name, config);
                case "INET":
                    return new Inets(name, config);
                case "INT":
                    return new Integers(name, config);
                case "VARINT":
                    return new BigIntegers(name, config);
                case "TIMESTAMP":
                    return new Dates(name, config);
                case "UUID":
                    return new UUIDs(name, config);
                case "TIMEUUID":
                    return new TimeUUIDs(name, config);
                case "TINYINT":
                    return new TinyInts(name, config);
                case "SMALLINT":
                    return new SmallInts(name, config);
                case "TIME":
                    return new Times(name, config);
                case "DATE":
                    return new LocalDates(name, config);
                case "SET":
                    return new Sets(name, getGenerator(name, collectionType, null, config), config);
                case "LIST":
                    return new Lists(name, getGenerator(name, collectionType, null, config), config);
                default:
                    throw new UnsupportedOperationException("Because of this name: "+name+" if you removed it from the yaml and are still seeing this, make sure to drop table");
            }
        }
    }

    public static StressProfile load(URI file) throws IOError
    {
        try
        {
            Constructor constructor = new Constructor(StressYaml.class);

            Yaml yaml = new Yaml(constructor);

            InputStream yamlStream = file.toURL().openStream();

            if (yamlStream.available() == 0)
                throw new IOException("Unable to load yaml file from: "+file);

            StressYaml profileYaml = yaml.loadAs(yamlStream, StressYaml.class);

            StressProfile profile = new StressProfile();
            profile.init(profileYaml);

            return profile;
        }
        catch (YAMLException | IOException | RequestValidationException e)
        {
            throw new IOError(e);
        }
    }

    static <V> void lowerCase(Map<String, V> map)
    {
        List<Map.Entry<String, V>> reinsert = new ArrayList<>();
        Iterator<Map.Entry<String, V>> iter = map.entrySet().iterator();
        while (iter.hasNext())
        {
            Map.Entry<String, V> e = iter.next();
            if (!e.getKey().equalsIgnoreCase(e.getKey()))
            {
                reinsert.add(e);
                iter.remove();
            }
        }
        for (Map.Entry<String, V> e : reinsert)
            map.put(e.getKey().toLowerCase(), e.getValue());
    }

    /* Quote a identifier if it contains uppercase letters */
    private static String quoteIdentifier(String identifier)
    {
        return lowercaseAlphanumeric.matcher(identifier).matches() ? identifier : '\"'+identifier+ '\"';
    }
}