/*
 * Copyright DataStax, Inc.
 *
 * 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.
 */
package com.datastax.driver.core.querybuilder;

import com.datastax.driver.core.*;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;

/**
 * A built BATCH statement.
 */
public class Batch extends BuiltStatement {

    private final List<RegularStatement> statements;
    private final boolean logged;
    private final Options usings;

    // Only used when we add at last one statement that is not a BuiltStatement subclass
    private int nonBuiltStatementValues;

    Batch(RegularStatement[] statements, boolean logged) {
        super(null, null, null);
        this.statements = statements.length == 0
                ? new ArrayList<RegularStatement>()
                : new ArrayList<RegularStatement>(statements.length);
        this.logged = logged;
        this.usings = new Options(this);

        for (int i = 0; i < statements.length; i++)
            add(statements[i]);
    }

    @Override
    StringBuilder buildQueryString(List<Object> variables, CodecRegistry codecRegistry) {
        StringBuilder builder = new StringBuilder();

        builder.append(isCounterOp()
                ? "BEGIN COUNTER BATCH"
                : (logged ? "BEGIN BATCH" : "BEGIN UNLOGGED BATCH"));

        if (!usings.usings.isEmpty()) {
            builder.append(" USING ");
            Utils.joinAndAppend(builder, codecRegistry, " AND ", usings.usings, variables);
        }
        builder.append(' ');

        for (int i = 0; i < statements.size(); i++) {
            RegularStatement stmt = statements.get(i);
            if (stmt instanceof BuiltStatement) {
                BuiltStatement bst = (BuiltStatement) stmt;
                builder.append(maybeAddSemicolon(bst.buildQueryString(variables, codecRegistry)));

            } else {
                String str = stmt.getQueryString(codecRegistry);
                builder.append(str);
                if (!str.trim().endsWith(";"))
                    builder.append(';');

                // Note that we force hasBindMarkers if there is any non-BuiltStatement, so we know
                // that we can only get there with variables == null
                assert variables == null;
            }
        }
        builder.append("APPLY BATCH;");
        return builder;
    }

    /**
     * Adds a new statement to this batch.
     *
     * @param statement the new statement to add.
     * @return this batch.
     * @throws IllegalArgumentException if counter and non-counter operations
     *                                  are mixed.
     */
    public Batch add(RegularStatement statement) {
        boolean isCounterOp = statement instanceof BuiltStatement && ((BuiltStatement) statement).isCounterOp();

        if (this.isCounterOp == null)
            setCounterOp(isCounterOp);
        else if (isCounterOp() != isCounterOp)
            throw new IllegalArgumentException("Cannot mix counter operations and non-counter operations in a batch statement");

        this.statements.add(statement);

        if (statement instanceof BuiltStatement) {
            this.hasBindMarkers |= ((BuiltStatement) statement).hasBindMarkers;
        } else {
            // For non-BuiltStatement, we cannot know if it includes a bind makers and we assume it does. In practice,
            // this means we will always serialize values as strings when there is non-BuiltStatement
            this.hasBindMarkers = true;
            this.nonBuiltStatementValues += ((SimpleStatement) statement).valuesCount();
        }

        setDirty();

        return this;
    }

    @Override
    public ByteBuffer[] getValues(ProtocolVersion protocolVersion, CodecRegistry codecRegistry) {
        // If there is some non-BuiltStatement inside the batch with values, we shouldn't
        // use super.getValues() since it will ignore the values of said non-BuiltStatement.
        // If that's the case, we just collects all those values (and we know
        // super.getValues() == null in that case since we've explicitely set this.hasBindMarker
        // to true). Otherwise, we simply call super.getValues().
        if (nonBuiltStatementValues == 0)
            return super.getValues(protocolVersion, codecRegistry);

        ByteBuffer[] values = new ByteBuffer[nonBuiltStatementValues];
        int i = 0;
        for (RegularStatement statement : statements) {
            if (statement instanceof BuiltStatement)
                continue;

            ByteBuffer[] statementValues = statement.getValues(protocolVersion, codecRegistry);
            System.arraycopy(statementValues, 0, values, i, statementValues.length);
            i += statementValues.length;
        }
        return values;
    }

    /**
     * Adds a new options for this BATCH statement.
     *
     * @param using the option to add.
     * @return the options of this BATCH statement.
     */
    public Options using(Using using) {
        return usings.and(using);
    }

    /**
     * Returns the first non-null routing key of the statements in this batch
     * or null otherwise.
     *
     * @return the routing key for this batch statement.
     */
    @Override
    public ByteBuffer getRoutingKey(ProtocolVersion protocolVersion, CodecRegistry codecRegistry) {
        for (RegularStatement statement : statements) {
            ByteBuffer routingKey = statement.getRoutingKey(protocolVersion, codecRegistry);
            if (routingKey != null) {
                return routingKey;
            }
        }
        return null;
    }

    /**
     * Returns the keyspace of the first statement in this batch.
     *
     * @return the keyspace of the first statement in this batch.
     */
    @Override
    public String getKeyspace() {
        return statements.isEmpty() ? null : statements.get(0).getKeyspace();
    }

    @Override
    public Boolean isIdempotent() {
        if (idempotent != null) {
            return idempotent;
        }
        return isBatchIdempotent(statements);
    }

    /**
     * The options of a BATCH statement.
     */
    public static class Options extends BuiltStatement.ForwardingStatement<Batch> {

        private final List<Using> usings = new ArrayList<Using>();

        Options(Batch statement) {
            super(statement);
        }

        /**
         * Adds the provided option.
         *
         * @param using a BATCH option.
         * @return this {@code Options} object.
         */
        public Options and(Using using) {
            usings.add(using);
            checkForBindMarkers(using);
            return this;
        }

        /**
         * Adds a new statement to the BATCH statement these options are part of.
         *
         * @param statement the statement to add.
         * @return the BATCH statement these options are part of.
         */
        public Batch add(RegularStatement statement) {
            return this.statement.add(statement);
        }
    }
}
