/*
 * 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.kafka.connect.integration;

import org.apache.kafka.connect.errors.DataException;
import org.apache.kafka.connect.sink.SinkRecord;

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

import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
import java.util.stream.IntStream;

/**
 * A handle to a connector executing in a Connect cluster.
 */
public class ConnectorHandle {

    private static final Logger log = LoggerFactory.getLogger(ConnectorHandle.class);

    private final String connectorName;
    private final Map<String, TaskHandle> taskHandles = new ConcurrentHashMap<>();
    private final StartAndStopCounter startAndStopCounter = new StartAndStopCounter();

    private CountDownLatch recordsRemainingLatch;
    private CountDownLatch recordsToCommitLatch;
    private int expectedRecords = -1;
    private int expectedCommits = -1;

    public ConnectorHandle(String connectorName) {
        this.connectorName = connectorName;
    }

    /**
     * Get or create a task handle for a given task id. The task need not be created when this method is called. If the
     * handle is called before the task is created, the task will bind to the handle once it starts (or restarts).
     *
     * @param taskId the task id
     * @return a non-null {@link TaskHandle}
     */
    public TaskHandle taskHandle(String taskId) {
        return taskHandle(taskId, null);
    }

    /**
     * Get or create a task handle for a given task id. The task need not be created when this method is called. If the
     * handle is called before the task is created, the task will bind to the handle once it starts (or restarts).
     *
     * @param taskId the task id
     * @param consumer A callback invoked when a sink task processes a record.
     * @return a non-null {@link TaskHandle}
     */
    public TaskHandle taskHandle(String taskId, Consumer<SinkRecord> consumer) {
        return taskHandles.computeIfAbsent(taskId, k -> new TaskHandle(this, taskId, consumer));
    }

    /**
     * Gets the start and stop counter corresponding to this handle.
     *
     * @return the start and stop counter
     */
    public StartAndStopCounter startAndStopCounter() {
        return startAndStopCounter;
    }

    /**
     * Get the connector's name corresponding to this handle.
     *
     * @return the connector's name
     */
    public String name() {
        return connectorName;
    }

    /**
     * Get the list of tasks handles monitored by this connector handle.
     *
     * @return the task handle list
     */
    public Collection<TaskHandle> tasks() {
        return taskHandles.values();
    }

    /**
     * Delete all task handles for this connector.
     */
    public void clearTasks() {
        log.info("Clearing {} existing task handles for connector {}", taskHandles.size(), connectorName);
        taskHandles.clear();
    }

    /**
     * Set the number of expected records for this connector.
     *
     * @param expected number of records
     */
    public void expectedRecords(int expected) {
        expectedRecords = expected;
        recordsRemainingLatch = new CountDownLatch(expected);
    }

    /**
     * Set the number of expected commits performed by this connector.
     *
     * @param expected number of commits
     */
    public void expectedCommits(int expected) {
        expectedCommits = expected;
        recordsToCommitLatch = new CountDownLatch(expected);
    }

    /**
     * Record a message arrival at the connector.
     */
    public void record() {
        if (recordsRemainingLatch != null) {
            recordsRemainingLatch.countDown();
        }
    }

    /**
     * Record arrival of a batch of messages at the connector.
     *
     * @param batchSize the number of messages
     */
    public void record(int batchSize) {
        if (recordsRemainingLatch != null) {
            IntStream.range(0, batchSize).forEach(i -> recordsRemainingLatch.countDown());
        }
    }

    /**
     * Record a message commit from the connector.
     */
    public void commit() {
        if (recordsToCommitLatch != null) {
            recordsToCommitLatch.countDown();
        }
    }

    /**
     * Record commit on a batch of messages from the connector.
     *
     * @param batchSize the number of messages
     */
    public void commit(int batchSize) {
        if (recordsToCommitLatch != null) {
            IntStream.range(0, batchSize).forEach(i -> recordsToCommitLatch.countDown());
        }
    }

