/*
 * 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.operations.predefined;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import com.datastax.driver.core.PreparedStatement;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.Row;
import com.google.common.base.Function;

import org.apache.cassandra.stress.generate.PartitionGenerator;
import org.apache.cassandra.stress.generate.SeedManager;
import org.apache.cassandra.stress.report.Timer;
import org.apache.cassandra.stress.settings.Command;
import org.apache.cassandra.stress.settings.ConnectionStyle;
import org.apache.cassandra.stress.settings.StressSettings;
import org.apache.cassandra.stress.util.JavaDriverClient;
import org.apache.cassandra.transport.SimpleClient;
import org.apache.cassandra.transport.messages.ResultMessage;
import org.apache.cassandra.utils.ByteBufferUtil;

public abstract class CqlOperation<V> extends PredefinedOperation
{

    public static final ByteBuffer[][] EMPTY_BYTE_BUFFERS = new ByteBuffer[0][];
    public static final byte[][] EMPTY_BYTE_ARRAYS = new byte[0][];

    protected abstract List<Object> getQueryParameters(byte[] key);
    protected abstract String buildQuery();
    protected abstract CqlRunOp<V> buildRunOp(QueryExecutor<?> queryExecutor, List<Object> params, ByteBuffer key);

    public CqlOperation(Command type, Timer timer, PartitionGenerator generator, SeedManager seedManager, StressSettings settings)
    {
        super(type, timer, generator, seedManager, settings);
        if (settings.columns.variableColumnCount)
            throw new IllegalStateException("Variable column counts are not implemented for CQL");
    }

    protected CqlRunOp<V> run(final QueryExecutor<?> queryExecutor, final List<Object> queryParams, final ByteBuffer key) throws IOException
    {
        final CqlRunOp<V> op = buildRunOp(queryExecutor, queryParams, key);
        timeWithRetry(op);
        return op;
    }

    protected void run(final QueryExecutor<?> queryExecutor) throws IOException
    {
        final byte[] key = getKey().array();
        final List<Object> queryParams = getQueryParameters(key);
        run(queryExecutor, queryParams, ByteBuffer.wrap(key));
    }

    // Classes to process Cql results

    // Always succeeds so long as the query executes without error; provides a keyCount to increment on instantiation
    protected final class CqlRunOpAlwaysSucceed extends CqlRunOp<Integer>
    {

        final int keyCount;

        protected CqlRunOpAlwaysSucceed(QueryExecutor<?> queryExecutor, List<Object> params, ByteBuffer key, int keyCount)
        {
            super(queryExecutor, RowCountHandler.INSTANCE, params, key);
            this.keyCount = keyCount;
        }

        @Override
        public boolean validate(Integer result)
        {
            return true;
        }

        @Override
        public int partitionCount()
        {
            return keyCount;
        }

        @Override
        public int rowCount()
        {
            return keyCount;
        }
    }

    // Succeeds so long as the result set is nonempty, and the query executes without error
    protected final class CqlRunOpTestNonEmpty extends CqlRunOp<Integer>
    {

        protected CqlRunOpTestNonEmpty(QueryExecutor<?> queryExecutor, List<Object> params, ByteBuffer key)
        {
            super(queryExecutor, RowCountHandler.INSTANCE, params, key);
        }

        @Override
        public boolean validate(Integer result)
        {
            return result > 0;
        }

        @Override
        public int partitionCount()
        {
            return result;
        }

        @Override
        public int rowCount()
        {
            return result;
        }
    }

    protected final class CqlRunOpMatchResults extends CqlRunOp<ByteBuffer[][]>
    {

        final List<List<ByteBuffer>> expect;

        // a null value for an item in expect means we just check the row is present
        protected CqlRunOpMatchResults(QueryExecutor<?> queryExecutor, List<Object> params, ByteBuffer key, List<List<ByteBuffer>> expect)
        {
            super(queryExecutor, RowsHandler.INSTANCE, params, key);
            this.expect = expect;
        }

        @Override
        public int partitionCount()
        {
            return result == null ? 0 : result.length;
        }

        @Override
        public int rowCount()
        {
            return result == null ? 0 : result.length;
        }

        public boolean validate(ByteBuffer[][] result)
        {
            if (!settings.errors.skipReadValidation)
            {
                if (result.length != expect.size())
                    return false;
                for (int i = 0; i < result.length; i++)
                    if (expect.get(i) != null && !expect.get(i).equals(Arrays.asList(result[i])))
                        return false;
            }
            return true;
        }
    }

    // Cql
    protected abstract class CqlRunOp<V> implements RunOp
    {

        final QueryExecutor<?> queryExecutor;
        final List<Object> params;
        final ByteBuffer key;
        final ResultHandler<V> handler;
        V result;

        private CqlRunOp(QueryExecutor<?> queryExecutor, ResultHandler<V> handler, List<Object> params, ByteBuffer key)
        {
            this.queryExecutor = queryExecutor;
            this.handler = handler;
            this.params = params;
            this.key = key;
        }

        @Override
        public boolean run() throws Exception
        {
            return validate(result = queryExecutor.execute(key, params, handler));
        }

        public abstract boolean validate(V result);

    }


    /// LOTS OF WRAPPING/UNWRAPPING NONSENSE


    @Override
    public void run(SimpleClient client) throws IOException
    {
        run(new SimpleClientQueryExecutor(client, buildQuery()));
    }

    @Override
    public void run(JavaDriverClient client) throws IOException
    {
        run(new JavaDriverQueryExecutor(client, buildQuery()));
    }

    protected abstract class QueryExecutor<PS>
    {
        private final String query;
        private PS preparedStatement;

        public QueryExecutor(String query)
        {
            this.query = query;
        }

        private boolean isPrepared()
        {
            return settings.mode.style == ConnectionStyle.CQL_PREPARED;
        }

        abstract protected PS createPreparedStatement(String query);

        private PS getPreparedStatement()
        {
            if (preparedStatement == null)
            {
                preparedStatement = createPreparedStatement(this.query);
            }
            return preparedStatement;
        }

        /**
         * Constructs a CQL query string by replacing instances of the character
         * '?', with the corresponding parameter.
         *
         * @param query base query string to format
         * @param parms sequence of string query parameters
         * @return formatted CQL query string
         */
        private String formatCqlQuery(String query, List<Object> parms)
        {
            int marker, position = 0;
            StringBuilder result = new StringBuilder();

            if (-1 == (marker = query.indexOf('?')) || parms.size() == 0)
                return query;

            for (Object parm : parms)
            {
                result.append(query.substring(position, marker));

                if (parm instanceof ByteBuffer)
                    result.append(getUnQuotedCqlBlob((ByteBuffer) parm));
                else if (parm instanceof Long)
                    result.append(parm);
                else throw new AssertionError();

                position = marker + 1;
                if (-1 == (marker = query.indexOf('?', position + 1)))
                    break;
            }

            if (position < query.length())
                result.append(query.substring(position));

            return result.toString();
        }

        <V> V execute(ByteBuffer key, List<Object> queryParams, ResultHandler<V> handler)
        {
            return isPrepared()
                   ? execute(getPreparedStatement(), key, queryParams, handler)
                   : execute(formatCqlQuery(this.query, queryParams), key, handler);
        }

        abstract <V> V execute(PS preparedStatement, ByteBuffer key, List<Object> queryParams, ResultHandler<V> handler);
        abstract <V> V execute(String formattedQuery, ByteBuffer key, ResultHandler<V> handler);
    }

    private final class JavaDriverQueryExecutor extends QueryExecutor<PreparedStatement>
    {
        final JavaDriverClient client;
        private JavaDriverQueryExecutor(JavaDriverClient client, String query)
        {
            super(query);
            this.client = client;
        }

        protected <V> V execute(String formattedQuery, ByteBuffer key, ResultHandler<V> handler)
        {
            return handler.javaDriverHandler().apply(client.execute(formattedQuery, settings.command.consistencyLevel));
        }

        protected <V> V execute(PreparedStatement preparedStatement, ByteBuffer key, List<Object> queryParams, ResultHandler<V> handler)
        {
            return handler.javaDriverHandler().apply(
                    client.executePrepared(
                            preparedStatement,
                            queryParams,
                            settings.command.consistencyLevel));
        }

        public PreparedStatement createPreparedStatement(String cqlQuery)
        {
            return client.prepare(cqlQuery);
        }
    }

    private final class SimpleClientQueryExecutor extends QueryExecutor<ResultMessage.Prepared>
    {
        final SimpleClient client;
        private SimpleClientQueryExecutor(SimpleClient client, String query)
        {
            super(query);
            this.client = client;
        }

        public <V> V execute(String formattedQuery, ByteBuffer key, ResultHandler<V> handler)
        {
            return handler.simpleClientHandler().apply(client.execute(formattedQuery, settings.command.consistencyLevel));
        }

        public <V> V execute(ResultMessage.Prepared preparedStatement, ByteBuffer key, List<Object> queryParams, ResultHandler<V> handler)
        {
            return handler.simpleClientHandler().apply(
                    client.executePrepared(
                            preparedStatement,
                            toByteBufferParams(queryParams),
                            settings.command.consistencyLevel));
        }

        public ResultMessage.Prepared createPreparedStatement(String cqlQuery)
        {
            return client.prepare(cqlQuery);
        }
    }


    // interface for building functions to standardise results from each client
    protected static interface ResultHandler<V>
    {
        Function<ResultSet, V> javaDriverHandler();
        Function<ResultMessage, V> simpleClientHandler();
    }

    protected static class RowCountHandler implements ResultHandler<Integer>
    {
        static final RowCountHandler INSTANCE = new RowCountHandler();

        @Override
        public Function<ResultSet, Integer> javaDriverHandler()
        {
            return new Function<ResultSet, Integer>()
            {
                @Override
                public Integer apply(ResultSet rows)
                {
                    if (rows == null)
                        return 0;
                    return rows.all().size();
                }
            };
        }

        @Override
        public Function<ResultMessage, Integer> simpleClientHandler()
        {
            return new Function<ResultMessage, Integer>()
            {
                @Override
                public Integer apply(ResultMessage result)
                {
                    return result instanceof ResultMessage.Rows ? ((ResultMessage.Rows) result).result.size() : 0;
                }
            };
        }
    }

    // Processes results from each client into an array of all key bytes returned
    protected static final class RowsHandler implements ResultHandler<ByteBuffer[][]>
    {
        static final RowsHandler INSTANCE = new RowsHandler();

        @Override
        public Function<ResultSet, ByteBuffer[][]> javaDriverHandler()
        {
            return new Function<ResultSet, ByteBuffer[][]>()
            {

                @Override
                public ByteBuffer[][] apply(ResultSet result)
                {
                    if (result == null)
                        return EMPTY_BYTE_BUFFERS;
                    List<Row> rows = result.all();

                    ByteBuffer[][] r = new ByteBuffer[rows.size()][];
                    for (int i = 0 ; i < r.length ; i++)
                    {
                        Row row = rows.get(i);
                        r[i] = new ByteBuffer[row.getColumnDefinitions().size()];
                        for (int j = 0 ; j < row.getColumnDefinitions().size() ; j++)
                            r[i][j] = row.getBytes(j);
                    }
                    return r;
                }
            };
        }

        @Override
        public Function<ResultMessage, ByteBuffer[][]> simpleClientHandler()
        {
            return new Function<ResultMessage, ByteBuffer[][]>()
            {

                @Override
                public ByteBuffer[][] apply(ResultMessage result)
                {
                    if (!(result instanceof ResultMessage.Rows))
                        return EMPTY_BYTE_BUFFERS;

                    ResultMessage.Rows rows = ((ResultMessage.Rows) result);
                    ByteBuffer[][] r = new ByteBuffer[rows.result.size()][];
                    for (int i = 0 ; i < r.length ; i++)
                    {
                        List<byte[]> row = rows.result.rows.get(i);
                        r[i] = new ByteBuffer[row.size()];
                        for (int j = 0 ; j < row.size() ; j++)
                        {
                            byte[] v = row.get(j);
                            r[i][j] = v == null ? null : ByteBuffer.wrap(v);
                        }
                    }
                    return r;
                }
            };
        }
    }
    // Processes results from each client into an array of all key bytes returned
    protected static final class KeysHandler implements ResultHandler<byte[][]>
    {
        static final KeysHandler INSTANCE = new KeysHandler();

        @Override
        public Function<ResultSet, byte[][]> javaDriverHandler()
        {
            return new Function<ResultSet, byte[][]>()
            {

                @Override
                public byte[][] apply(ResultSet result)
                {

                    if (result == null)
                        return EMPTY_BYTE_ARRAYS;
                    List<Row> rows = result.all();
                    byte[][] r = new byte[rows.size()][];
                    for (int i = 0 ; i < r.length ; i++)
                        r[i] = rows.get(i).getBytes(0).array();
                    return r;
                }
            };
        }

        @Override
        public Function<ResultMessage, byte[][]> simpleClientHandler()
        {
            return new Function<ResultMessage, byte[][]>()
            {

                @Override
                public byte[][] apply(ResultMessage result)
                {
                    if (result instanceof ResultMessage.Rows)
                    {
                        ResultMessage.Rows rows = ((ResultMessage.Rows) result);
                        byte[][] r = new byte[rows.result.size()][];
                        for (int i = 0 ; i < r.length ; i++)
                            r[i] = rows.result.rows.get(i).get(0);
                        return r;
                    }
                    return null;
                }
            };
        }
    }

    private static String getUnQuotedCqlBlob(ByteBuffer term)
    {
        return "0x" + ByteBufferUtil.bytesToHex(term);
    }

    private static List<ByteBuffer> toByteBufferParams(List<Object> params)
    {
        List<ByteBuffer> r = new ArrayList<>();
        for (Object param : params)
        {
            if (param instanceof ByteBuffer)
                r.add((ByteBuffer) param);
            else if (param instanceof Long)
                r.add(ByteBufferUtil.bytes((Long) param));
            else throw new AssertionError();
        }
        return r;
    }

    protected String wrapInQuotes(String string)
    {
        return "\"" + string + "\"";
    }

}
