/*
 *  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.gora.cassandra.serializers;

import com.datastax.driver.core.AbstractGettableData;
import com.datastax.driver.core.ColumnDefinitions;
import com.datastax.driver.core.ConsistencyLevel;
import com.datastax.driver.core.DataType;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.Row;
import com.datastax.driver.core.SimpleStatement;
import com.datastax.driver.core.UDTValue;
import com.datastax.driver.core.UserType;
import org.apache.avro.Schema;
import org.apache.avro.specific.SpecificData;
import org.apache.commons.lang.ArrayUtils;
import org.apache.gora.cassandra.bean.CassandraKey;
import org.apache.gora.cassandra.bean.Field;
import org.apache.gora.cassandra.query.CassandraResultSet;
import org.apache.gora.cassandra.store.CassandraClient;
import org.apache.gora.cassandra.store.CassandraMapping;
import org.apache.gora.persistency.Persistent;
import org.apache.gora.persistency.impl.PersistentBase;
import org.apache.gora.query.Query;
import org.apache.gora.query.Result;
import org.apache.gora.store.DataStore;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

/**
 * This class contains the operations relates to Avro Serialization.
 */
class AvroSerializer<K, T extends PersistentBase> extends CassandraSerializer {


  private static final Logger LOG = LoggerFactory.getLogger(AvroSerializer.class);

  private DataStore<K, T> cassandraDataStore;

  private Schema persistentSchema;

  AvroSerializer(CassandraClient cassandraClient, DataStore<K, T> dataStore, CassandraMapping mapping) {
    super(cassandraClient, dataStore.getKeyClass(), dataStore.getPersistentClass(), mapping);
    if (PersistentBase.class.isAssignableFrom(dataStore.getPersistentClass())) {
      persistentSchema = ((PersistentBase) dataStore.getBeanFactory().getCachedPersistent()).getSchema();
    } else {
      persistentSchema = null;
    }
    this.cassandraDataStore = dataStore;
    try {
      analyzePersistent();
    } catch (Exception e) {
      throw new RuntimeException("Error occurred while analyzing the persistent class, :" + e.getMessage());
    }
  }

  /**
   * {@inheritDoc}
   *
   * @throws Exception
   */
  protected void analyzePersistent() throws Exception {
    userDefineTypeMaps = new HashMap<>();
    for (Field field : mapping.getFieldList()) {
      String fieldType = field.getType();
      if (fieldType.contains("frozen")) {
        String udtType = fieldType.substring(fieldType.indexOf("<") + 1, fieldType.indexOf(">"));
        if (PersistentBase.class.isAssignableFrom(persistentClass)) {
          Schema fieldSchema = persistentSchema.getField(field.getFieldName()).schema();
          if (fieldSchema.getType().equals(Schema.Type.UNION)) {
            for (Schema currentSchema : fieldSchema.getTypes()) {
              if (currentSchema.getType().equals(Schema.Type.RECORD)) {
                fieldSchema = currentSchema;
                break;
              }
            }
          }
          String createQuery = CassandraQueryFactory.getCreateUDTTypeForAvro(mapping, udtType, fieldSchema);
          userDefineTypeMaps.put(udtType, createQuery);
        } else {
          throw new RuntimeException("Unsupported Class for User Define Types, Please use PersistentBase class. field : " + udtType);
        }
      }
    }
  }

  /**
   * {@inheritDoc}
   *
   * @param query
   * @return
   */
  @Override
  public boolean updateByQuery(Query query) {
    List<Object> objectArrayList = new ArrayList<>();
    String cqlQuery = CassandraQueryFactory.getUpdateByQueryForAvro(mapping, query, objectArrayList, persistentSchema);
    ResultSet results;
    SimpleStatement statement;
    if (objectArrayList.size() == 0) {
      statement = new SimpleStatement(cqlQuery);
    } else {
      statement = new SimpleStatement(cqlQuery, objectArrayList.toArray());
    }
    if (writeConsistencyLevel != null) {
      statement.setConsistencyLevel(ConsistencyLevel.valueOf(writeConsistencyLevel));
    }
    results = client.getSession().execute(statement);
    return results.wasApplied();
  }

