/*
 * 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.arrow.vector.types.pojo;

import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.io.File;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
import java.nio.channels.SeekableByteChannel;
import java.nio.channels.WritableByteChannel;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.nio.file.StandardOpenOption;
import java.util.Collections;
import java.util.UUID;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.memory.util.hash.ArrowBufHasher;
import org.apache.arrow.vector.ExtensionTypeVector;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.FixedSizeBinaryVector;
import org.apache.arrow.vector.Float4Vector;
import org.apache.arrow.vector.UuidVector;
import org.apache.arrow.vector.ValueIterableVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.compare.Range;
import org.apache.arrow.vector.compare.RangeEqualsVisitor;
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.ipc.ArrowFileReader;
import org.apache.arrow.vector.ipc.ArrowFileWriter;
import org.apache.arrow.vector.types.FloatingPointPrecision;
import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType;
import org.apache.arrow.vector.util.VectorBatchAppender;
import org.apache.arrow.vector.validate.ValidateVectorVisitor;
import org.junit.jupiter.api.Test;

public class TestExtensionType {
  /** Test that a custom UUID type can be round-tripped through a temporary file. */
  @Test
  public void roundtripUuid() throws IOException {
    ExtensionTypeRegistry.register(new UuidType());
    final Schema schema =
        new Schema(Collections.singletonList(Field.nullable("a", new UuidType())));
    try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
        final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
      UUID u1 = UUID.randomUUID();
      UUID u2 = UUID.randomUUID();
      UuidVector vector = (UuidVector) root.getVector("a");
      vector.setValueCount(2);
      vector.set(0, u1);
      vector.set(1, u2);
      root.setRowCount(2);

      final File file = File.createTempFile("uuidtest", ".arrow");
      try (final WritableByteChannel channel =
              FileChannel.open(Paths.get(file.getAbsolutePath()), StandardOpenOption.WRITE);
          final ArrowFileWriter writer = new ArrowFileWriter(root, null, channel)) {
        writer.start();
        writer.writeBatch();
        writer.end();
      }

      try (final SeekableByteChannel channel =
              Files.newByteChannel(Paths.get(file.getAbsolutePath()));
          final ArrowFileReader reader = new ArrowFileReader(channel, allocator)) {
        reader.loadNextBatch();
        final VectorSchemaRoot readerRoot = reader.getVectorSchemaRoot();
        assertEquals(root.getSchema(), readerRoot.getSchema());

        final Field field = readerRoot.getSchema().getFields().get(0);
        final UuidType expectedType = new UuidType();
        assertEquals(
            field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME),
            expectedType.extensionName());
        assertEquals(
            field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_METADATA),
            expectedType.serialize());

        final ExtensionTypeVector deserialized =
            (ExtensionTypeVector) readerRoot.getFieldVectors().get(0);
        assertEquals(vector.getValueCount(), deserialized.getValueCount());
        for (int i = 0; i < vector.getValueCount(); i++) {
          assertEquals(vector.isNull(i), deserialized.isNull(i));
          if (!vector.isNull(i)) {
            assertEquals(vector.getObject(i), deserialized.getObject(i));
          }
        }
      }
    }
  }

  /** Test that a custom UUID type can be read as its underlying type. */
  @Test
  public void readUnderlyingType() throws IOException {
    ExtensionTypeRegistry.register(new UuidType());
    final Schema schema =
        new Schema(Collections.singletonList(Field.nullable("a", new UuidType())));
    try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
        final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
      UUID u1 = UUID.randomUUID();
      UUID u2 = UUID.randomUUID();
      UuidVector vector = (UuidVector) root.getVector("a");
      vector.setValueCount(2);
      vector.set(0, u1);
      vector.set(1, u2);
      root.setRowCount(2);

      final File file = File.createTempFile("uuidtest", ".arrow");
      try (final WritableByteChannel channel =
              FileChannel.open(Paths.get(file.getAbsolutePath()), StandardOpenOption.WRITE);
          final ArrowFileWriter writer = new ArrowFileWriter(root, null, channel)) {
        writer.start();
        writer.writeBatch();
        writer.end();
      }

      ExtensionTypeRegistry.unregister(new UuidType());

      try (final SeekableByteChannel channel =
              Files.newByteChannel(Paths.get(file.getAbsolutePath()));
          final ArrowFileReader reader = new ArrowFileReader(channel, allocator)) {
        reader.loadNextBatch();
        final VectorSchemaRoot readerRoot = reader.getVectorSchemaRoot();
        assertEquals(1, readerRoot.getSchema().getFields().size());
        assertEquals("a", readerRoot.getSchema().getFields().get(0).getName());
        assertTrue(
            readerRoot.getSchema().getFields().get(0).getType()
                instanceof ArrowType.FixedSizeBinary);
        assertEquals(
            16,
            ((ArrowType.FixedSizeBinary) readerRoot.getSchema().getFields().get(0).getType())
                .getByteWidth());

        final Field field = readerRoot.getSchema().getFields().get(0);
        final UuidType expectedType = new UuidType();
        assertEquals(
            field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME),
            expectedType.extensionName());
        assertEquals(
            field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_METADATA),
            expectedType.serialize());

        final FixedSizeBinaryVector deserialized =
            (FixedSizeBinaryVector) readerRoot.getFieldVectors().get(0);
        assertEquals(vector.getValueCount(), deserialized.getValueCount());
        for (int i = 0; i < vector.getValueCount(); i++) {
          assertEquals(vector.isNull(i), deserialized.isNull(i));
          if (!vector.isNull(i)) {
            final UUID uuid = vector.getObject(i);
            final ByteBuffer bb = ByteBuffer.allocate(16);
            bb.putLong(uuid.getMostSignificantBits());
            bb.putLong(uuid.getLeastSignificantBits());
            assertArrayEquals(bb.array(), deserialized.get(i));
          }
        }
      }
    }
  }

  @Test
  public void testNullCheck() {
    NullPointerException e =
        assertThrows(
            NullPointerException.class,
            () -> {
              try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
                  final ExtensionTypeVector vector = new UuidVector("uuid", allocator, null)) {
                vector.getField();
                vector.allocateNewSafe();
              }
            });
    assertTrue(e.getMessage().contains("underlyingVector cannot be null."));
  }

  /** Test that a custom Location type can be round-tripped through a temporary file. */
  @Test
  public void roundtripLocation() throws IOException {
    ExtensionTypeRegistry.register(new LocationType());
    final Schema schema =
        new Schema(Collections.singletonList(Field.nullable("location", new LocationType())));
    try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
        final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
      LocationVector vector = (LocationVector) root.getVector("location");
      vector.allocateNew();
      vector.set(0, 34.073814f, -118.240784f);
      vector.set(2, 37.768056f, -122.3875f);
      vector.set(3, 40.739716f, -73.840782f);
      vector.setValueCount(4);
      root.setRowCount(4);

      final File file = File.createTempFile("locationtest", ".arrow");
      try (final WritableByteChannel channel =
              FileChannel.open(Paths.get(file.getAbsolutePath()), StandardOpenOption.WRITE);
          final ArrowFileWriter writer = new ArrowFileWriter(root, null, channel)) {
        writer.start();
        writer.writeBatch();
        writer.end();
      }

      try (final SeekableByteChannel channel =
              Files.newByteChannel(Paths.get(file.getAbsolutePath()));
          final ArrowFileReader reader = new ArrowFileReader(channel, allocator)) {
        reader.loadNextBatch();
        final VectorSchemaRoot readerRoot = reader.getVectorSchemaRoot();
        assertEquals(root.getSchema(), readerRoot.getSchema());

        final Field field = readerRoot.getSchema().getFields().get(0);
        final LocationType expectedType = new LocationType();
        assertEquals(
            field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME),
            expectedType.extensionName());
        assertEquals(
            field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_METADATA),
            expectedType.serialize());

        final ExtensionTypeVector deserialized =
            (ExtensionTypeVector) readerRoot.getFieldVectors().get(0);
        assertTrue(deserialized instanceof LocationVector);
        assertEquals("location", deserialized.getName());
        StructVector deserStruct = (StructVector) deserialized.getUnderlyingVector();
        assertNotNull(deserStruct.getChild("Latitude"));
        assertNotNull(deserStruct.getChild("Longitude"));
        assertEquals(vector.getValueCount(), deserialized.getValueCount());
        for (int i = 0; i < vector.getValueCount(); i++) {
          assertEquals(vector.isNull(i), deserialized.isNull(i));
          if (!vector.isNull(i)) {
            assertEquals(vector.getObject(i), deserialized.getObject(i));
          }
        }
      }
    }
  }

  @Test
  public void testVectorCompare() {
    UuidType uuidType = new UuidType();
    ExtensionTypeRegistry.register(uuidType);
    try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
        UuidVector a1 =
            (UuidVector) uuidType.getNewVector("a", FieldType.nullable(uuidType), allocator);
        UuidVector a2 =
            (UuidVector) uuidType.getNewVector("a", FieldType.nullable(uuidType), allocator);
        UuidVector bb =
            (UuidVector) uuidType.getNewVector("a", FieldType.nullable(uuidType), allocator)) {
      UUID u1 = UUID.randomUUID();
      UUID u2 = UUID.randomUUID();

      // Test out type and vector validation visitors for an ExtensionTypeVector
      ValidateVectorVisitor validateVisitor = new ValidateVectorVisitor();
      validateVisitor.visit(a1, null);

      a1.setValueCount(2);
      a1.set(0, u1);
      a1.set(1, u2);

      a2.setValueCount(2);
      a2.set(0, u1);
      a2.set(1, u2);

      bb.setValueCount(2);
      bb.set(0, u2);
      bb.set(1, u1);

      Range range = new Range(0, 0, a1.getValueCount());
      RangeEqualsVisitor visitor = new RangeEqualsVisitor(a1, a2);
      assertTrue(visitor.rangeEquals(range));

      visitor = new RangeEqualsVisitor(a1, bb);
      assertFalse(visitor.rangeEquals(range));

      // Test out vector appender
      VectorBatchAppender.batchAppend(a1, a2, bb);
      assertEquals(6, a1.getValueCount());
      validateVisitor.visit(a1, null);
    }
  }

  static class LocationType extends ExtensionType {

    @Override
    public ArrowType storageType() {
      return Struct.INSTANCE;
    }

    @Override
    public String extensionName() {
      return "location";
    }

    @Override
    public boolean extensionEquals(ExtensionType other) {
      return other instanceof LocationType;
    }

    @Override
    public ArrowType deserialize(ArrowType storageType, String serializedData) {
      if (!storageType.equals(storageType())) {
        throw new UnsupportedOperationException(
            "Cannot construct LocationType from underlying type " + storageType);
      }
      return new LocationType();
    }

    @Override
    public String serialize() {
      return "";
    }

    @Override
    public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) {
      return new LocationVector(name, allocator);
    }
  }

  public static class LocationVector extends ExtensionTypeVector<StructVector>
      implements ValueIterableVector<java.util.Map<String, ?>> {

    private static StructVector buildUnderlyingVector(String name, BufferAllocator allocator) {
      final StructVector underlyingVector =
          new StructVector(name, allocator, FieldType.nullable(ArrowType.Struct.INSTANCE), null);
      underlyingVector.addOrGet(
          "Latitude",
          FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)),
          Float4Vector.class);
      underlyingVector.addOrGet(
          "Longitude",
          FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)),
          Float4Vector.class);
      return underlyingVector;
    }

    public LocationVector(String name, BufferAllocator allocator) {
      super(name, allocator, buildUnderlyingVector(name, allocator));
    }

    @Override
    public int hashCode(int index) {
      return hashCode(index, null);
    }

    @Override
    public int hashCode(int index, ArrowBufHasher hasher) {
      return getUnderlyingVector().hashCode(index, hasher);
    }

    @Override
    public java.util.Map<String, ?> getObject(int index) {
      return getUnderlyingVector().getObject(index);
    }

    public void set(int index, float latitude, float longitude) {
      getUnderlyingVector().getChild("Latitude", Float4Vector.class).set(index, latitude);
      getUnderlyingVector().getChild("Longitude", Float4Vector.class).set(index, longitude);
      getUnderlyingVector().setIndexDefined(index);
    }
  }
}
