/*
 * 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.xtable.kernel;

import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import io.delta.kernel.types.ArrayType;
import io.delta.kernel.types.BinaryType;
import io.delta.kernel.types.BooleanType;
import io.delta.kernel.types.DataType;
import io.delta.kernel.types.DateType;
import io.delta.kernel.types.DecimalType;
import io.delta.kernel.types.DoubleType;
import io.delta.kernel.types.FieldMetadata;
import io.delta.kernel.types.FloatType;
import io.delta.kernel.types.IntegerType;
import io.delta.kernel.types.LongType;
import io.delta.kernel.types.MapType;
import io.delta.kernel.types.StringType;
import io.delta.kernel.types.StructField;
import io.delta.kernel.types.StructType;
import io.delta.kernel.types.TimestampNTZType;
import io.delta.kernel.types.TimestampType;

import org.apache.xtable.collectors.CustomCollectors;
import org.apache.xtable.delta.DeltaPartitionExtractor;
import org.apache.xtable.exception.NotSupportedException;
import org.apache.xtable.model.schema.InternalField;
import org.apache.xtable.model.schema.InternalSchema;
import org.apache.xtable.model.schema.InternalType;
import org.apache.xtable.schema.SchemaUtils;

public class DeltaKernelSchemaExtractor {

  private static final String DELTA_COLUMN_MAPPING_ID = "delta.columnMapping.id";
  private static final DeltaKernelSchemaExtractor INSTANCE = new DeltaKernelSchemaExtractor();
  private static final Map<InternalSchema.MetadataKey, Object>
      DEFAULT_TIMESTAMP_PRECISION_METADATA =
          Collections.singletonMap(
              InternalSchema.MetadataKey.TIMESTAMP_PRECISION, InternalSchema.MetadataValue.MICROS);
  static final String DELTA_GENERATION_EXPRESSION = "delta.generationExpression";

  public static DeltaKernelSchemaExtractor getInstance() {
    return INSTANCE;
  }

  public InternalSchema toInternalSchema(StructType structType) {
    return toInternalSchema(structType, null, false, null, null);
  }

  private InternalSchema toInternalSchema(
      DataType dataType,
      String parentPath,
      boolean nullable,
      String comment,
      FieldMetadata originalMetadata) {

    String trimmedTypeName = "";
    InternalType type = null;
    Map<InternalSchema.MetadataKey, Object> metadata = null;
    List<InternalField> fields = null;

    if (dataType instanceof IntegerType) {
      type = InternalType.INT;
      trimmedTypeName = "int";
    } else if (dataType instanceof StringType) {
      type = InternalType.STRING;
      trimmedTypeName = "string";
    } else if (dataType instanceof BooleanType) {
      type = InternalType.BOOLEAN;
      trimmedTypeName = "boolean";
    } else if (dataType instanceof FloatType) {
      type = InternalType.FLOAT;
      trimmedTypeName = "float";
    } else if (dataType instanceof DoubleType) {
      type = InternalType.DOUBLE;
      trimmedTypeName = "double";
    } else if (dataType instanceof BinaryType) {
      if (originalMetadata != null
          && originalMetadata.contains(InternalSchema.XTABLE_LOGICAL_TYPE)
          && "uuid".equals(originalMetadata.getString(InternalSchema.XTABLE_LOGICAL_TYPE))) {
        type = InternalType.UUID;
        trimmedTypeName = "binary";
      } else {
        type = InternalType.BYTES;
        trimmedTypeName = "binary";
      }
    } else if (dataType instanceof LongType) {
      type = InternalType.LONG;
      trimmedTypeName = "long";
    } else if (dataType instanceof DateType) {
      type = InternalType.DATE;
      trimmedTypeName = "date";
    } else if (dataType instanceof TimestampType) {
      type = InternalType.TIMESTAMP;
      metadata = DEFAULT_TIMESTAMP_PRECISION_METADATA;
      trimmedTypeName = "timestamp";
    } else if (dataType instanceof TimestampNTZType) {
      type = InternalType.TIMESTAMP_NTZ;
      metadata = DEFAULT_TIMESTAMP_PRECISION_METADATA;
      trimmedTypeName = "timestamp_ntz";
    } else if (dataType instanceof StructType) {
      // Handle StructType
      StructType structType = (StructType) dataType;
      // your logic here

      fields =
          structType.fields().stream()
              .filter(
                  field ->
                      !field
                          .getMetadata()
                          .contains(DeltaPartitionExtractor.DELTA_GENERATION_EXPRESSION))
              .map(
                  field -> {
                    Integer fieldId =
                        field.getMetadata().contains(DELTA_COLUMN_MAPPING_ID)
                            ? Long.valueOf(field.getMetadata().getLong(DELTA_COLUMN_MAPPING_ID))
                                .intValue()
                            : null;
                    String fieldComment =
                        field.getMetadata().contains("comment")
                            ? field.getMetadata().getString("comment")
                            : null;
                    InternalSchema schema =
                        toInternalSchema(
                            field.getDataType(),
                            SchemaUtils.getFullyQualifiedPath(parentPath, field.getName()),
                            field.isNullable(),
                            fieldComment,
                            field.getMetadata());
                    return InternalField.builder()
                        .name(field.getName())
                        .fieldId(fieldId)
                        .parentPath(parentPath)
                        .schema(schema)
                        .defaultValue(
                            field.isNullable() ? InternalField.Constants.NULL_DEFAULT_VALUE : null)
                        .build();
                  })
              .collect(CustomCollectors.toList(structType.fields().size()));
      type = InternalType.RECORD;
      trimmedTypeName = "struct";
    } else if (dataType instanceof DecimalType) {
      DecimalType decimalType = (DecimalType) dataType;
      metadata = new HashMap<>(2, 1.0f);
      metadata.put(InternalSchema.MetadataKey.DECIMAL_PRECISION, decimalType.getPrecision());
      metadata.put(InternalSchema.MetadataKey.DECIMAL_SCALE, decimalType.getScale());
      type = InternalType.DECIMAL;
      trimmedTypeName = "decimal";

    } else if (dataType instanceof ArrayType) {
      ArrayType arrayType = (ArrayType) dataType;
      InternalSchema elementSchema =
          toInternalSchema(
              arrayType.getElementType(),
              SchemaUtils.getFullyQualifiedPath(
                  parentPath, InternalField.Constants.ARRAY_ELEMENT_FIELD_NAME),
              arrayType.containsNull(),
              null,
              null);
      InternalField elementField =
          InternalField.builder()
              .name(InternalField.Constants.ARRAY_ELEMENT_FIELD_NAME)
              .parentPath(parentPath)
              .schema(elementSchema)
              .build();
      type = InternalType.LIST;
      fields = Collections.singletonList(elementField);
      trimmedTypeName = "array";
    } else if (dataType instanceof MapType) {
      MapType mapType = (MapType) dataType;
      InternalSchema keySchema =
          toInternalSchema(
              mapType.getKeyType(),
              SchemaUtils.getFullyQualifiedPath(
                  parentPath, InternalField.Constants.MAP_KEY_FIELD_NAME),
              false,
              null,
              null);
      InternalField keyField =
          InternalField.builder()
              .name(InternalField.Constants.MAP_KEY_FIELD_NAME)
              .parentPath(parentPath)
              .schema(keySchema)
              .build();
      InternalSchema valueSchema =
          toInternalSchema(
              mapType.getValueType(),
              SchemaUtils.getFullyQualifiedPath(
                  parentPath, InternalField.Constants.MAP_VALUE_FIELD_NAME),
              mapType.isValueContainsNull(),
              null,
              null);
      InternalField valueField =
          InternalField.builder()
              .name(InternalField.Constants.MAP_VALUE_FIELD_NAME)
              .parentPath(parentPath)
              .schema(valueSchema)
              .build();
      type = InternalType.MAP;
      fields = Arrays.asList(keyField, valueField);
      trimmedTypeName = "map";
    } else {
      throw new NotSupportedException("Unsupported type: " + dataType);
    }
    return InternalSchema.builder()
        .name(trimmedTypeName)
        .dataType(type)
        .comment(comment)
        .isNullable(nullable)
        .metadata(metadata)
        .fields(fields)
        .build();
  }

  /**
   * Converts an InternalSchema to Delta Kernel StructType.
   *
   * @param internalSchema the internal schema representation
   * @return Delta Kernel StructType
   */
  public StructType fromInternalSchema(InternalSchema internalSchema) {
    List<StructField> fields =
        internalSchema.getFields().stream()
            .map(
                field ->
                    new StructField(
                        field.getName(),
                        convertFieldType(field),
                        field.getSchema().isNullable(),
                        getFieldMetadata(field)))
            .collect(CustomCollectors.toList(internalSchema.getFields().size()));
    return new StructType(fields);
  }

  /**
   * Converts an InternalField to Delta Kernel DataType.
   *
   * @param field the internal field
   * @return Delta Kernel DataType
   */
  private DataType convertFieldType(InternalField field) {
    switch (field.getSchema().getDataType()) {
      case STRING:
      case ENUM:
        return StringType.STRING;
      case INT:
        return IntegerType.INTEGER;
      case LONG:
        return LongType.LONG;
      case BYTES:
      case FIXED:
      case UUID:
        return BinaryType.BINARY;
      case BOOLEAN:
        return BooleanType.BOOLEAN;
      case FLOAT:
        return FloatType.FLOAT;
      case DATE:
        return DateType.DATE;
      case TIMESTAMP:
        return TimestampType.TIMESTAMP;
      case TIMESTAMP_NTZ:
        return TimestampNTZType.TIMESTAMP_NTZ;
      case DOUBLE:
        return DoubleType.DOUBLE;
      case DECIMAL:
        Map<InternalSchema.MetadataKey, Object> metadata = field.getSchema().getMetadata();
        if (metadata == null) {
          throw new IllegalStateException(
              String.format(
                  "DECIMAL field '%s' missing metadata - cannot determine precision/scale",
                  field.getName()));
        }
        Object precisionObj = metadata.get(InternalSchema.MetadataKey.DECIMAL_PRECISION);
        Object scaleObj = metadata.get(InternalSchema.MetadataKey.DECIMAL_SCALE);
        if (!(precisionObj instanceof Number) || !(scaleObj instanceof Number)) {
          throw new IllegalStateException(
              String.format(
                  "DECIMAL field '%s' missing precision/scale metadata - found precision=%s, scale=%s",
                  field.getName(), precisionObj, scaleObj));
        }
        int precision = ((Number) precisionObj).intValue();
        int scale = ((Number) scaleObj).intValue();
        return new DecimalType(precision, scale);
      case RECORD:
        return fromInternalSchema(field.getSchema());
      case MAP:
        InternalField key =
            field.getSchema().getFields().stream()
                .filter(
                    mapField ->
                        InternalField.Constants.MAP_KEY_FIELD_NAME.equals(mapField.getName()))
                .findFirst()
                .orElseThrow(() -> new IllegalStateException("Invalid map schema"));
        InternalField value =
            field.getSchema().getFields().stream()
                .filter(
                    mapField ->
                        InternalField.Constants.MAP_VALUE_FIELD_NAME.equals(mapField.getName()))
                .findFirst()
                .orElseThrow(() -> new IllegalStateException("Invalid map schema"));
        return new MapType(
            convertFieldType(key), convertFieldType(value), value.getSchema().isNullable());
      case LIST:
        InternalField element =
            field.getSchema().getFields().stream()
                .filter(
                    arrayField ->
                        InternalField.Constants.ARRAY_ELEMENT_FIELD_NAME.equals(
                            arrayField.getName()))
                .findFirst()
                .orElseThrow(() -> new IllegalStateException("Invalid array schema"));
        return new ArrayType(convertFieldType(element), element.getSchema().isNullable());
      default:
        throw new NotSupportedException("Unsupported type: " + field.getSchema().getDataType());
    }
  }

  /**
   * Creates Delta Kernel FieldMetadata from InternalField.
   *
   * @param field the internal field
   * @return Delta Kernel FieldMetadata
   */
  private FieldMetadata getFieldMetadata(InternalField field) {
    FieldMetadata.Builder metadataBuilder = FieldMetadata.builder();

    // Propagate column mapping ID to maintain Delta column mapping
    if (field.getFieldId() != null) {
      metadataBuilder.putLong(DELTA_COLUMN_MAPPING_ID, field.getFieldId().longValue());
    }

    InternalSchema schema = field.getSchema();

    // Handle UUID type
    InternalType type = schema.getDataType();
    if (type == InternalType.UUID) {
      metadataBuilder.putString(InternalSchema.XTABLE_LOGICAL_TYPE, "uuid");
    }

    // Handle comment
    if (schema.getComment() != null) {
      metadataBuilder.putString("comment", schema.getComment());
    }

    return metadataBuilder.build();
  }
}
