/*
 * Copyright (c) 2019 the Eclipse Milo Authors
 *
 * This program and the accompanying materials are made
 * available under the terms of the Eclipse Public License 2.0
 * which is available at https://www.eclipse.org/legal/epl-2.0/
 *
 * SPDX-License-Identifier: EPL-2.0
 */

package org.eclipse.milo.opcua.sdk.client.subscriptions;

import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;
import java.util.stream.Collectors;

import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.ListMultimap;
import com.google.common.collect.Lists;
import com.google.common.collect.Multimaps;
import org.eclipse.milo.opcua.sdk.client.OpcUaClient;
import org.eclipse.milo.opcua.sdk.client.nodes.UaVariableNode;
import org.eclipse.milo.opcua.sdk.core.util.GroupMapCollate;
import org.eclipse.milo.opcua.stack.core.AttributeId;
import org.eclipse.milo.opcua.stack.core.Identifiers;
import org.eclipse.milo.opcua.stack.core.StatusCodes;
import org.eclipse.milo.opcua.stack.core.UaException;
import org.eclipse.milo.opcua.stack.core.types.builtin.ExtensionObject;
import org.eclipse.milo.opcua.stack.core.types.builtin.StatusCode;
import org.eclipse.milo.opcua.stack.core.types.builtin.unsigned.UInteger;
import org.eclipse.milo.opcua.stack.core.types.enumerated.TimestampsToReturn;
import org.eclipse.milo.opcua.stack.core.types.structured.MonitoredItemModifyRequest;
import org.eclipse.milo.opcua.stack.core.types.structured.MonitoringParameters;
import org.eclipse.milo.opcua.stack.core.util.FutureUtils;
import org.jetbrains.annotations.Nullable;

import static org.eclipse.milo.opcua.stack.core.types.builtin.unsigned.Unsigned.uint;


/**
 * Manages the batch execution of one or more modifications to one or more MonitoredItems.
 * <p>
 * Handles splitting the execution into multiple service calls as needed based on the requested modifications and the
 * operation limits reported by the server.
 */
public class BatchModifyMonitoredItems {

    /**
     * MonitoredItems used as a parameter in add() mapped to their builder containing the modifications to make.
     */
    private final Map<OpcUaMonitoredItem, BatchModifyParametersBuilder> buildersByItem =
        Collections.synchronizedMap(new LinkedHashMap<>());

    /**
     * Holds possibly multiple CompletableFutures per-item. A MonitoredItem may have multiple futures if it was used as
     * the parameter in add() more than once, which is expected if multiple monitoring parameters are being modified.
     * <p>
     * There will only be as many underlying operation results as there are MonitoredItems, which is potentially less
     * than the total number futures returned by add(), so this ListMultimap makes it possible to multiplex those
     * underlying results onto all of the returned futures.
     */
    private final ListMultimap<OpcUaMonitoredItem, CompletableFuture<ModifyMonitoredItemResult>> futuresByItem =
        Multimaps.synchronizedListMultimap(ArrayListMultimap.create());

    /**
     * Every CompletableFuture returned by a call to add().
     * <p>
     * Used to ensure execute() can return CompletableFutures in the same order they were created by add() calls.
     */
    private final List<CompletableFuture<ModifyMonitoredItemResult>> resultFutures =
        Collections.synchronizedList(new ArrayList<>());

    private final AtomicInteger serviceInvocationCount = new AtomicInteger(0);

    private final OpcUaClient client;
    private final OpcUaSubscription subscription;

    public BatchModifyMonitoredItems(ManagedSubscription subscription) {
        this(subscription.getClient(), subscription.getSubscription());
    }

    public BatchModifyMonitoredItems(OpcUaClient client, OpcUaSubscription subscription) {
        this.client = client;
        this.subscription = subscription;
    }

    /**
     * Get the number of service invocations that were needed to execute this batch.
     *
     * @return the number of service invocations that were needed to execute this batch.
     */
    public int getServiceInvocationCount() {
        return serviceInvocationCount.get();
    }

    /**
     * Add an {@link OpcUaMonitoredItem} to be modified as part of a batch.
     *
     * @param monitoredItem   the {@link OpcUaMonitoredItem} to modify.
     * @param builderConsumer a consumer that will set changes on a {@link BatchModifyParametersBuilder}.
     * @return a {@link CompletableFuture} that always completes successfully with a {@link ModifyMonitoredItemResult}.
     */
    public CompletableFuture<ModifyMonitoredItemResult> add(
        OpcUaMonitoredItem monitoredItem,
        Consumer<BatchModifyParametersBuilder> builderConsumer
    ) {

        BatchModifyParametersBuilder builder = buildersByItem.computeIfAbsent(
            monitoredItem,
            BatchModifyMonitoredItems::parametersBuilderFromItem
        );

        builderConsumer.accept(builder);

        CompletableFuture<ModifyMonitoredItemResult> future = new CompletableFuture<>();
        futuresByItem.put(monitoredItem, future);
        resultFutures.add(future);
        return future;
    }

