/*
 * 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.iotdb.db.schemaengine.schemaregion.attribute.update;

import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
import org.apache.iotdb.common.rpc.thrift.TSStatus;
import org.apache.iotdb.commons.client.exception.ClientManagerException;
import org.apache.iotdb.commons.client.request.AsyncRequestContext;
import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory;
import org.apache.iotdb.commons.concurrent.ThreadName;
import org.apache.iotdb.commons.conf.CommonDescriptor;
import org.apache.iotdb.commons.consensus.SchemaRegionId;
import org.apache.iotdb.commons.exception.StartupException;
import org.apache.iotdb.commons.service.AbstractPeriodicalServiceWithAdvance;
import org.apache.iotdb.commons.service.IService;
import org.apache.iotdb.commons.service.ServiceType;
import org.apache.iotdb.commons.utils.PathUtils;
import org.apache.iotdb.confignode.rpc.thrift.TShowClusterResp;
import org.apache.iotdb.db.conf.IoTDBConfig;
import org.apache.iotdb.db.conf.IoTDBDescriptor;
import org.apache.iotdb.db.protocol.client.ConfigNodeClient;
import org.apache.iotdb.db.protocol.client.ConfigNodeClientManager;
import org.apache.iotdb.db.protocol.client.ConfigNodeInfo;
import org.apache.iotdb.db.protocol.client.dn.DnToDnInternalServiceAsyncRequestManager;
import org.apache.iotdb.db.protocol.client.dn.DnToDnRequestType;
import org.apache.iotdb.db.queryengine.execution.executor.RegionWriteExecutor;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeId;
import org.apache.iotdb.db.queryengine.plan.relational.planner.node.schema.TableDeviceAttributeCommitUpdateNode;
import org.apache.iotdb.db.schemaengine.schemaregion.ISchemaRegion;
import org.apache.iotdb.db.schemaengine.schemaregion.impl.SchemaRegionMemoryImpl;
import org.apache.iotdb.mpp.rpc.thrift.TAttributeUpdateReq;
import org.apache.iotdb.mpp.rpc.thrift.TSchemaRegionAttributeInfo;
import org.apache.iotdb.rpc.TSStatusCode;

import org.apache.thrift.TException;
import org.apache.tsfile.utils.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;

import java.nio.ByteBuffer;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

public class GeneralRegionAttributeSecurityService extends AbstractPeriodicalServiceWithAdvance
    implements IService {
  private static final Logger LOGGER =
      LoggerFactory.getLogger(GeneralRegionAttributeSecurityService.class);

  private static final IoTDBConfig iotdbConfig = IoTDBDescriptor.getInstance().getConfig();
  private final Map<Integer, Pair<Long, Integer>> dataNodeId2FailureDurationAndTimesMap =
      new HashMap<>();
  private final Set<ISchemaRegion> regionLeaders =
      Collections.newSetFromMap(new ConcurrentHashMap<>());
  private final Map<SchemaRegionId, String> regionId2DatabaseMap = new ConcurrentHashMap<>();

  public void startBroadcast(final ISchemaRegion schemaRegion) {
    if (schemaRegion instanceof SchemaRegionMemoryImpl
        && PathUtils.isTableModelDatabase(schemaRegion.getDatabaseFullPath())) {
      regionId2DatabaseMap.put(
          schemaRegion.getSchemaRegionId(), schemaRegion.getDatabaseFullPath());
      regionLeaders.add(schemaRegion);
    }
  }

  public void stopBroadcast(final ISchemaRegion schemaRegion) {
    regionLeaders.remove(schemaRegion);
    // Reserve the database info for concurrency simplicity
  }

  @Override
  protected void executeTask() {
    // All the "detailContainer"'s size will add up to at most "limit"
    // UpdateClearContainer and version / TEndPoint are not calculated
    final AtomicInteger limit =
        new AtomicInteger(
            CommonDescriptor.getInstance()
                .getConfig()
                .getPipeConnectorRequestSliceThresholdBytes());

    final AtomicBoolean hasRemaining = new AtomicBoolean(false);
    final Map<SchemaRegionId, Pair<Long, Map<TDataNodeLocation, byte[]>>> attributeUpdateCommitMap =
        new HashMap<>();
    for (final ISchemaRegion regionLeader : regionLeaders) {
      final Pair<Long, Map<TDataNodeLocation, byte[]>> currentResult =
          regionLeader.getAttributeUpdateInfo(limit, hasRemaining);
      if (currentResult.getRight().isEmpty()) {
        break;
      }
      attributeUpdateCommitMap.put(regionLeader.getSchemaRegionId(), currentResult);
    }

    if (hasRemaining.get()) {
      skipNextSleep = true;
    }

    if (!attributeUpdateCommitMap.isEmpty()) {
      // Send & may shrink
      final Map<SchemaRegionId, Set<TDataNodeLocation>> shrinkMap =
          sendUpdateRequestAndMayShrink(attributeUpdateCommitMap);

      // Commit
      attributeUpdateCommitMap.forEach(
          (schemaRegionId, pair) -> {
            if (!new RegionWriteExecutor()
                .execute(
                    schemaRegionId,
                    new TableDeviceAttributeCommitUpdateNode(
                        new PlanNodeId(""),
                        pair.getLeft(),
                        pair.getRight(),
                        shrinkMap.getOrDefault(schemaRegionId, Collections.emptySet()),
                        new TDataNodeLocation(
                            iotdbConfig.getDataNodeId(),
                            null,
                            new TEndPoint(
                                iotdbConfig.getInternalAddress(), iotdbConfig.getInternalPort()),
                            null,
                            null,
                            null)))
                .isAccepted()) {
              // May fail due to region shutdown, migration or other reasons
              // Just ignore
              skipNextSleep = false;
              LOGGER.warn("Failed to write attribute commit message to region {}.", schemaRegionId);
            }
          });
    }
  }

  // Shrink the locations which do not exist in cluster dataNodes
  private @Nonnull Map<SchemaRegionId, Set<TDataNodeLocation>> detectNodeShrinkage(
      final Map<SchemaRegionId, Pair<Long, Map<TDataNodeLocation, byte[]>>>
          attributeUpdateCommitMap) {
    final TShowClusterResp showClusterResp;
    try (final ConfigNodeClient client =
        ConfigNodeClientManager.getInstance().borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
      showClusterResp = client.showCluster();
    } catch (final ClientManagerException | TException e) {
      LOGGER.warn("Failed to fetch dataNodeLocations, will retry.");
      return Collections.emptyMap();
    }
    dataNodeId2FailureDurationAndTimesMap.clear();
    final Set<TDataNodeLocation> realDataNodeLocations = new HashSet<>();
    showClusterResp
        .getDataNodeList()
        .forEach(
            location -> {
              location.setDataRegionConsensusEndPoint(null);
              location.setMPPDataExchangeEndPoint(null);
              location.setSchemaRegionConsensusEndPoint(null);
              location.setClientRpcEndPoint(null);
              realDataNodeLocations.add(location);
            });
    return attributeUpdateCommitMap.entrySet().stream()
        .filter(
            entry -> {
              entry.getValue().getRight().keySet().removeIf(realDataNodeLocations::contains);
              return !entry.getValue().getRight().isEmpty();
            })
        .collect(
            Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().getRight().keySet()));
  }

  private @Nonnull Map<SchemaRegionId, Set<TDataNodeLocation>> sendUpdateRequestAndMayShrink(
      final Map<SchemaRegionId, Pair<Long, Map<TDataNodeLocation, byte[]>>>
          attributeUpdateCommitMap) {
    final AsyncRequestContext<TAttributeUpdateReq, TSStatus, DnToDnRequestType, TDataNodeLocation>
        clientHandler = new AsyncRequestContext<>(DnToDnRequestType.UPDATE_ATTRIBUTE);

    attributeUpdateCommitMap.forEach(
        (id, pair) ->
            pair.getRight()
                .forEach(
                    (location, bytes) -> {
                      clientHandler.putNodeLocation(location.getDataNodeId(), location);
                      clientHandler.putRequestIfAbsent(
                          location.getDataNodeId(), new TAttributeUpdateReq(new HashMap<>()));
                      clientHandler
                          .getRequest(location.getDataNodeId())
                          .getAttributeUpdateMap()
                          .put(
                              id.getId(),
                              new TSchemaRegionAttributeInfo(
                                  pair.getLeft(),
                                  regionId2DatabaseMap.get(id),
                                  ByteBuffer.wrap(bytes)));
                    }));

    DnToDnInternalServiceAsyncRequestManager.getInstance()
        .sendAsyncRequestWithTimeoutInMs(
            clientHandler,
            IoTDBDescriptor.getInstance()
                .getConfig()
                .getGeneralRegionAttributeSecurityServiceTimeoutSeconds());

    final AtomicBoolean needFetch = new AtomicBoolean(false);

    final Set<Integer> failedDataNodes =
        clientHandler.getResponseMap().entrySet().stream()
            .filter(
                entry -> {
                  final boolean failed =
                      entry.getValue().getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode();
                  if (failed) {
                    dataNodeId2FailureDurationAndTimesMap.compute(
                        entry.getKey(),
                        (k, v) -> {
                          if (Objects.isNull(v)) {
                            return new Pair<>(System.currentTimeMillis(), 1);
                          }
                          v.setRight(v.getRight() + 1);
                          if (System.currentTimeMillis() - v.getLeft()
                                  >= iotdbConfig
                                      .getGeneralRegionAttributeSecurityServiceFailureDurationSecondsToFetch()
                              || v.getRight()
                                  >= iotdbConfig
                                      .getGeneralRegionAttributeSecurityServiceFailureTimesToFetch()) {
                            needFetch.set(true);
                          }
                          return v;
                        });
                  } else {
                    dataNodeId2FailureDurationAndTimesMap.remove(entry.getKey());
                  }
                  return failed;
                })
            .map(Map.Entry::getKey)
            .collect(Collectors.toSet());

    // Compute node shrinkage before failure removal in commit
    final Map<SchemaRegionId, Set<TDataNodeLocation>> result =
        needFetch.get() ? detectNodeShrinkage(attributeUpdateCommitMap) : Collections.emptyMap();

    for (final Iterator<Map.Entry<SchemaRegionId, Pair<Long, Map<TDataNodeLocation, byte[]>>>> it =
            attributeUpdateCommitMap.entrySet().iterator();
        it.hasNext(); ) {
      final Map.Entry<SchemaRegionId, Pair<Long, Map<TDataNodeLocation, byte[]>>> currentEntry =
          it.next();
      final Map<TDataNodeLocation, byte[]> dataNodeLocationMap = currentEntry.getValue().getRight();
      dataNodeLocationMap
          .entrySet()
          .removeIf(
              locationEntry -> failedDataNodes.contains(locationEntry.getKey().getDataNodeId()));
      // Reserve the schemaIds with shrinkage for commit convenience
      if (dataNodeLocationMap.isEmpty() && !result.containsKey(currentEntry.getKey())) {
        it.remove();
      }
    }
    return result;
  }

  /////////////////////////////// IService ///////////////////////////////

  @Override
  public void start() throws StartupException {
    startService();
  }

  @Override
  public void stop() {
    stopService();
    securityServiceExecutor.shutdown();
  }

  @Override
  public ServiceType getID() {
    return ServiceType.GENERAL_REGION_ATTRIBUTE_SECURITY_SERVICE;
  }

  /////////////////////////////// SingleTon ///////////////////////////////

  private GeneralRegionAttributeSecurityService() {
    super(
        IoTDBThreadPoolFactory.newSingleThreadExecutor(
            ThreadName.GENERAL_REGION_ATTRIBUTE_SECURITY_SERVICE.getName()),
        iotdbConfig.getGeneralRegionAttributeSecurityServiceIntervalSeconds() * 1000L);
  }

  private static final class GeneralRegionAttributeSecurityServiceHolder {
    private static final GeneralRegionAttributeSecurityService INSTANCE =
        new GeneralRegionAttributeSecurityService();

    private GeneralRegionAttributeSecurityServiceHolder() {}
  }

  public static GeneralRegionAttributeSecurityService getInstance() {
    return GeneralRegionAttributeSecurityServiceHolder.INSTANCE;
  }
}