  /**
   * {@inheritDoc}
   *
   * @param key
   * @param fields
   * @return
   */
  @Override
  public Persistent get(Object key, String[] fields) {
    if (fields == null) {
      fields = getFields();
    }
    ArrayList<String> cassandraKeys = new ArrayList<>();
    ArrayList<Object> cassandraValues = new ArrayList<>();
    AvroCassandraUtils.processKeys(mapping, key, cassandraKeys, cassandraValues);
    String cqlQuery = CassandraQueryFactory.getSelectObjectWithFieldsQuery(mapping, fields, cassandraKeys);
    SimpleStatement statement = new SimpleStatement(cqlQuery, cassandraValues.toArray());
    if (readConsistencyLevel != null) {
      statement.setConsistencyLevel(ConsistencyLevel.valueOf(readConsistencyLevel));
    }
    ResultSet resultSet = this.client.getSession().execute(statement);
    Iterator<Row> iterator = resultSet.iterator();
    ColumnDefinitions definitions = resultSet.getColumnDefinitions();
    T obj = null;
    if (iterator.hasNext()) {
      obj = cassandraDataStore.newPersistent();
      AbstractGettableData row = (AbstractGettableData) iterator.next();
      populateValuesToPersistent(row, definitions, obj, fields);
    }
    return obj;
  }

  /**
   * {@inheritDoc}
   *
   * @param key
   * @param persistent
   */
  @Override
  public void put(Object key, Persistent persistent) {
    if (persistent instanceof PersistentBase) {
      if (persistent.isDirty()) {
        PersistentBase persistentBase = (PersistentBase) persistent;
        ArrayList<String> fields = new ArrayList<>();
        ArrayList<Object> values = new ArrayList<>();
        AvroCassandraUtils.processKeys(mapping, key, fields, values);
        for (Schema.Field f : persistentBase.getSchema().getFields()) {
          String fieldName = f.name();
          Field field = mapping.getFieldFromFieldName(fieldName);
          if (field == null) {
            LOG.debug("Ignoring {} adding field, {} field can't find in {} mapping", new Object[]{fieldName, fieldName, persistentClass});
            continue;
          }
          if (persistent.isDirty(f.pos()) || mapping.getInlinedDefinedPartitionKey().equals(mapping.getFieldFromFieldName(fieldName))) {
            Object value = persistentBase.get(f.pos());
            String fieldType = field.getType();
            if (fieldType.contains("frozen")) {
              fieldType = fieldType.substring(fieldType.indexOf("<") + 1, fieldType.indexOf(">"));
              UserType userType = client.getSession().getCluster().getMetadata().getKeyspace(mapping.getKeySpace().getName()).getUserType(fieldType);
              UDTValue udtValue = userType.newValue();
              Schema udtSchema = f.schema();
              if (udtSchema.getType().equals(Schema.Type.UNION)) {
                for (Schema schema : udtSchema.getTypes()) {
                  if (schema.getType().equals(Schema.Type.RECORD)) {
                    udtSchema = schema;
                    break;
                  }
                }
              }
              PersistentBase udtObjectBase = (PersistentBase) value;
              for (Schema.Field udtField : udtSchema.getFields()) {
                Object udtFieldValue = AvroCassandraUtils.getFieldValueFromAvroBean(udtField.schema(), udtField.schema().getType(), udtObjectBase.get(udtField.name()), field);
                if (udtField.schema().getType().equals(Schema.Type.MAP)) {
                  udtValue.setMap(udtField.name(), (Map) udtFieldValue);
                } else if (udtField.schema().getType().equals(Schema.Type.ARRAY)) {
                  udtValue.setList(udtField.name(), (List) udtFieldValue);
                } else {
                  udtValue.set(udtField.name(), udtFieldValue, (Class) udtFieldValue.getClass());
                }
              }
              value = udtValue;
            } else {
              value = AvroCassandraUtils.getFieldValueFromAvroBean(f.schema(), f.schema().getType(), value, field);
            }
            values.add(value);
            fields.add(fieldName);
          }
        }
        String cqlQuery = CassandraQueryFactory.getInsertDataQuery(mapping, fields);
        SimpleStatement statement = new SimpleStatement(cqlQuery, values.toArray());
        if (writeConsistencyLevel != null) {
          statement.setConsistencyLevel(ConsistencyLevel.valueOf(writeConsistencyLevel));
        }
        client.getSession().execute(statement);
      } else {
        LOG.info("Ignored putting persistent bean {} in the store as it is neither "
                + "new, neither dirty.", new Object[]{persistent});
      }
    } else {
      LOG.error("{} Persistent bean isn't extended by {} .", new Object[]{this.persistentClass, PersistentBase.class});
    }
  }