    /**
     * Execute this batch operation and return a List of {@link ModifyMonitoredItemResult}s the same size and order as
     * calls to {@link #add(OpcUaMonitoredItem, Consumer)} were made.
     *
     * @return a List of {@link ModifyMonitoredItemResult}s the same size and order as calls to
     * {@link #add(OpcUaMonitoredItem, Consumer)} were made.
     */
    public List<ModifyMonitoredItemResult> execute() throws InterruptedException {
        try {
            return executeAsync().get();
        } catch (ExecutionException e) {
            throw new IllegalStateException(e);
        }
    }

    /**
     * Execute this batch operation and return a {@link CompletableFuture} with a List of
     * {@link ModifyMonitoredItemResult}s the same size and order as calls to
     * {@link #add(OpcUaMonitoredItem, Consumer)} were made.
     * <p>
     * The returned CompletableFuture always completes successfully.
     *
     * @return a {@link CompletableFuture} with a List of {@link ModifyMonitoredItemResult}s the same size and order as
     * calls to {@link #add(OpcUaMonitoredItem, Consumer)} were made.
     */
    public CompletableFuture<List<ModifyMonitoredItemResult>> executeAsync() {
        return readOperationLimit(client).thenCompose(this::executeAsync);
    }

    private CompletableFuture<List<ModifyMonitoredItemResult>> executeAsync(UInteger operationLimit) {
        List<BatchModifyParameters> allMonitoringParameters = buildersByItem.values().stream()
            .map(BatchModifyParametersBuilder::build)
            .collect(Collectors.toList());

        CompletableFuture<List<ModifyMonitoredItemResult>> resultsFuture = GroupMapCollate.groupMapCollate(
            allMonitoringParameters,
            parametersItem -> parametersItem.timestamps,
            (TimestampsToReturn timestampsKey) -> parameterGroup -> {
                List<MonitoredItemModifyRequest> itemsToModify = parameterGroup.stream()
                    .map(parameters -> new MonitoredItemModifyRequest(
                        parameters.item.getMonitoredItemId(),
                        new MonitoringParameters(
                            parameters.clientHandle,
                            parameters.samplingInterval,
                            parameters.filter,
                            parameters.queueSize,
                            parameters.discardOldest
                        )
                    ))
                    .collect(Collectors.toList());

                List<CompletableFuture<List<ModifyMonitoredItemResult>>> partitionFutures =
                    Lists.partition(itemsToModify, operationLimit.intValue())
                        .stream()
                        .map(partition -> modifyItemsAsync(timestampsKey, partition))
                        .collect(Collectors.toList());

                return FutureUtils.flatSequence(partitionFutures);
            }
        );

        return resultsFuture.thenCompose(results -> {
            List<OpcUaMonitoredItem> items = new ArrayList<>(buildersByItem.keySet());

            assert items.size() == results.size();

            for (int i = 0; i < items.size(); i++) {
                OpcUaMonitoredItem item = items.get(i);
                ModifyMonitoredItemResult result = results.get(i);

                List<CompletableFuture<ModifyMonitoredItemResult>> futures;
                synchronized (futuresByItem) {
                    futures = new ArrayList<>(futuresByItem.get(item));
                }
                futures.forEach(f -> f.complete(result));
            }

            return FutureUtils.sequence(resultFutures);
        });
    }

    /**
     * Make the ModifyMonitoredItems service call for a partition of MonitoredItemModifyRequest.
     * <p>
     * The returned CompletableFuture always completes successfully.
     * <p>
     * The corresponding operation result for each item to modify may be either the StatusCode from the service call
     * if the call failed or the StatusCode from the operation if the service call succeeded.
     *
     * @param timestamps    the {@link TimestampsToReturn} to use in the  modify call.
     * @param itemsToModify the {@link MonitoredItemModifyRequest} for the items to modify.
     * @return a {@link CompletableFuture} that is always completed successfully.
     */
    private CompletableFuture<List<ModifyMonitoredItemResult>> modifyItemsAsync(
        TimestampsToReturn timestamps,
        List<MonitoredItemModifyRequest> itemsToModify
    ) {

        serviceInvocationCount.incrementAndGet();

        CompletableFuture<List<StatusCode>> modifyResults =
            subscription.modifyMonitoredItems(timestamps, itemsToModify);

        CompletableFuture<List<ModifyMonitoredItemResult>> batchItemResults =
            modifyResults.thenApply(statusCodes ->
                statusCodes.stream()
                    .map(statusCode -> new ModifyMonitoredItemResult(StatusCode.GOOD, statusCode))
                    .collect(Collectors.toList())
            );

        return batchItemResults.exceptionally(ex -> {
            StatusCode serviceResult = UaException
                .extractStatusCode(ex)
                .orElse(new StatusCode(StatusCodes.Bad_UnexpectedError));

            ModifyMonitoredItemResult result = new ModifyMonitoredItemResult(serviceResult);

            return Collections.nCopies(itemsToModify.size(), result);
        });
    }

