/*
 *  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.iceberg.service.dispatcher;

import java.io.IOException;
import java.time.Instant;
import java.util.Optional;
import org.apache.gravitino.Entity;
import org.apache.gravitino.EntityStore;
import org.apache.gravitino.GravitinoEnv;
import org.apache.gravitino.NameIdentifier;
import org.apache.gravitino.catalog.TableDispatcher;
import org.apache.gravitino.exceptions.NoSuchEntityException;
import org.apache.gravitino.iceberg.common.utils.IcebergIdentifierUtils;
import org.apache.gravitino.iceberg.service.authorization.IcebergRESTServerContext;
import org.apache.gravitino.listener.api.event.IcebergRequestContext;
import org.apache.gravitino.meta.AuditInfo;
import org.apache.gravitino.meta.TableEntity;
import org.apache.gravitino.utils.HierarchicalSchemaUtil;
import org.apache.gravitino.utils.PrincipalUtils;
import org.apache.iceberg.UpdateRequirement;
import org.apache.iceberg.catalog.Namespace;
import org.apache.iceberg.catalog.TableIdentifier;
import org.apache.iceberg.rest.requests.CreateTableRequest;
import org.apache.iceberg.rest.requests.PlanTableScanRequest;
import org.apache.iceberg.rest.requests.RenameTableRequest;
import org.apache.iceberg.rest.requests.UpdateTableRequest;
import org.apache.iceberg.rest.responses.ListTablesResponse;
import org.apache.iceberg.rest.responses.LoadCredentialsResponse;
import org.apache.iceberg.rest.responses.LoadTableResponse;
import org.apache.iceberg.rest.responses.PlanTableScanResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class IcebergTableHookDispatcher implements IcebergTableOperationDispatcher {

  private static final Logger LOG = LoggerFactory.getLogger(IcebergTableHookDispatcher.class);

  private final IcebergTableOperationDispatcher dispatcher;
  private final IcebergNamespaceOperationDispatcher namespaceDispatcher;
  private String metalake;

  public IcebergTableHookDispatcher(
      IcebergTableOperationDispatcher dispatcher,
      IcebergNamespaceOperationDispatcher namespaceDispatcher) {
    this.dispatcher = dispatcher;
    this.namespaceDispatcher = namespaceDispatcher;
    this.metalake = IcebergRESTServerContext.getInstance().metalakeName();
  }

  @Override
  public LoadTableResponse createTable(
      IcebergRequestContext context, Namespace namespace, CreateTableRequest createTableRequest) {
    LoadTableResponse response = dispatcher.createTable(context, namespace, createTableRequest);
    // Skip import and ownership for staged creates (credential vending).
    // The table will be imported when the staged table is committed via updateTable.
    if (!createTableRequest.stageCreate()) {
      importTableAndSetOwner(context, namespace, createTableRequest.name());
    }

    return response;
  }

  @Override
  public LoadTableResponse updateTable(
      IcebergRequestContext context,
      TableIdentifier tableIdentifier,
      UpdateTableRequest updateTableRequest) {
    LoadTableResponse response =
        dispatcher.updateTable(context, tableIdentifier, updateTableRequest);
    // Import the table and set ownership only when committing a staged table create.
    // A staged create commit is identified by the AssertTableDoesNotExist requirement,
    // which is set by UpdateRequirements.forCreateTable().
    boolean isStagedCreateCommit =
        updateTableRequest.requirements().stream()
            .anyMatch(UpdateRequirement.AssertTableDoesNotExist.class::isInstance);
    if (isStagedCreateCommit) {
      importTableAndSetOwner(context, tableIdentifier.namespace(), tableIdentifier.name());
    }
    return response;
  }

  @Override
  public void dropTable(
      IcebergRequestContext context, TableIdentifier tableIdentifier, boolean purgeRequested) {
    dispatcher.dropTable(context, tableIdentifier, purgeRequested);
    // Reconcile against Iceberg backend state — without a distributed TreeLock,
    // another node may recreate the same table between the drop above and the
    // EntityStore delete, leaving a stale Gravitino entity if we blindly delete.
    bestEffortReconcileTableEntity(context, tableIdentifier);
    IcebergOrphanSchemaCleanup.bestEffortCleanUp(
        metalake, namespaceDispatcher, context, tableIdentifier.namespace());
  }

  @Override
  public LoadTableResponse loadTable(
      IcebergRequestContext context, TableIdentifier tableIdentifier) {
    return dispatcher.loadTable(context, tableIdentifier);
  }

  @Override
  public ListTablesResponse listTable(IcebergRequestContext context, Namespace namespace) {
    return dispatcher.listTable(context, namespace);
  }

  @Override
  public boolean tableExists(IcebergRequestContext context, TableIdentifier tableIdentifier) {
    return dispatcher.tableExists(context, tableIdentifier);
  }

  @Override
  public void renameTable(IcebergRequestContext context, RenameTableRequest renameTableRequest) {
    dispatcher.renameTable(context, renameTableRequest);
    String separator = HierarchicalSchemaUtil.schemaSeparator();
    NameIdentifier tableSource =
        IcebergIdentifierUtils.toGravitinoTableIdentifier(
            metalake, context.catalogName(), renameTableRequest.source(), separator);
    NameIdentifier tableDest =
        IcebergIdentifierUtils.toGravitinoTableIdentifier(
            metalake, context.catalogName(), renameTableRequest.destination(), separator);
    EntityStore store = GravitinoEnv.getInstance().entityStore();
    try {
      if (store != null) {
        // Update the entity for the destination table.
        store.update(
            tableSource,
            TableEntity.class,
            Entity.EntityType.TABLE,
            tableEntity ->
                TableEntity.builder()
                    .withId(tableEntity.id())
                    .withName(tableDest.name())
                    .withNamespace(tableDest.namespace())
                    .withColumns(tableEntity.columns())
                    .withAuditInfo(
                        AuditInfo.builder()
                            .withCreator(tableEntity.auditInfo().creator())
                            .withCreateTime(tableEntity.auditInfo().createTime())
                            .withLastModifier(PrincipalUtils.getCurrentPrincipal().getName())
                            .withLastModifiedTime(Instant.now())
                            .build())
                    .build());
      }
    } catch (NoSuchEntityException ignore) {
      // Ignore if the source table entity does not exist.
    } catch (IOException ioe) {
      throw new RuntimeException("io exception when renaming table entity", ioe);
    }

    // IRC rename can race with another node's drop/create on either name.
    // Reconcile both ends against the Iceberg backend so we don't leave a
    // stale entity on the source or miss importing a re-created destination.
    bestEffortReconcileTableEntity(context, renameTableRequest.source());
    bestEffortReconcileTableEntity(context, renameTableRequest.destination());
  }

  @Override
  public LoadCredentialsResponse getTableCredentials(
      IcebergRequestContext context, TableIdentifier tableIdentifier) {
    return dispatcher.getTableCredentials(context, tableIdentifier);
  }

  /**
   * Plan table scan and return scan tasks.
   *
   * <p>This method performs server-side scan planning to optimize query performance by reducing
   * client-side metadata loading and enabling parallel task execution.
   *
   * @param context Iceberg REST request context information.
   * @param tableIdentifier The Iceberg table identifier.
   * @param scanRequest The scan request parameters including filters, projections, snapshot-id,
   *     etc.
   * @return A PlanTableScanResponse containing the scan plan with plan-id and tasks. The response
   *     format follows Iceberg REST API specification.
   */
  @Override
  public PlanTableScanResponse planTableScan(
      IcebergRequestContext context,
      TableIdentifier tableIdentifier,
      PlanTableScanRequest scanRequest) {
    return dispatcher.planTableScan(context, tableIdentifier, scanRequest);
  }

  @Override
  public Optional<String> getTableMetadataLocation(
      IcebergRequestContext context, TableIdentifier tableIdentifier) {
    return dispatcher.getTableMetadataLocation(context, tableIdentifier);
  }

  private void importTableAndSetOwner(
      IcebergRequestContext context, Namespace namespace, String tableName) {
    // Import is intentionally NOT wrapped in try-catch: if it fails the table exists in Iceberg
    // but not in Gravitino, and silently swallowing that would mislead callers into thinking the
    // entity is registered. Surface the failure so the caller can react.
    importTableEntity(context.catalogName(), namespace, tableName);
    IcebergOwnershipUtils.setTableOwner(
        metalake,
        context.catalogName(),
        namespace,
        tableName,
        context.userName(),
        GravitinoEnv.getInstance().internalOwnerDispatcher());
  }

  private void importTableEntity(String catalogName, Namespace namespace, String tableName) {
    TableDispatcher tableDispatcher = GravitinoEnv.getInstance().internalTableDispatcher();
    if (tableDispatcher != null) {
      tableDispatcher.loadTable(
          IcebergIdentifierUtils.toGravitinoTableIdentifier(
              metalake,
              catalogName,
              TableIdentifier.of(namespace, tableName),
              HierarchicalSchemaUtil.schemaSeparator()));
    }
  }

  private void reconcileTableEntity(
      IcebergRequestContext context, TableIdentifier tableIdentifier) {
    // IRC requests can be served by different Gravitino nodes. Without a distributed TreeLock,
    // another node may drop or recreate the same Iceberg table between the backend operation and
    // this hook's EntityStore mutation. Reconcile the local Gravitino entity with the Iceberg
    // backend state to avoid leaving stale/orphan table metadata in multi-node deployments.
    if (dispatcher.tableExists(context, tableIdentifier)) {
      importTableEntity(context.catalogName(), tableIdentifier.namespace(), tableIdentifier.name());
      return;
    }

    deleteTableEntity(context.catalogName(), tableIdentifier);

    if (dispatcher.tableExists(context, tableIdentifier)) {
      importTableEntity(context.catalogName(), tableIdentifier.namespace(), tableIdentifier.name());
    }
  }

  private void bestEffortReconcileTableEntity(
      IcebergRequestContext context, TableIdentifier tableIdentifier) {
    try {
      reconcileTableEntity(context, tableIdentifier);
    } catch (RuntimeException e) {
      LOG.warn(
          "Failed to reconcile Gravitino table entity after the Iceberg backend operation "
              + "succeeded. catalog={}, table={}",
          context.catalogName(),
          tableIdentifier,
          e);
    }
  }

  private void deleteTableEntity(String catalogName, TableIdentifier tableIdentifier) {
    EntityStore store = GravitinoEnv.getInstance().entityStore();
    try {
      if (store != null) {
        store.delete(
            IcebergIdentifierUtils.toGravitinoTableIdentifier(
                metalake, catalogName, tableIdentifier, HierarchicalSchemaUtil.schemaSeparator()),
            Entity.EntityType.TABLE);
      }
    } catch (NoSuchEntityException ignore) {
      // Ignore if the table entity does not exist.
    } catch (IOException ioe) {
      throw new RuntimeException("io exception when deleting table entity", ioe);
    }
  }
}