    /**
     * Wait for this connector to meet the expected number of records as defined by {@code
     * expectedRecords}.
     *
     * @param timeout max duration to wait for records
     * @throws InterruptedException if another threads interrupts this one while waiting for records
     */
    public void awaitRecords(long timeout) throws InterruptedException {
        if (recordsRemainingLatch == null || expectedRecords < 0) {
            throw new IllegalStateException("expectedRecords() was not set for this connector?");
        }
        if (!recordsRemainingLatch.await(timeout, TimeUnit.MILLISECONDS)) {
            String msg = String.format(
                    "Insufficient records seen by connector %s in %d millis. Records expected=%d, actual=%d",
                    connectorName,
                    timeout,
                    expectedRecords,
                    expectedRecords - recordsRemainingLatch.getCount());
            throw new DataException(msg);
        }
    }

     /**
     * Wait for this connector to meet the expected number of commits as defined by {@code
     * expectedCommits}.
     *
     * @param  timeout duration to wait for commits
     * @throws InterruptedException if another threads interrupts this one while waiting for commits
     */
    public void awaitCommits(long timeout) throws InterruptedException {
        if (recordsToCommitLatch == null || expectedCommits < 0) {
            throw new IllegalStateException("expectedCommits() was not set for this connector?");
        }
        if (!recordsToCommitLatch.await(timeout, TimeUnit.MILLISECONDS)) {
            String msg = String.format(
                    "Insufficient records committed by connector %s in %d millis. Records expected=%d, actual=%d",
                    connectorName,
                    timeout,
                    expectedCommits,
                    expectedCommits - recordsToCommitLatch.getCount());
            throw new DataException(msg);
        }
    }

    /**
     * Record that this connector has been started. This should be called by the connector under
     * test.
     *
     * @see #expectedStarts(int)
     */
    public void recordConnectorStart() {
        startAndStopCounter.recordStart();
    }

    /**
     * Record that this connector has been stopped. This should be called by the connector under
     * test.
     *
     * @see #expectedStarts(int)
     */
    public void recordConnectorStop() {
        startAndStopCounter.recordStop();
    }

    /**
     * Obtain a {@link StartAndStopLatch} that can be used to wait until the connector using this handle
     * and all tasks using {@link TaskHandle} have completed the expected number of
     * starts, starting the counts at the time this method is called.
     *
     * <p>A test can call this method, specifying the number of times the connector and tasks
     * will each be stopped and started from that point (typically {@code expectedStarts(1)}).
     * The test should then change the connector or otherwise cause the connector to restart one or
     * more times, and then can call {@link StartAndStopLatch#await(long, TimeUnit)} to wait up to a
     * specified duration for the connector and all tasks to be started at least the specified
     * number of times.
     *
     * <p>This method does not track the number of times the connector and tasks are stopped, and
     * only tracks the number of times the connector and tasks are <em>started</em>.
     *
     * @param expectedStarts the minimum number of starts that are expected once this method is
     *                       called
     * @return the latch that can be used to wait for the starts to complete; never null
     */
    public StartAndStopLatch expectedStarts(int expectedStarts) {
        return expectedStarts(expectedStarts, true);
    }

    /**
     * Obtain a {@link StartAndStopLatch} that can be used to wait until the connector using this handle
     * and optionally all tasks using {@link TaskHandle} have completed the expected number of
     * starts, starting the counts at the time this method is called.
     *
     * <p>A test can call this method, specifying the number of times the connector and tasks
     * will each be stopped and started from that point (typically {@code expectedStarts(1)}).
     * The test should then change the connector or otherwise cause the connector to restart one or
     * more times, and then can call {@link StartAndStopLatch#await(long, TimeUnit)} to wait up to a
     * specified duration for the connector and all tasks to be started at least the specified
     * number of times.
     *
     * <p>This method does not track the number of times the connector and tasks are stopped, and
     * only tracks the number of times the connector and tasks are <em>started</em>.
     *
     * @param expectedStarts the minimum number of starts that are expected once this method is
     *                       called
     * @param includeTasks  true if the latch should also wait for the tasks to be stopped the
     *                      specified minimum number of times
     * @return the latch that can be used to wait for the starts to complete; never null
     */
    public StartAndStopLatch expectedStarts(int expectedStarts, boolean includeTasks) {
        List<StartAndStopLatch> taskLatches = includeTasks
                ? taskHandles.values().stream()
                .map(task -> task.expectedStarts(expectedStarts))
                .toList()
                : List.of();
        return startAndStopCounter.expectedStarts(expectedStarts, taskLatches);
    }