    private static BatchModifyParametersBuilder parametersBuilderFromItem(OpcUaMonitoredItem item) {
        return new BatchModifyParametersBuilder(item)
            .setTimestamps(item.getTimestamps())
            .setClientHandle(item.getClientHandle())
            .setSamplingInterval(item.getRevisedSamplingInterval())
            .setFilter(item.getMonitoringFilter())
            .setQueueSize(item.getRevisedQueueSize())
            .setDiscardOldest(item.getDiscardOldest());
    }

    private static CompletableFuture<UInteger> readOperationLimit(OpcUaClient client) {
        CompletableFuture<UaVariableNode> nodeFuture = client.getAddressSpace().getVariableNodeAsync(
            Identifiers.Server_ServerCapabilities_OperationLimits_MaxMonitoredItemsPerCall
        );

        return nodeFuture.thenCompose(
            variableNode ->
                variableNode.readAttributeAsync(AttributeId.Value)
                    .thenApply(v -> (UInteger) v.getValue().getValue())
                    .exceptionally(ex -> uint(1000))
        );
    }

    /**
     * Holds the operation- and service-level results of an operation that is part of a
     * {@link BatchModifyMonitoredItems} execution.
     */
    public static class ModifyMonitoredItemResult implements BatchItemResult<StatusCode> {

        private final StatusCode serviceResult;
        private final StatusCode operationResult;

        public ModifyMonitoredItemResult(StatusCode serviceResult) {
            this(serviceResult, null);
        }

        ModifyMonitoredItemResult(StatusCode serviceResult, @Nullable StatusCode operationResult) {
            this.serviceResult = serviceResult;
            this.operationResult = operationResult;
        }

        @Override
        public StatusCode serviceResult() {
            return serviceResult;
        }

        @Override
        public Optional<StatusCode> operationResult() {
            return Optional.ofNullable(operationResult);
        }

        public boolean isServiceResultGood() {
            return serviceResult.isGood();
        }

        public boolean isOperationResultGood() {
            return operationResult().map(StatusCode::isGood).orElse(false);
        }

    }

    private static class BatchModifyParameters {

        final OpcUaMonitoredItem item;
        final TimestampsToReturn timestamps;
        final UInteger clientHandle;
        final Double samplingInterval;
        final ExtensionObject filter;
        final UInteger queueSize;
        final Boolean discardOldest;

        BatchModifyParameters(
            OpcUaMonitoredItem item, TimestampsToReturn timestamps,
            UInteger clientHandle,
            Double samplingInterval,
            ExtensionObject filter,
            UInteger queueSize,
            Boolean discardOldest
        ) {

            this.item = item;
            this.timestamps = timestamps;
            this.clientHandle = clientHandle;
            this.samplingInterval = samplingInterval;
            this.filter = filter;
            this.queueSize = queueSize;
            this.discardOldest = discardOldest;
        }
    }

    public static class BatchModifyParametersBuilder {

        private TimestampsToReturn timestamps;
        private UInteger clientHandle;
        private Double samplingInterval;
        private ExtensionObject filter;
        private UInteger queueSize;
        private Boolean discardOldest;

        private final OpcUaMonitoredItem item;

        public BatchModifyParametersBuilder(OpcUaMonitoredItem item) {
            this.item = item;
        }

        public BatchModifyParametersBuilder setTimestamps(TimestampsToReturn timestamps) {
            this.timestamps = timestamps;
            return this;
        }

        public BatchModifyParametersBuilder setClientHandle(UInteger clientHandle) {
            this.clientHandle = clientHandle;
            return this;
        }

        public BatchModifyParametersBuilder setSamplingInterval(double samplingInterval) {
            this.samplingInterval = samplingInterval;
            return this;
        }

        public BatchModifyParametersBuilder setFilter(ExtensionObject filter) {
            this.filter = filter;
            return this;
        }

        public BatchModifyParametersBuilder setQueueSize(UInteger queueSize) {
            this.queueSize = queueSize;
            return this;
        }

        public BatchModifyParametersBuilder setDiscardOldest(boolean discardOldest) {
            this.discardOldest = discardOldest;
            return this;
        }

        private BatchModifyParameters build() {
            return new BatchModifyParameters(
                item,
                timestamps,
                clientHandle,
                samplingInterval,
                filter,
                queueSize,
                discardOldest
            );
        }

    }

}
