/*
 * 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.gravitino.trino.connector;

import io.airlift.slice.Slice;
import io.trino.spi.connector.ColumnMetadata;
import io.trino.spi.connector.ConnectorInsertTableHandle;
import io.trino.spi.connector.ConnectorMergeTableHandle;
import io.trino.spi.connector.ConnectorOutputMetadata;
import io.trino.spi.connector.ConnectorOutputTableHandle;
import io.trino.spi.connector.ConnectorSession;
import io.trino.spi.connector.ConnectorTableExecuteHandle;
import io.trino.spi.connector.ConnectorTableHandle;
import io.trino.spi.connector.RetryMode;
import io.trino.spi.connector.SchemaTableName;
import io.trino.spi.statistics.ComputedStatistics;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;
import org.apache.gravitino.trino.connector.catalog.CatalogConnectorMetadata;
import org.apache.gravitino.trino.connector.catalog.CatalogConnectorMetadataAdapter;
import org.apache.gravitino.trino.connector.metadata.GravitinoColumn;

public class GravitinoMetadata446 extends GravitinoMetadata {

  public GravitinoMetadata446(
      CatalogConnectorMetadata catalogConnectorMetadata,
      CatalogConnectorMetadataAdapter metadataAdapter,
      io.trino.spi.connector.ConnectorMetadata internalMetadata) {
    super(catalogConnectorMetadata, metadataAdapter, internalMetadata);
  }

  @Override
  public void addColumn(
      ConnectorSession session, ConnectorTableHandle tableHandle, ColumnMetadata column) {
    GravitinoColumn gravitinoColumn = metadataAdapter.createColumn(column);
    catalogConnectorMetadata.addColumn(getTableName(tableHandle), gravitinoColumn);
  }

  @Override
  public Optional<ConnectorTableExecuteHandle> getTableHandleForExecute(
      ConnectorSession session,
      ConnectorTableHandle tableHandle,
      String procedureName,
      Map<String, Object> executeProperties,
      RetryMode retryMode) {
    return internalMetadata
        .getTableHandleForExecute(
            session,
            GravitinoHandle.unWrap(tableHandle),
            procedureName,
            executeProperties,
            retryMode)
        .map(GravitinoTableExecuteHandle::new);
  }

  @Override
  public void executeTableExecute(
      ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle) {
    internalMetadata.executeTableExecute(session, GravitinoHandle.unWrap(tableExecuteHandle));
  }

  @Override
  public Optional<ConnectorOutputMetadata> finishInsert(
      ConnectorSession session,
      ConnectorInsertTableHandle insertHandle,
      List<ConnectorTableHandle> sourceTableHandles,
      Collection<Slice> fragments,
      Collection<ComputedStatistics> computedStatistics) {
    return internalMetadata.finishInsert(
        session,
        GravitinoHandle.unWrap(insertHandle),
        sourceTableHandles.stream().map(GravitinoHandle::unWrap).collect(Collectors.toList()),
        fragments,
        computedStatistics);
  }

  @Override
  public Optional<ConnectorOutputMetadata> finishCreateTable(
      ConnectorSession session,
      ConnectorOutputTableHandle tableHandle,
      Collection<Slice> fragments,
      Collection<ComputedStatistics> computedStatistics) {
    ConnectorInsertTableHandle insertHandle =
        ((GravitinoOutputTableHandle) tableHandle).getInternalHandle();
    return internalMetadata.finishInsert(
        session, insertHandle, List.of(), fragments, computedStatistics);
  }

  @Override
  public ConnectorMergeTableHandle beginMerge(
      ConnectorSession session, ConnectorTableHandle tableHandle, RetryMode retryMode) {
    ConnectorMergeTableHandle connectorMergeTableHandle =
        internalMetadata.beginMerge(session, GravitinoHandle.unWrap(tableHandle), retryMode);
    SchemaTableName tableName = getTableName(tableHandle);

    return new GravitinoMergeTableHandle(
        tableName.getSchemaName(), tableName.getTableName(), connectorMergeTableHandle);
  }

  @Override
  public void finishMerge(
      ConnectorSession session,
      ConnectorMergeTableHandle mergeTableHandle,
      Collection<Slice> fragments,
      Collection<ComputedStatistics> computedStatistics) {
    internalMetadata.finishMerge(
        session, GravitinoHandle.unWrap(mergeTableHandle), fragments, computedStatistics);
  }
}
