/*
 * 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 com.datastax.driver.core;

import com.datastax.driver.core.exceptions.InvalidTypeException;
import java.nio.ByteBuffer;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/** A simple {@code RegularStatement} implementation built directly from a query string. */
public class SimpleStatement extends RegularStatement {

  private final String query;
  private final Object[] values;
  private final Map<String, Object> namedValues;

  private volatile ByteBuffer routingKey;
  private volatile String keyspace;

  /**
   * Creates a new {@code SimpleStatement} with the provided query string (and no values).
   *
   * @param query the query string.
   */
  public SimpleStatement(String query) {
    this(query, (Object[]) null);
  }

  /**
   * Creates a new {@code SimpleStatement} with the provided query string and values.
   *
   * <p>This version of SimpleStatement is useful when you want to execute a query only once (and
   * thus do not want to resort to prepared statement), but do not want to convert all column values
   * to string (typically, if you have blob values, encoding them to a hexadecimal string is not
   * very efficient). In that case, you can provide a query string with bind markers to this
   * constructor along with the values for those bind variables. When executed, the server will
   * prepare the provided, bind the provided values to that prepare statement and execute the
   * resulting statement. Thus,
   *
   * <pre>
   *   session.execute(new SimpleStatement(query, value1, value2, value3));
   * </pre>
   *
   * is functionally equivalent to
   *
   * <pre>
   *   PreparedStatement ps = session.prepare(query);
   *   session.execute(ps.bind(value1, value2, value3));
   * </pre>
   *
   * except that the former version:
   *
   * <ul>
   *   <li>Requires only one round-trip to a Cassandra node.
   *   <li>Does not left any prepared statement stored in memory (neither client or server side)
   *       once it has been executed.
   * </ul>
   *
   * <p>Note that the types of the {@code values} provided to this method will not be validated by
   * the driver as is done by {@link BoundStatement#bind} since {@code query} is not parsed (and
   * hence the driver cannot know what those values should be). The codec to serialize each value
   * will be chosen in the codec registry associated with the cluster executing this statement,
   * based on the value's Java type (this is the equivalent to calling {@link
   * CodecRegistry#codecFor(Object)}). If too many or too few values are provided, or if a value is
   * not a valid one for the variable it is bound to, an {@link
   * com.datastax.driver.core.exceptions.InvalidQueryException} will be thrown by Cassandra at
   * execution time. A {@code CodecNotFoundException} may be thrown by this constructor however, if
   * the codec registry does not know how to handle one of the values.
   *
   * <p>If you have a single value of type {@code Map<String, Object>}, you can't call this
   * constructor using the varargs syntax, because the signature collides with {@link
   * #SimpleStatement(String, Map)}. To prevent this, pass an explicit array object:
   *
   * <pre>
   * new SimpleStatement("...", new Object[]{m});
   * </pre>
   *
   * @param query the query string.
   * @param values values required for the execution of {@code query}.
   * @throws IllegalArgumentException if the number of values is greater than 65535.
   */
  public SimpleStatement(String query, Object... values) {
    if (values != null && values.length > 65535)
      throw new IllegalArgumentException("Too many values, the maximum allowed is 65535");
    this.query = query;
    this.values = values;
    this.namedValues = null;
  }

  /**
   * Creates a new {@code SimpleStatement} with the provided query string and named values.
   *
   * <p>This constructor requires that the query string use named placeholders, for example:
   *
   * <pre>{@code
   * new SimpleStatement("SELECT * FROM users WHERE id = :i", ImmutableMap.<String, Object>of("i", 1));
   * }</pre>
   *
   * Make sure that the map is correctly typed {@code Map<String, Object>}, otherwise you might
   * accidentally call {@link #SimpleStatement(String, Object...)} with a positional value of type
   * map.
   *
   * <p>The types of the values will be handled the same way as with anonymous placeholders (see
   * {@link #SimpleStatement(String, Object...)}).
   *
   * <p>Simple statements with named values are only supported starting with native protocol {@link
   * ProtocolVersion#V3 v3}. With earlier versions, an {@link
   * com.datastax.driver.core.exceptions.UnsupportedFeatureException} will be thrown at execution
   * time.
   *
   * @param query the query string.
   * @param values named values required for the execution of {@code query}.
   * @throws IllegalArgumentException if the number of values is greater than 65535.
   */
  public SimpleStatement(String query, Map<String, Object> values) {
    if (values.size() > 65535)
      throw new IllegalArgumentException("Too many values, the maximum allowed is 65535");
    this.query = query;
    this.values = null;
    this.namedValues = values;
  }