  /**
   * {@inheritDoc}
   *
   * @param key
   * @return
   */
  @Override
  public Persistent get(Object key) {
    ArrayList<String> cassandraKeys = new ArrayList<>();
    ArrayList<Object> cassandraValues = new ArrayList<>();
    AvroCassandraUtils.processKeys(mapping, key, cassandraKeys, cassandraValues);
    String cqlQuery = CassandraQueryFactory.getSelectObjectQuery(mapping, cassandraKeys);
    SimpleStatement statement = new SimpleStatement(cqlQuery, cassandraValues.toArray());
    if (readConsistencyLevel != null) {
      statement.setConsistencyLevel(ConsistencyLevel.valueOf(readConsistencyLevel));
    }
    ResultSet resultSet = client.getSession().execute(statement);
    Iterator<Row> iterator = resultSet.iterator();
    ColumnDefinitions definitions = resultSet.getColumnDefinitions();
    T obj = null;
    if (iterator.hasNext()) {
      obj = cassandraDataStore.newPersistent();
      AbstractGettableData row = (AbstractGettableData) iterator.next();
      populateValuesToPersistent(row, definitions, obj, mapping.getFieldNames());
    }
    return obj;
  }

  /**
   * This method wraps result set data in to DataEntry and creates a list of DataEntry.
   **/
  private void populateValuesToPersistent(AbstractGettableData row, ColumnDefinitions columnDefinitions, PersistentBase base, String[] fields) {
    Object paramValue;
    for (String fieldName : fields) {
      Schema.Field avroField = base.getSchema().getField(fieldName);
      Field field = mapping.getFieldFromFieldName(fieldName);
      //to ignore unspecified fields in the mapping
      if (field == null || avroField == null) {
        continue;
      }
      Schema fieldSchema = avroField.schema();
      String columnName = field.getColumnName();
      paramValue = getValue(row, columnDefinitions.getType(columnName), columnName, fieldSchema);
      Object value = AvroCassandraUtils.getAvroFieldValue(paramValue, fieldSchema);
      base.put(avroField.pos(), value);
    }
  }

  private Object getValue(AbstractGettableData row, DataType columnType, String columnName, Schema schema) {
    Object paramValue;
    String dataType;
    switch (columnType.getName()) {
      case ASCII:
        paramValue = row.getString(columnName);
        break;
      case BIGINT:
        paramValue = row.isNull(columnName) ? null : row.getLong(columnName);
        break;
      case BLOB:
        paramValue = row.isNull(columnName) ? null : row.getBytes(columnName);
        break;
      case BOOLEAN:
        paramValue = row.isNull(columnName) ? null : row.getBool(columnName);
        break;
      case COUNTER:
        paramValue = row.isNull(columnName) ? null : row.getLong(columnName);
        break;
      case DECIMAL:
        paramValue = row.isNull(columnName) ? null : row.getDecimal(columnName);
        break;
      case DOUBLE:
        paramValue = row.isNull(columnName) ? null : row.getDouble(columnName);
        break;
      case FLOAT:
        paramValue = row.isNull(columnName) ? null : row.getFloat(columnName);
        break;
      case INET:
        paramValue = row.isNull(columnName) ? null : row.getInet(columnName).toString();
        break;
      case INT:
        paramValue = row.isNull(columnName) ? null : row.getInt(columnName);
        break;
      case TEXT:
        paramValue = row.getString(columnName);
        break;
      case TIMESTAMP:
        paramValue = row.isNull(columnName) ? null : row.getDate(columnName);
        break;
      case UUID:
        paramValue = row.isNull(columnName) ? null : row.getUUID(columnName);
        break;
      case VARCHAR:
        paramValue = row.getString(columnName);
        break;
      case VARINT:
        paramValue = row.isNull(columnName) ? null : row.getVarint(columnName);
        break;
      case TIMEUUID:
        paramValue = row.isNull(columnName) ? null : row.getUUID(columnName);
        break;
      case LIST:
        dataType = columnType.getTypeArguments().get(0).toString();
        paramValue = row.isNull(columnName) ? null : row.getList(columnName, AvroCassandraUtils.getRelevantClassForCassandraDataType(dataType));
        break;
      case SET:
        dataType = columnType.getTypeArguments().get(0).toString();
        paramValue = row.isNull(columnName) ? null : row.getList(columnName, AvroCassandraUtils.getRelevantClassForCassandraDataType(dataType));
        break;
      case MAP:
        dataType = columnType.getTypeArguments().get(1).toString();
        // Avro supports only String for keys
        paramValue = row.isNull(columnName) ? null : row.getMap(columnName, String.class, AvroCassandraUtils.getRelevantClassForCassandraDataType(dataType));
        break;
      case UDT:
        paramValue = row.isNull(columnName) ? null : row.getUDTValue(columnName);
        if (paramValue != null) {
          try {
            PersistentBase udtObject = (PersistentBase) SpecificData.newInstance(Class.forName(schema.getFullName()), schema);
            for (Schema.Field f : udtObject.getSchema().getFields()) {
              DataType dType = ((UDTValue) paramValue).getType().getFieldType(f.name());
              Object fieldValue = getValue((UDTValue) paramValue, dType, f.name(), f.schema());
              udtObject.put(f.pos(), fieldValue);
            }
            paramValue = udtObject;
          } catch (ClassNotFoundException e) {
            throw new RuntimeException("Error occurred while populating data to " + schema.getFullName() + " : " + e.getMessage());
          }
        }
        break;
      case TUPLE:
        paramValue = row.isNull(columnName) ? null : row.getTupleValue(columnName).toString();
        break;
      case CUSTOM:
        paramValue = row.isNull(columnName) ? null : row.getBytes(columnName);
        break;
      default:
        paramValue = row.getString(columnName);
        break;
    }
    return paramValue;
  }