    public StartAndStopLatch expectedStarts(int expectedStarts, Map<String, Integer> expectedTasksStarts, boolean includeTasks) {
        List<StartAndStopLatch> taskLatches = includeTasks
                ? taskHandles.values().stream()
                .map(task -> task.expectedStarts(expectedTasksStarts.get(task.taskId())))
                .toList()
                : List.of();
        return startAndStopCounter.expectedStarts(expectedStarts, taskLatches);
    }

    /**
     * Obtain a {@link StartAndStopLatch} that can be used to wait until the connector using this handle
     * and optionally all tasks using {@link TaskHandle} have completed the minimum number of
     * stops, starting the counts at the time this method is called.
     *
     * <p>A test can call this method, specifying the number of times the connector and tasks
     * will each be stopped from that point (typically {@code expectedStops(1)}).
     * The test should then change the connector or otherwise cause the connector to stop (or
     * restart) one or more times, and then can call
     * {@link StartAndStopLatch#await(long, TimeUnit)} to wait up to a specified duration for the
     * connector and all tasks to be started at least the specified number of times.
     *
     * <p>This method does not track the number of times the connector and tasks are started, and
     * only tracks the number of times the connector and tasks are <em>stopped</em>.
     *
     * @param expectedStops the minimum number of starts that are expected once this method is
     *                      called
     * @return the latch that can be used to wait for the starts to complete; never null
     */
    public StartAndStopLatch expectedStops(int expectedStops) {
        return expectedStops(expectedStops, true);
    }

    /**
     * Obtain a {@link StartAndStopLatch} that can be used to wait until the connector using this handle
     * and optionally all tasks using {@link TaskHandle} have completed the minimum number of
     * stops, starting the counts at the time this method is called.
     *
     * <p>A test can call this method, specifying the number of times the connector and tasks
     * will each be stopped from that point (typically {@code expectedStops(1)}).
     * The test should then change the connector or otherwise cause the connector to stop (or
     * restart) one or more times, and then can call
     * {@link StartAndStopLatch#await(long, TimeUnit)} to wait up to a specified duration for the
     * connector and all tasks to be started at least the specified number of times.
     *
     * <p>This method does not track the number of times the connector and tasks are started, and
     * only tracks the number of times the connector and tasks are <em>stopped</em>.
     *
     * @param expectedStops the minimum number of starts that are expected once this method is
     *                      called
     * @param includeTasks  true if the latch should also wait for the tasks to be stopped the
     *                      specified minimum number of times
     * @return the latch that can be used to wait for the starts to complete; never null
     */
    public StartAndStopLatch expectedStops(int expectedStops, boolean includeTasks) {
        List<StartAndStopLatch> taskLatches = includeTasks
                ? taskHandles.values().stream()
                .map(task -> task.expectedStops(expectedStops))
                .toList()
                : List.of();
        return startAndStopCounter.expectedStops(expectedStops, taskLatches);
    }

    public StartAndStopLatch expectedStops(int expectedStops, Map<String, Integer> expectedTasksStops, boolean includeTasks) {
        List<StartAndStopLatch> taskLatches = includeTasks
                ? taskHandles.values().stream()
                .map(task -> task.expectedStops(expectedTasksStops.get(task.taskId())))
                .toList()
                : List.of();
        return startAndStopCounter.expectedStops(expectedStops, taskLatches);
    }

    @Override
    public String toString() {
        return "ConnectorHandle{" +
                "connectorName='" + connectorName + '\'' +
                '}';
    }
}
