/*
 * 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.amoro.server.optimizing;

import org.apache.amoro.ServerTableIdentifier;
import org.apache.amoro.optimizing.MetricsSummary;
import org.apache.amoro.optimizing.OptimizingType;
import org.apache.amoro.process.ProcessStatus;
import org.apache.amoro.server.persistence.PersistentBase;
import org.apache.amoro.server.persistence.mapper.TableMetaMapper;
import org.apache.amoro.server.persistence.mapper.TableProcessMapper;
import org.apache.amoro.server.process.TableProcessMeta;
import org.apache.amoro.shade.guava32.com.google.common.collect.Sets;
import org.apache.amoro.table.TableIdentifier;
import org.apache.iceberg.data.Record;
import org.junit.Assert;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.TimeoutException;
import java.util.function.Supplier;
import java.util.stream.Collectors;

public class BaseOptimizingChecker extends PersistentBase {

  private static final Logger LOG = LoggerFactory.getLogger(BaseOptimizingChecker.class);
  private static final long WAIT_SUCCESS_TIMEOUT = 150_000;
  private static final long CHECK_TIMEOUT = 1_000;
  private final TableIdentifier tableIdentifier;
  private long lastProcessId;

  public BaseOptimizingChecker(TableIdentifier tableIdentifier) {
    this.tableIdentifier = tableIdentifier;
    this.lastProcessId = System.currentTimeMillis();
  }

  protected void assertIds(List<Record> actualRows, Object... expectIds) {
    assertRecordValues(actualRows, 0, expectIds);
  }

  protected void assertIdRange(List<Record> actualRows, int from, int to) {
    assertRecordValues(actualRows, 0, range(from, to).toArray());
  }

  protected List<Integer> range(int from, int to) {
    List<Integer> ids = new ArrayList<>();
    for (int i = from; i <= to; i++) {
      ids.add(i);
    }
    return ids;
  }

  protected void assertNames(List<Record> actualRows, Object... expectNames) {
    assertRecordValues(actualRows, 1, expectNames);
  }

  protected void assertRecordValues(List<Record> actualRows, int index, Object... expectValues) {
    Set<Object> actualValueSet = Sets.newHashSet();
    int cnt = 0;
    for (Record r : actualRows) {
      actualValueSet.add(r.get(index));
      cnt++;
    }
    Assert.assertEquals(cnt, expectValues.length);
    for (Object id : expectValues) {
      if (!actualValueSet.contains(id)) {
        throw new AssertionError("assert id contain " + id + ", but not found");
      }
    }
  }

  protected void assertOptimizingProcess(
      TableProcessMeta optimizingProcess,
      OptimizingType optimizeType,
      int fileCntBefore,
      int fileCntAfter) {
    Assert.assertNotNull(optimizingProcess);
    Assert.assertEquals(optimizeType.name(), optimizingProcess.getProcessType());
    MetricsSummary summary = MetricsSummary.fromMap(optimizingProcess.getSummary());
    Assert.assertEquals(
        fileCntBefore,
        summary.getRewritePosDataFileCnt()
            + summary.getRewriteDataFileCnt()
            + summary.getEqDeleteFileCnt()
            + summary.getPosDeleteFileCnt());
    Assert.assertEquals(fileCntAfter, summary.getNewDataFileCnt() + summary.getNewDeleteFileCnt());
  }

  protected TableProcessMeta waitOptimizeResult() {
    boolean success;
    ServerTableIdentifier identifier =
        getAs(
            TableMetaMapper.class,
            m ->
                m.selectTableIdentifier(
                    tableIdentifier.getCatalog(),
                    tableIdentifier.getDatabase(),
                    tableIdentifier.getTableName()));
    try {
      success =
          waitUntilFinish(
              () -> {
                List<TableProcessMeta> tableOptimizingProcesses =
                    getAs(
                        TableProcessMapper.class,
                        mapper -> mapper.listProcessMeta(identifier.getId(), null, null));
                if (tableOptimizingProcesses == null || tableOptimizingProcesses.isEmpty()) {
                  LOG.info("optimize history is empty");
                  return Status.RUNNING;
                }
                Optional<TableProcessMeta> any =
                    tableOptimizingProcesses.stream()
                        .filter(p -> p.getProcessId() > lastProcessId)
                        .filter(p -> p.getStatus().equals(ProcessStatus.SUCCESS))
                        .findAny();

                if (any.isPresent()) {
                  return Status.SUCCESS;
                } else {
                  LOG.info(
                      "optimize max process id {}",
                      tableOptimizingProcesses.stream()
                          .map(TableProcessMeta::getProcessId)
                          .max(Comparator.naturalOrder())
                          .get());
                  return Status.RUNNING;
                }
              },
              WAIT_SUCCESS_TIMEOUT);
    } catch (TimeoutException e) {
      throw new IllegalStateException(
          "wait optimize result timeout expectRecordId " + lastProcessId, e);
    }

    if (success) {
      List<TableProcessMeta> result =
          getAs(
                  TableProcessMapper.class,
                  mapper -> mapper.listProcessMeta(identifier.getId(), null, null))
              .stream()
              .filter(p -> p.getProcessId() > lastProcessId)
              .filter(p -> p.getStatus().equals(ProcessStatus.SUCCESS))
              .collect(Collectors.toList());
      if (result.size() == 1) {
        this.lastProcessId = result.get(0).getProcessId();
        return result.get(0);
      } else {
        throw new RuntimeException("optimize result size " + result.size());
      }
    } else {
      return null;
    }
  }

  protected void assertOptimizeHangUp() {
    try {
      Thread.sleep(CHECK_TIMEOUT);
    } catch (InterruptedException e) {
      throw new IllegalStateException("waiting result was interrupted");
    }
    ServerTableIdentifier identifier =
        getAs(
            TableMetaMapper.class,
            m ->
                m.selectTableIdentifier(
                    tableIdentifier.getCatalog(),
                    tableIdentifier.getDatabase(),
                    tableIdentifier.getTableName()));
    List<TableProcessMeta> tableOptimizingProcesses =
        getAs(
                TableProcessMapper.class,
                mapper -> mapper.listProcessMeta(identifier.getId(), null, null))
            .stream()
            .filter(p -> p.getProcessId() > lastProcessId)
            .collect(Collectors.toList());
    Assert.assertTrue("optimize is not stopped", tableOptimizingProcesses.isEmpty());
  }

  protected boolean waitUntilFinish(Supplier<Status> statusSupplier, final long timeout)
      throws TimeoutException {
    long start = System.currentTimeMillis();
    while (true) {
      long duration = System.currentTimeMillis() - start;
      if (duration > timeout) {
        throw new TimeoutException("wait exceed timeout, " + duration + "ms > " + timeout + "ms");
      }
      Status status = statusSupplier.get();
      if (status == Status.FAILED) {
        return false;
      } else if (status == Status.RUNNING) {
        try {
          Thread.sleep(1000);
        } catch (InterruptedException e) {
          throw new IllegalStateException("waiting result was interrupted");
        }
      } else {
        return true;
      }
    }
  }

  private enum Status {
    SUCCESS,
    FAILED,
    RUNNING
  }
}
