/*
 * 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.datafusion.proto;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import org.apache.arrow.vector.types.DateUnit;
import org.apache.arrow.vector.types.FloatingPointPrecision;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.types.pojo.Schema;

import datafusion_common.DatafusionCommon;

/**
 * Convert between Arrow Java {@link Schema} and the {@code datafusion_common.Schema} protobuf shape
 * used by DataFusion plan messages such as {@code ListingTableScanNode.schema}.
 *
 * <p>Supports the primitive Arrow types this project's tests exercise (Bool, signed/unsigned Int
 * 8..64, Float32/64, Utf8, Utf8View, LargeUtf8, Date32, Decimal128). Anything else raises {@link
 * UnsupportedOperationException} with a message naming the offending type.
 */
public final class SchemaConverter {

  private SchemaConverter() {}

  public static DatafusionCommon.Schema toProto(Schema arrow) {
    DatafusionCommon.Schema.Builder builder = DatafusionCommon.Schema.newBuilder();
    for (Field f : arrow.getFields()) {
      builder.addColumns(fieldToProto(f));
    }
    if (arrow.getCustomMetadata() != null) {
      builder.putAllMetadata(arrow.getCustomMetadata());
    }
    return builder.build();
  }

  public static Schema fromProto(DatafusionCommon.Schema proto) {
    List<Field> fields = new ArrayList<>(proto.getColumnsCount());
    for (DatafusionCommon.Field f : proto.getColumnsList()) {
      fields.add(fieldFromProto(f));
    }
    Map<String, String> metadata = new LinkedHashMap<>(proto.getMetadataMap());
    return new Schema(fields, metadata);
  }

  static DatafusionCommon.Field fieldToProto(Field f) {
    DatafusionCommon.Field.Builder b =
        DatafusionCommon.Field.newBuilder()
            .setName(f.getName())
            .setArrowType(arrowTypeToProto(f.getType()))
            .setNullable(f.isNullable());
    if (f.getMetadata() != null) {
      b.putAllMetadata(f.getMetadata());
    }
    return b.build();
  }

  static Field fieldFromProto(DatafusionCommon.Field f) {
    ArrowType type = arrowTypeFromProto(f.getArrowType());
    FieldType ft = new FieldType(f.getNullable(), type, null, f.getMetadataMap());
    return new Field(f.getName(), ft, null);
  }

  static DatafusionCommon.ArrowType arrowTypeToProto(ArrowType t) {
    DatafusionCommon.EmptyMessage empty = DatafusionCommon.EmptyMessage.getDefaultInstance();
    DatafusionCommon.ArrowType.Builder b = DatafusionCommon.ArrowType.newBuilder();

    if (t instanceof ArrowType.Bool) {
      return b.setBOOL(empty).build();
    }
    if (t instanceof ArrowType.Int) {
      ArrowType.Int i = (ArrowType.Int) t;
      if (i.getIsSigned()) {
        switch (i.getBitWidth()) {
          case 8:
            return b.setINT8(empty).build();
          case 16:
            return b.setINT16(empty).build();
          case 32:
            return b.setINT32(empty).build();
          case 64:
            return b.setINT64(empty).build();
          default:
            throw new UnsupportedOperationException(
                "Arrow type Int signed width "
                    + i.getBitWidth()
                    + " not yet supported by SchemaConverter");
        }
      } else {
        switch (i.getBitWidth()) {
          case 8:
            return b.setUINT8(empty).build();
          case 16:
            return b.setUINT16(empty).build();
          case 32:
            return b.setUINT32(empty).build();
          case 64:
            return b.setUINT64(empty).build();
          default:
            throw new UnsupportedOperationException(
                "Arrow type Int unsigned width "
                    + i.getBitWidth()
                    + " not yet supported by SchemaConverter");
        }
      }
    }
    if (t instanceof ArrowType.FloatingPoint) {
      FloatingPointPrecision p = ((ArrowType.FloatingPoint) t).getPrecision();
      if (p == FloatingPointPrecision.SINGLE) return b.setFLOAT32(empty).build();
      if (p == FloatingPointPrecision.DOUBLE) return b.setFLOAT64(empty).build();
      throw new UnsupportedOperationException(
          "Arrow type FloatingPoint " + p + " not yet supported by SchemaConverter");
    }
    if (t instanceof ArrowType.Utf8) {
      return b.setUTF8(empty).build();
    }
    if (t instanceof ArrowType.Utf8View) {
      return b.setUTF8VIEW(empty).build();
    }
    if (t instanceof ArrowType.LargeUtf8) {
      return b.setLARGEUTF8(empty).build();
    }
    if (t instanceof ArrowType.Date) {
      DateUnit u = ((ArrowType.Date) t).getUnit();
      if (u == DateUnit.DAY) return b.setDATE32(empty).build();
      throw new UnsupportedOperationException(
          "Arrow type Date " + u + " not yet supported by SchemaConverter");
    }
    if (t instanceof ArrowType.Decimal) {
      ArrowType.Decimal d = (ArrowType.Decimal) t;
      if (d.getBitWidth() != 128) {
        throw new UnsupportedOperationException(
            "Arrow type Decimal bit width "
                + d.getBitWidth()
                + " not yet supported by SchemaConverter");
      }
      return b.setDECIMAL128(
              DatafusionCommon.Decimal128Type.newBuilder()
                  .setPrecision(d.getPrecision())
                  .setScale(d.getScale())
                  .build())
          .build();
    }
    throw new UnsupportedOperationException(
        "Arrow type " + t.getClass().getSimpleName() + " not yet supported by SchemaConverter");
  }

  static ArrowType arrowTypeFromProto(DatafusionCommon.ArrowType p) {
    switch (p.getArrowTypeEnumCase()) {
      case BOOL:
        return ArrowType.Bool.INSTANCE;
      case INT8:
        return new ArrowType.Int(8, true);
      case INT16:
        return new ArrowType.Int(16, true);
      case INT32:
        return new ArrowType.Int(32, true);
      case INT64:
        return new ArrowType.Int(64, true);
      case UINT8:
        return new ArrowType.Int(8, false);
      case UINT16:
        return new ArrowType.Int(16, false);
      case UINT32:
        return new ArrowType.Int(32, false);
      case UINT64:
        return new ArrowType.Int(64, false);
      case FLOAT32:
        return new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE);
      case FLOAT64:
        return new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE);
      case UTF8:
        return ArrowType.Utf8.INSTANCE;
      case UTF8_VIEW:
        return ArrowType.Utf8View.INSTANCE;
      case LARGE_UTF8:
        return ArrowType.LargeUtf8.INSTANCE;
      case DATE32:
        return new ArrowType.Date(DateUnit.DAY);
      case DECIMAL128:
        DatafusionCommon.Decimal128Type d = p.getDECIMAL128();
        return new ArrowType.Decimal(d.getPrecision(), d.getScale(), 128);
      default:
        throw new UnsupportedOperationException(
            "datafusion_common.ArrowType "
                + p.getArrowTypeEnumCase()
                + " not yet supported by SchemaConverter");
    }
  }
}
