/*
 * 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.confignode.manager.pipe.agent.task;

import org.apache.iotdb.commons.exception.pipe.PipeRuntimeException;
import org.apache.iotdb.commons.pipe.agent.plugin.builtin.BuiltinPipePlugin;
import org.apache.iotdb.commons.pipe.agent.task.meta.PipeTaskMeta;
import org.apache.iotdb.commons.pipe.agent.task.progress.PipeEventCommitManager;
import org.apache.iotdb.commons.pipe.agent.task.subtask.PipeAbstractSinkSubtask;
import org.apache.iotdb.commons.pipe.config.constant.PipeProcessorConstant;
import org.apache.iotdb.commons.pipe.config.plugin.configuraion.PipeTaskRuntimeConfiguration;
import org.apache.iotdb.commons.pipe.config.plugin.env.PipeTaskProcessorRuntimeEnvironment;
import org.apache.iotdb.commons.pipe.config.plugin.env.PipeTaskSinkRuntimeEnvironment;
import org.apache.iotdb.commons.pipe.config.plugin.env.PipeTaskSourceRuntimeEnvironment;
import org.apache.iotdb.commons.pipe.event.EnrichedEvent;
import org.apache.iotdb.commons.pipe.event.ProgressReportEvent;
import org.apache.iotdb.confignode.manager.pipe.agent.PipeConfigNodeAgent;
import org.apache.iotdb.confignode.manager.pipe.metric.sink.PipeConfigRegionSinkMetrics;
import org.apache.iotdb.confignode.manager.pipe.source.IoTDBConfigRegionSource;
import org.apache.iotdb.pipe.api.PipeExtractor;
import org.apache.iotdb.pipe.api.PipeProcessor;
import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameterValidator;
import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameters;
import org.apache.iotdb.pipe.api.event.Event;
import org.apache.iotdb.pipe.api.exception.PipeException;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Map;

import static org.apache.iotdb.db.protocol.client.ConfigNodeInfo.CONFIG_REGION_ID;

public class PipeConfigNodeSubtask extends PipeAbstractSinkSubtask {

  private static final Logger LOGGER = LoggerFactory.getLogger(PipeConfigNodeSubtask.class);

  private final String pipeName;
  private final PipeTaskMeta pipeTaskMeta;

  // Pipe plugins for this subtask
  private PipeExtractor source;

  // TODO: currently unused
  @SuppressWarnings("unused")
  private PipeProcessor processor;

  public PipeConfigNodeSubtask(
      final String pipeName,
      final long creationTime,
      final Map<String, String> sourceAttributes,
      final Map<String, String> processorAttributes,
      final Map<String, String> sinkAttributes,
      final PipeTaskMeta pipeTaskMeta)
      throws Exception {
    // We initialize outputPipeConnector by initConnector()
    super(pipeName + "_" + creationTime, creationTime, null);
    this.pipeName = pipeName;
    this.pipeTaskMeta = pipeTaskMeta;

    initSource(sourceAttributes);
    initProcessor(processorAttributes);
    initSink(sinkAttributes);

    PipeConfigRegionSinkMetrics.getInstance().register(this);
    PipeEventCommitManager.getInstance()
        .register(pipeName, creationTime, CONFIG_REGION_ID.getId(), pipeName + "_" + creationTime);
  }

  private void initSource(final Map<String, String> sourceAttributes) throws Exception {
    final PipeParameters sourceParameters = new PipeParameters(sourceAttributes);

    // 1. Construct extractor
    source = PipeConfigNodeAgent.plugin().reflectSource(sourceParameters);

    try {
      // 2. Validate extractor parameters
      source.validate(new PipeParameterValidator(sourceParameters));

      // 3. Customize extractor
      final PipeTaskRuntimeConfiguration runtimeConfiguration =
          new PipeTaskRuntimeConfiguration(
              new PipeTaskSourceRuntimeEnvironment(
                  pipeName, creationTime, CONFIG_REGION_ID.getId(), pipeTaskMeta));
      source.customize(sourceParameters, runtimeConfiguration);
    } catch (final Exception e) {
      try {
        source.close();
      } catch (Exception closeException) {
        LOGGER.warn(
            "Failed to close extractor after failed to initialize extractor. "
                + "Ignore this exception.",
            closeException);
      }
      throw e;
    }
  }

  private void initProcessor(final Map<String, String> processorAttributes) {
    final PipeParameters processorParameters = new PipeParameters(processorAttributes);

    final PipeTaskRuntimeConfiguration runtimeConfiguration =
        new PipeTaskRuntimeConfiguration(
            new PipeTaskProcessorRuntimeEnvironment(
                pipeName, creationTime, CONFIG_REGION_ID.getId(), null));

    processor =
        PipeConfigNodeAgent.plugin()
            .getConfiguredProcessor(
                processorParameters.getStringOrDefault(
                    PipeProcessorConstant.PROCESSOR_KEY,
                    BuiltinPipePlugin.DO_NOTHING_PROCESSOR.getPipePluginName()),
                processorParameters,
                runtimeConfiguration);
  }

  private void initSink(final Map<String, String> sinkAttributes) throws Exception {
    final PipeParameters sinkParameters = new PipeParameters(sinkAttributes);

    // 1. Construct connector
    outputPipeSink = PipeConfigNodeAgent.plugin().reflectSink(sinkParameters);

    try {
      // 2. Validate connector parameters
      outputPipeSink.validate(new PipeParameterValidator(sinkParameters));

      // 3. Customize connector
      final PipeTaskRuntimeConfiguration runtimeConfiguration =
          new PipeTaskRuntimeConfiguration(
              new PipeTaskSinkRuntimeEnvironment(pipeName, creationTime, CONFIG_REGION_ID.getId()));
      outputPipeSink.customize(sinkParameters, runtimeConfiguration);

      // 4. Handshake
      outputPipeSink.handshake();
    } catch (final Exception e) {
      try {
        outputPipeSink.close();
      } catch (final Exception closeException) {
        LOGGER.warn(
            "Failed to close sink after failed to initialize it. Ignore this exception.",
            closeException);
      }
      throw e;
    }
  }

  /**
   * Try to consume an {@link Event} by the {@link IoTDBConfigRegionSource}.
   *
   * @return {@code true} if the {@link Event} is consumed successfully, {@code false} if no more
   *     {@link Event} can be consumed
   * @throws Exception if any error occurs when consuming the {@link Event}
   */
  @Override
  protected boolean executeOnce() throws Exception {
    if (isClosed.get()) {
      return false;
    }

    final Event event = lastEvent != null ? lastEvent : source.supply();
    // Record the last event for retry when exception occurs
    setLastEvent(event);

    try {
      if (event == null) {
        return false;
      }

      if (!(event instanceof ProgressReportEvent)) {
        outputPipeSink.transfer(event);
        PipeConfigRegionSinkMetrics.getInstance().markConfigEvent(taskID);
      }
      decreaseReferenceCountAndReleaseLastEvent(event, true);

    } catch (final PipeException e) {
      setLastExceptionEvent(event);
      if (!isClosed.get()) {
        throw e;
      } else {
        LOGGER.info(
            "{} in pipe transfer, ignored because pipe is dropped.",
            e.getClass().getSimpleName(),
            e);
        clearReferenceCountAndReleaseLastEvent(event);
      }
    } catch (final Exception e) {
      setLastExceptionEvent(event);
      if (!isClosed.get()) {
        throw new PipeException(
            String.format(
                "Exception in pipe transfer, subtask: %s, last event: %s", taskID, lastEvent),
            e);
      } else {
        LOGGER.info("Exception in pipe transfer, ignored because pipe is dropped.", e);
        clearReferenceCountAndReleaseLastEvent(event);
      }
    }

    return true;
  }

  @Override
  public void close() {
    isClosed.set(true);

    PipeEventCommitManager.getInstance()
        .deregister(pipeName, creationTime, CONFIG_REGION_ID.getId());
    PipeConfigRegionSinkMetrics.getInstance().deregister(taskID);

    try {
      source.close();
    } catch (final Exception e) {
      LOGGER.info("Error occurred during closing PipeExtractor.", e);
    }

    try {
      processor.close();
    } catch (final Exception e) {
      LOGGER.info("Error occurred during closing PipeProcessor.", e);
    }

    try {
      outputPipeSink.close();
    } catch (final Exception e) {
      LOGGER.info("Error occurred during closing PipeConnector.", e);
    } finally {
      // Should be after connector.close()
      super.close();
    }
  }

  //////////////////////////// Error report ////////////////////////////

  @Override
  protected String getRootCause(final Throwable throwable) {
    return throwable != null ? throwable.getMessage() : null;
  }

  @Override
  protected void report(final EnrichedEvent event, final PipeRuntimeException exception) {
    PipeConfigNodeAgent.runtime().report(event, exception);
  }

  //////////////////////////// APIs provided for metric framework ////////////////////////////

  public String getPipeName() {
    return pipeName;
  }
}