  @Override
  public String getQueryString(CodecRegistry codecRegistry) {
    return query;
  }

  @Override
  public ByteBuffer[] getValues(ProtocolVersion protocolVersion, CodecRegistry codecRegistry) {
    if (values == null) return null;
    return convert(values, protocolVersion, codecRegistry);
  }

  @Override
  public Map<String, ByteBuffer> getNamedValues(
      ProtocolVersion protocolVersion, CodecRegistry codecRegistry) {
    if (namedValues == null) return null;
    return convert(namedValues, protocolVersion, codecRegistry);
  }

  /**
   * The number of values for this statement, that is the size of the array that will be returned by
   * {@code getValues}.
   *
   * @return the number of values.
   */
  public int valuesCount() {
    if (values != null) return values.length;
    else if (namedValues != null) return namedValues.size();
    else return 0;
  }

  @Override
  public boolean hasValues(CodecRegistry codecRegistry) {
    return (values != null && values.length > 0) || (namedValues != null && namedValues.size() > 0);
  }

  @Override
  public boolean usesNamedValues() {
    return namedValues != null && namedValues.size() > 0;
  }

  /**
   * Returns the {@code i}th positional value as the Java type matching its CQL type.
   *
   * <p>Note that, unlike with other driver types like {@link Row}, you can't retrieve named values
   * by position. This getter will throw an exception if the statement was created with named values
   * (or no values at all). Call {@link #usesNamedValues()} to check the type of values, and {@link
   * #getObject(String)} if they are positional.
   *
   * @param i the index to retrieve.
   * @return the {@code i}th value of this statement.
   * @throws IllegalStateException if this statement does not have positional values.
   * @throws IndexOutOfBoundsException if {@code i} is not a valid index for this object.
   */
  public Object getObject(int i) {
    if (values == null)
      throw new IllegalStateException("This statement does not have positional values");
    if (i < 0 || i >= values.length) throw new ArrayIndexOutOfBoundsException(i);
    return values[i];
  }

  /**
   * Returns a named value as the Java type matching its CQL type.
   *
   * @param name the name of the value to retrieve.
   * @return the value that matches the name, or {@code null} if there is no such name.
   * @throws IllegalStateException if this statement does not have named values.
   */
  public Object getObject(String name) {
    if (namedValues == null)
      throw new IllegalStateException("This statement does not have named values");
    return namedValues.get(name);
  }

  /**
   * Returns the names of the named values of this statement.
   *
   * @return the names of the named values of this statement.
   * @throws IllegalStateException if this statement does not have named values.
   */
  public Set<String> getValueNames() {
    if (namedValues == null)
      throw new IllegalStateException("This statement does not have named values");
    return Collections.unmodifiableSet(namedValues.keySet());
  }

  /**
   * Returns the routing key for the query.
   *
   * <p>Unless the routing key has been explicitly set through {@link #setRoutingKey}, this method
   * will return {@code null} to avoid having to parse the query string to retrieve the partition
   * key.
   *
   * @param protocolVersion unused by this implementation (no internal serialization is required to
   *     compute the key).
   * @param codecRegistry unused by this implementation (no internal serialization is required to
   *     compute the key).
   * @return the routing key set through {@link #setRoutingKey} if such a key was set, {@code null}
   *     otherwise.
   * @see Statement#getRoutingKey
   */
  @Override
  public ByteBuffer getRoutingKey(ProtocolVersion protocolVersion, CodecRegistry codecRegistry) {
    return routingKey;
  }

  /**
   * Sets the routing key for this query.
   *
   * <p>This method allows you to manually provide a routing key for this query. It is thus optional
   * since the routing key is only an hint for token aware load balancing policy but is never
   * mandatory.
   *
   * <p>If the partition key for the query is composite, use the {@link
   * #setRoutingKey(ByteBuffer...)} method instead to build the routing key.
   *
   * @param routingKey the raw (binary) value to use as routing key.
   * @return this {@code SimpleStatement} object.
   * @see Statement#getRoutingKey
   */
  public SimpleStatement setRoutingKey(ByteBuffer routingKey) {
    this.routingKey = routingKey;
    return this;
  }

  /**
   * Returns the keyspace this query operates on.
   *
   * <p>Unless the keyspace has been explicitly set through {@link #setKeyspace}, this method will
   * return {@code null} to avoid having to parse the query string.
   *
   * @return the keyspace set through {@link #setKeyspace} if such keyspace was set, {@code null}
   *     otherwise.
   * @see Statement#getKeyspace
   */
  @Override
  public String getKeyspace() {
    return keyspace;
  }

