/**
 * 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.hcatalog.pig;


import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Properties;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hive.conf.HiveConf;
import org.apache.hadoop.hive.metastore.HiveMetaStoreClient;
import org.apache.hadoop.hive.metastore.MetaStoreUtils;
import org.apache.hadoop.hive.metastore.api.NoSuchObjectException;
import org.apache.hadoop.hive.ql.metadata.Table;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hcatalog.common.HCatConstants;
import org.apache.hcatalog.common.HCatException;
import org.apache.hcatalog.common.HCatUtil;
import org.apache.hcatalog.data.HCatRecord;
import org.apache.hcatalog.data.Pair;
import org.apache.hcatalog.data.schema.HCatFieldSchema;
import org.apache.hcatalog.data.schema.HCatFieldSchema.Type;
import org.apache.hcatalog.data.schema.HCatSchema;
import org.apache.pig.LoadPushDown.RequiredField;
import org.apache.pig.PigException;
import org.apache.pig.ResourceSchema;
import org.apache.pig.ResourceSchema.ResourceFieldSchema;
import org.apache.pig.data.DataBag;
import org.apache.pig.data.DataByteArray;
import org.apache.pig.data.DataType;
import org.apache.pig.data.DefaultDataBag;
import org.apache.pig.data.Tuple;
import org.apache.pig.data.TupleFactory;
import org.apache.pig.impl.util.UDFContext;

public class PigHCatUtil {

  static final int PIG_EXCEPTION_CODE = 1115; // http://wiki.apache.org/pig/PigErrorHandlingFunctionalSpecification#Error_codes
  private static final String DEFAULT_DB = MetaStoreUtils.DEFAULT_DATABASE_NAME;

  private final  Map<Pair<String,String>, Table> hcatTableCache =
    new HashMap<Pair<String,String>, Table>();

  private static final TupleFactory tupFac = TupleFactory.getInstance();

  static public Pair<String, String> getDBTableNames(String location) throws IOException {
    // the location string will be of the form:
    // <database name>.<table name> - parse it and
    // communicate the information to HCatInputFormat

    try {
      return HCatUtil.getDbAndTableName(location);
    } catch (IOException e) {
      String locationErrMsg = "The input location in load statement " +
      "should be of the form " +
      "<databasename>.<table name> or <table name>. Got " + location;
      throw new PigException(locationErrMsg, PIG_EXCEPTION_CODE);
    }
  }

  static public String getHCatServerUri(Job job) {

    return job.getConfiguration().get(HiveConf.ConfVars.METASTOREURIS.varname);
  }

  static public String getHCatServerPrincipal(Job job) {

    return job.getConfiguration().get(HCatConstants.HCAT_METASTORE_PRINCIPAL);
  }

  private static HiveMetaStoreClient getHiveMetaClient(String serverUri,
                                                       String serverKerberosPrincipal, Class<?> clazz) throws Exception {
    HiveConf hiveConf = new HiveConf(clazz);

    if (serverUri != null){
      hiveConf.set("hive.metastore.local", "false");
      hiveConf.setVar(HiveConf.ConfVars.METASTOREURIS, serverUri.trim());
    }

    if (serverKerberosPrincipal != null){
      hiveConf.setBoolVar(HiveConf.ConfVars.METASTORE_USE_THRIFT_SASL, true);
      hiveConf.setVar(HiveConf.ConfVars.METASTORE_KERBEROS_PRINCIPAL, serverKerberosPrincipal);
    }

    try {
        return HCatUtil.getHiveClient(hiveConf);
    } catch (Exception e){
      throw new Exception("Could not instantiate a HiveMetaStoreClient connecting to server uri:["+serverUri+"]",e);
    }
  }


  HCatSchema getHCatSchema(List<RequiredField> fields, String signature, Class<?> classForUDFCLookup) throws IOException {
    if(fields == null) {
      return null;
    }

    Properties props = UDFContext.getUDFContext().getUDFProperties(
        classForUDFCLookup, new String[] {signature});
    HCatSchema hcatTableSchema = (HCatSchema) props.get(HCatConstants.HCAT_TABLE_SCHEMA);

    ArrayList<HCatFieldSchema> fcols = new ArrayList<HCatFieldSchema>();
    for(RequiredField rf: fields) {
      fcols.add(hcatTableSchema.getFields().get(rf.getIndex()));
    }
    return new HCatSchema(fcols);
  }

  public Table getTable(String location, String hcatServerUri, String hcatServerPrincipal) throws IOException{
    Pair<String, String> loc_server = new Pair<String,String>(location, hcatServerUri);
    Table hcatTable = hcatTableCache.get(loc_server);
    if(hcatTable != null){
      return hcatTable;
    }

    Pair<String, String> dbTablePair = PigHCatUtil.getDBTableNames(location);
    String dbName = dbTablePair.first;
    String tableName = dbTablePair.second;
    Table table = null;
    HiveMetaStoreClient client = null;
    try {
      client = getHiveMetaClient(hcatServerUri, hcatServerPrincipal, PigHCatUtil.class);
      table = HCatUtil.getTable(client, dbName, tableName);
    } catch (NoSuchObjectException nsoe){
      throw new PigException("Table not found : " + nsoe.getMessage(), PIG_EXCEPTION_CODE); // prettier error messages to frontend
    } catch (Exception e) {
      throw new IOException(e);
    } finally {
        HCatUtil.closeHiveClientQuietly(client);
    }
    hcatTableCache.put(loc_server, table);
    return table;
  }

  public static ResourceSchema getResourceSchema(HCatSchema hcatSchema) throws IOException {

    List<ResourceFieldSchema> rfSchemaList = new ArrayList<ResourceFieldSchema>();
    for (HCatFieldSchema hfs : hcatSchema.getFields()){
      ResourceFieldSchema rfSchema;
      rfSchema = getResourceSchemaFromFieldSchema(hfs);
      rfSchemaList.add(rfSchema);
    }
    ResourceSchema rSchema = new ResourceSchema();
    rSchema.setFields(rfSchemaList.toArray(new ResourceFieldSchema[0]));
    return rSchema;

  }

  private static ResourceFieldSchema getResourceSchemaFromFieldSchema(HCatFieldSchema hfs)
      throws IOException {
    ResourceFieldSchema rfSchema;
    // if we are dealing with a bag or tuple column - need to worry about subschema
    if(hfs.getType() == Type.STRUCT) {
        rfSchema = new ResourceFieldSchema()
          .setName(hfs.getName())
          .setDescription(hfs.getComment())
          .setType(getPigType( hfs))
          .setSchema(getTupleSubSchema(hfs));
    } else if(hfs.getType() == Type.ARRAY) {
        rfSchema = new ResourceFieldSchema()
          .setName(hfs.getName())
          .setDescription(hfs.getComment())
          .setType(getPigType( hfs))
          .setSchema(getBagSubSchema(hfs));
    } else {
      rfSchema = new ResourceFieldSchema()
          .setName(hfs.getName())
          .setDescription(hfs.getComment())
          .setType(getPigType( hfs))
          .setSchema(null); // no munging inner-schemas
    }
    return rfSchema;
  }

  protected static ResourceSchema getBagSubSchema(HCatFieldSchema hfs) throws IOException {
    // there are two cases - array<Type> and array<struct<...>>
    // in either case the element type of the array is represented in a
    // tuple field schema in the bag's field schema - the second case (struct)
    // more naturally translates to the tuple - in the first case (array<Type>)
    // we simulate the tuple by putting the single field in a tuple

    Properties props = UDFContext.getUDFContext().getClientSystemProps();
    String innerTupleName = HCatConstants.HCAT_PIG_INNER_TUPLE_NAME_DEFAULT;
    if (props != null && props.containsKey(HCatConstants.HCAT_PIG_INNER_TUPLE_NAME)) {
      innerTupleName = props.getProperty(HCatConstants.HCAT_PIG_INNER_TUPLE_NAME)
          .replaceAll("FIELDNAME", hfs.getName());
    }
    String innerFieldName = HCatConstants.HCAT_PIG_INNER_FIELD_NAME_DEFAULT;
    if (props != null && props.containsKey(HCatConstants.HCAT_PIG_INNER_FIELD_NAME)) {
      innerFieldName = props.getProperty(HCatConstants.HCAT_PIG_INNER_FIELD_NAME)
          .replaceAll("FIELDNAME", hfs.getName());
    }

    ResourceFieldSchema[] bagSubFieldSchemas = new ResourceFieldSchema[1];
    bagSubFieldSchemas[0] = new ResourceFieldSchema().setName(innerTupleName)
      .setDescription("The tuple in the bag")
      .setType(DataType.TUPLE);
    HCatFieldSchema arrayElementFieldSchema = hfs.getArrayElementSchema().get(0);
    if(arrayElementFieldSchema.getType() == Type.STRUCT) {
      bagSubFieldSchemas[0].setSchema(getTupleSubSchema(arrayElementFieldSchema));
    } else if(arrayElementFieldSchema.getType() == Type.ARRAY) {
      ResourceSchema s = new ResourceSchema();
      List<ResourceFieldSchema> lrfs = Arrays.asList(getResourceSchemaFromFieldSchema(arrayElementFieldSchema));
      s.setFields(lrfs.toArray(new ResourceFieldSchema[0]));
      bagSubFieldSchemas[0].setSchema(s);
    } else {
      ResourceFieldSchema[] innerTupleFieldSchemas = new ResourceFieldSchema[1];
      innerTupleFieldSchemas[0] = new ResourceFieldSchema().setName(innerFieldName)
        .setDescription("The inner field in the tuple in the bag")
        .setType(getPigType(arrayElementFieldSchema))
        .setSchema(null); // the element type is not a tuple - so no subschema
      bagSubFieldSchemas[0].setSchema(new ResourceSchema().setFields(innerTupleFieldSchemas));
    }
    ResourceSchema s = new ResourceSchema().setFields(bagSubFieldSchemas);
    return s;

  }

  private static ResourceSchema getTupleSubSchema(HCatFieldSchema hfs) throws IOException {
    // for each struct subfield, create equivalent ResourceFieldSchema
    ResourceSchema s = new ResourceSchema();
    List<ResourceFieldSchema> lrfs = new ArrayList<ResourceFieldSchema>();
    for(HCatFieldSchema subField : hfs.getStructSubSchema().getFields()) {
      lrfs.add(getResourceSchemaFromFieldSchema(subField));
    }
    s.setFields(lrfs.toArray(new ResourceFieldSchema[0]));
    return s;
  }

/**
   * @param hfs the field schema of the column
   * @return corresponding pig type
   * @throws IOException
   */
  static public byte getPigType(HCatFieldSchema hfs) throws IOException {
    return getPigType(hfs.getType());
  }

  static public byte getPigType(Type type) throws IOException {
    String errMsg;

    if (type == Type.STRING){
      return DataType.CHARARRAY;
    }

    if ( (type == Type.INT) || (type == Type.SMALLINT) || (type == Type.TINYINT)){
      return DataType.INTEGER;
    }

    if (type == Type.ARRAY){
      return DataType.BAG;
    }

    if (type == Type.STRUCT){
      return DataType.TUPLE;
    }

    if (type == Type.MAP){
      return DataType.MAP;
    }

    if (type == Type.BIGINT){
      return DataType.LONG;
    }

    if (type == Type.FLOAT){
      return DataType.FLOAT;
    }

    if (type == Type.DOUBLE){
      return DataType.DOUBLE;
    }

    if (type == Type.BINARY){
        return DataType.BYTEARRAY;
    }

    if (type == Type.BOOLEAN){
      errMsg = "HCatalog column type 'BOOLEAN' is not supported in " +
      "Pig as a column type";
      throw new PigException(errMsg, PIG_EXCEPTION_CODE);
    }

    errMsg = "HCatalog column type '"+ type.toString() +"' is not supported in Pig as a column type";
    throw new PigException(errMsg, PIG_EXCEPTION_CODE);
  }

  public static Tuple transformToTuple(HCatRecord hr, HCatSchema hs) throws Exception {
    if (hr == null){
      return null;
    }
    return transformToTuple(hr.getAll(),hs);
  }

  @SuppressWarnings("unchecked")
  public static Object extractPigObject(Object o, HCatFieldSchema hfs) throws Exception {
    Object result;
    Type itemType = hfs.getType();
    switch (itemType){
    case BINARY:
      result = (o == null) ? null : new DataByteArray((byte[])o);
      break;
    case STRUCT:
      result = transformToTuple((List<Object>)o,hfs);
      break;
    case ARRAY:
      result = transformToBag((List<? extends Object>) o,hfs);
      break;
    case MAP:
      result = transformToPigMap((Map<String, Object>)o,hfs);
      break;
    default:
      result = o;
      break;
    }
    return result;
  }

  public static Tuple transformToTuple(List<? extends Object> objList, HCatFieldSchema hfs) throws Exception {
      try {
          return transformToTuple(objList,hfs.getStructSubSchema());
      } catch (Exception e){
          if (hfs.getType() != Type.STRUCT){
              throw new Exception("Expected Struct type, got "+hfs.getType(), e);
          } else {
              throw e;
          }
      }
  }

  public static Tuple transformToTuple(List<? extends Object> objList, HCatSchema hs) throws Exception {
    if (objList == null){
      return null;
    }
    Tuple t = tupFac.newTuple(objList.size());
    List<HCatFieldSchema> subFields = hs.getFields();
    for (int i = 0; i < subFields.size(); i++){
      t.set(i,extractPigObject(objList.get(i), subFields.get(i)));
    }
    return t;
  }

  public static Map<String,Object> transformToPigMap(Map<String,Object> map, HCatFieldSchema hfs) throws Exception {
    if (map == null) {
      return null;
    }

    Map<String,Object> result = new HashMap<String, Object>();
    for (Entry<String, Object> entry : map.entrySet()) {
      result.put(entry.getKey(), extractPigObject(entry.getValue(), hfs.getMapValueSchema().get(0)));
    }
    return result;
  }

  @SuppressWarnings("unchecked")
  public static DataBag transformToBag(List<? extends Object> list, HCatFieldSchema hfs) throws Exception {
    if (list == null){
      return null;
    }

    HCatFieldSchema elementSubFieldSchema = hfs.getArrayElementSchema().getFields().get(0);
    DataBag db = new DefaultDataBag();
    for (Object o : list){
      Tuple tuple;
      if (elementSubFieldSchema.getType() == Type.STRUCT){
        tuple = transformToTuple((List<Object>)o, elementSubFieldSchema);
      } else {
        // bags always contain tuples
        tuple = tupFac.newTuple(extractPigObject(o, elementSubFieldSchema));
      }
      db.add(tuple);
    }
    return db;
  }


  private static void validateHCatSchemaFollowsPigRules(HCatSchema tblSchema) throws PigException {
    for(HCatFieldSchema hcatField : tblSchema.getFields()){
      validateHcatFieldFollowsPigRules(hcatField);
    }
  }

  private static void validateHcatFieldFollowsPigRules(HCatFieldSchema hcatField) throws PigException {
    try {
      Type hType = hcatField.getType();
      switch(hType){
      case BOOLEAN:
        throw new PigException("Incompatible type found in hcat table schema: "+hcatField, PigHCatUtil.PIG_EXCEPTION_CODE);
      case ARRAY:
        validateHCatSchemaFollowsPigRules(hcatField.getArrayElementSchema());
        break;
      case STRUCT:
        validateHCatSchemaFollowsPigRules(hcatField.getStructSubSchema());
        break;
      case MAP:
        // key is only string
        validateHCatSchemaFollowsPigRules(hcatField.getMapValueSchema());
        break;
      }
    } catch (HCatException e) {
      throw new PigException("Incompatible type found in hcat table schema: "+hcatField, PigHCatUtil.PIG_EXCEPTION_CODE, e);
    }
  }


  public static void validateHCatTableSchemaFollowsPigRules(HCatSchema hcatTableSchema) throws IOException {
    validateHCatSchemaFollowsPigRules(hcatTableSchema);
  }

  public static void getConfigFromUDFProperties(Properties p, Configuration config, String propName) {
    if(p.getProperty(propName) != null){
      config.set(propName, p.getProperty(propName));
    }
  }

  public static void saveConfigIntoUDFProperties(Properties p, Configuration config, String propName) {
    if(config.get(propName) != null){
      p.setProperty(propName, config.get(propName));
    }
  }

}