  /**
   * {@inheritDoc}
   *
   * @param key
   * @return
   */
  @Override
  public boolean delete(Object key) {
    ArrayList<String> cassandraKeys = new ArrayList<>();
    ArrayList<Object> cassandraValues = new ArrayList<>();
    AvroCassandraUtils.processKeys(mapping, key, cassandraKeys, cassandraValues);
    String cqlQuery = CassandraQueryFactory.getDeleteDataQuery(mapping, cassandraKeys);
    SimpleStatement statement = new SimpleStatement(cqlQuery, cassandraValues.toArray());
    if (writeConsistencyLevel != null) {
      statement.setConsistencyLevel(ConsistencyLevel.valueOf(writeConsistencyLevel));
    }
    ResultSet resultSet = client.getSession().execute(statement);
    return resultSet.wasApplied();
  }

  /**
   * {@inheritDoc}
   *
   * @param dataStore
   * @param query
   * @return
   */
  @Override
  public Result execute(DataStore dataStore, Query query) {
    List<Object> objectArrayList = new ArrayList<>();
    String[] fields = query.getFields();
    if (fields != null) {
      fields = (String[]) ArrayUtils.addAll(fields, mapping.getAllKeys());
    } else {
      fields = mapping.getAllFieldsIncludingKeys();
    }
    CassandraResultSet<K, T> cassandraResult = new CassandraResultSet<>(dataStore, query);
    String cqlQuery = CassandraQueryFactory.getExecuteQuery(mapping, query, objectArrayList, fields);
    ResultSet results;
    SimpleStatement statement;
    if (objectArrayList.size() == 0) {
      statement = new SimpleStatement(cqlQuery);
    } else {
      statement = new SimpleStatement(cqlQuery, objectArrayList.toArray());
    }
    if (readConsistencyLevel != null) {
      statement.setConsistencyLevel(ConsistencyLevel.valueOf(readConsistencyLevel));
    }
    results = client.getSession().execute(statement);
    Iterator<Row> iterator = results.iterator();
    ColumnDefinitions definitions = results.getColumnDefinitions();
    T obj;
    K keyObject;
    CassandraKey cassandraKey = mapping.getCassandraKey();
    while (iterator.hasNext()) {
      AbstractGettableData row = (AbstractGettableData) iterator.next();
      obj = cassandraDataStore.newPersistent();
      keyObject = cassandraDataStore.newKey();
      populateValuesToPersistent(row, definitions, obj, fields);
      if (cassandraKey != null) {
        populateValuesToPersistent(row, definitions, (PersistentBase) keyObject, cassandraKey.getFieldNames());
      } else {
        Field key = mapping.getInlinedDefinedPartitionKey();
        keyObject = (K) getValue(row, definitions.getType(key.getColumnName()), key.getColumnName(), null);
      }
      cassandraResult.addResultElement(keyObject, obj);
    }
    return cassandraResult;
  }

}