  /**
   * Sets the keyspace this query operates on.
   *
   * <p>This method allows you to manually provide a keyspace for this query. It is thus optional
   * since the value returned by this method is only an hint for token aware load balancing policy
   * but is never mandatory.
   *
   * <p>Do note that if the query does not use a fully qualified keyspace, then you do not need to
   * set the keyspace through that method as the currently logged in keyspace will be used.
   *
   * @param keyspace the name of the keyspace this query operates on.
   * @return this {@code SimpleStatement} object.
   * @see Statement#getKeyspace
   */
  public SimpleStatement setKeyspace(String keyspace) {
    this.keyspace = keyspace;
    return this;
  }

  /**
   * Sets the routing key for this query.
   *
   * <p>See {@link #setRoutingKey(ByteBuffer)} for more information. This method is a variant for
   * when the query partition key is composite and thus the routing key must be built from multiple
   * values.
   *
   * @param routingKeyComponents the raw (binary) values to compose to obtain the routing key.
   * @return this {@code SimpleStatement} object.
   * @see Statement#getRoutingKey
   */
  public SimpleStatement setRoutingKey(ByteBuffer... routingKeyComponents) {
    this.routingKey = compose(routingKeyComponents);
    return this;
  }

  /*
   * This method performs a best-effort heuristic to guess which codec to use.
   * Note that this is not particularly efficient as the codec registry needs to iterate over
   * the registered codecs until it finds a suitable one.
   */
  private static ByteBuffer[] convert(
      Object[] values, ProtocolVersion protocolVersion, CodecRegistry codecRegistry) {
    ByteBuffer[] serializedValues = new ByteBuffer[values.length];
    for (int i = 0; i < values.length; i++) {
      Object value = values[i];
      if (value == null) {
        // impossible to locate the right codec when object is null,
        // so forcing the result to null
        serializedValues[i] = null;
      } else {
        if (value instanceof Token) {
          // bypass CodecRegistry for Token instances
          serializedValues[i] = ((Token) value).serialize(protocolVersion);
        } else {
          try {
            TypeCodec<Object> codec = codecRegistry.codecFor(value);
            serializedValues[i] = codec.serialize(value, protocolVersion);
          } catch (Exception e) {
            // Catch and rethrow to provide a more helpful error message (one that include which
            // value is bad)
            throw new InvalidTypeException(
                String.format(
                    "Value %d of type %s does not correspond to any CQL3 type",
                    i, value.getClass()),
                e);
          }
        }
      }
    }
    return serializedValues;
  }

  private static Map<String, ByteBuffer> convert(
      Map<String, Object> values, ProtocolVersion protocolVersion, CodecRegistry codecRegistry) {
    Map<String, ByteBuffer> serializedValues = new HashMap<String, ByteBuffer>();
    for (Map.Entry<String, Object> entry : values.entrySet()) {
      String name = entry.getKey();
      Object value = entry.getValue();
      if (value == null) {
        // impossible to locate the right codec when object is null,
        // so forcing the result to null
        serializedValues.put(name, null);
      } else {
        if (value instanceof Token) {
          // bypass CodecRegistry for Token instances
          serializedValues.put(name, ((Token) value).serialize(protocolVersion));
        } else {
          try {
            TypeCodec<Object> codec = codecRegistry.codecFor(value);
            serializedValues.put(name, codec.serialize(value, protocolVersion));
          } catch (Exception e) {
            // Catch and rethrow to provide a more helpful error message (one that include which
            // value is bad)
            throw new InvalidTypeException(
                String.format(
                    "Value '%s' of type %s does not correspond to any CQL3 type",
                    name, value.getClass()),
                e);
          }
        }
      }
    }
    return serializedValues;
  }

  /**
   * Utility method to assemble different routing key components into a single {@link ByteBuffer}.
   * Mainly intended for statements that need to generate a routing key out of their current values.
   *
   * @param buffers the components of the routing key.
   * @return A ByteBuffer containing the serialized routing key
   */
  static ByteBuffer compose(ByteBuffer... buffers) {
    if (buffers.length == 1) return buffers[0];

    int totalLength = 0;
    for (ByteBuffer bb : buffers) totalLength += 2 + bb.remaining() + 1;

    ByteBuffer out = ByteBuffer.allocate(totalLength);
    for (ByteBuffer buffer : buffers) {
      ByteBuffer bb = buffer.duplicate();
      putShortLength(out, bb.remaining());
      out.put(bb);
      out.put((byte) 0);
    }
    out.flip();
    return out;
  }

  static void putShortLength(ByteBuffer bb, int length) {
    bb.put((byte) ((length >> 8) & 0xFF));
    bb.put((byte) (length & 0xFF));
  }
}
