/*
 * 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.nifi.components.connector;

import org.apache.nifi.asset.AssetManager;
import org.apache.nifi.components.connector.components.FlowContext;
import org.apache.nifi.components.connector.secrets.SecretsManager;
import org.apache.nifi.flow.Bundle;
import org.apache.nifi.flow.VersionedControllerService;
import org.apache.nifi.flow.VersionedExternalFlow;
import org.apache.nifi.flow.VersionedProcessGroup;
import org.apache.nifi.flow.VersionedProcessor;
import org.apache.nifi.logging.ComponentLog;

import java.util.List;
import java.util.function.Consumer;

public class StandardConnectorInitializationContext implements FrameworkConnectorInitializationContext {
    private final String identifier;
    private final String name;
    private final ComponentLog componentLog;
    private final SecretsManager secretsManager;
    private final AssetManager assetManager;
    private final ComponentBundleLookup componentBundleLookup;


    protected StandardConnectorInitializationContext(final Builder builder) {
        this.identifier = builder.identifier;
        this.name = builder.name;
        this.componentLog = builder.componentLog;
        this.secretsManager = builder.secretsManager;
        this.assetManager = builder.assetManager;
        this.componentBundleLookup = builder.componentBundleLookup;
    }

    @Override
    public String getIdentifier() {
        return identifier;
    }

    @Override
    public String getName() {
        return name;
    }

    @Override
    public ComponentLog getLogger() {
        return componentLog;
    }

    @Override
    public ComponentBundleLookup getComponentBundleLookup() {
        return componentBundleLookup;
    }

    @Override
    public SecretsManager getSecretsManager() {
        return secretsManager;
    }

    @Override
    public AssetManager getAssetManager() {
        return assetManager;
    }

    @Override
    public void updateFlow(final FlowContext flowContext, final VersionedExternalFlow versionedExternalFlow,
                           final BundleCompatibility bundleCompatability) throws FlowUpdateException {
        if (!(flowContext instanceof final FrameworkFlowContext frameworkFlowContext)) {
            throw new IllegalArgumentException("FlowContext is not an instance provided by the framework");
        }

        resolveBundles(versionedExternalFlow.getFlowContents(), bundleCompatability);
        frameworkFlowContext.updateFlow(versionedExternalFlow, assetManager);
    }

    protected void resolveBundles(final VersionedProcessGroup group, final BundleCompatibility bundleCompatability) {
        if (bundleCompatability == BundleCompatibility.REQUIRE_EXACT_BUNDLE) {
            return;
        }

        if (group.getProcessors() != null) {
            for (final VersionedProcessor processor : group.getProcessors()) {
                resolveBundle(processor.getType(), processor.getBundle(), bundleCompatability, processor::setBundle);
            }
        }

        if (group.getControllerServices() != null) {
            for (final VersionedControllerService service : group.getControllerServices()) {
                resolveBundle(service.getType(), service.getBundle(), bundleCompatability, service::setBundle);
            }
        }

        if (group.getProcessGroups() != null) {
            for (final VersionedProcessGroup childGroup : group.getProcessGroups()) {
                resolveBundles(childGroup, bundleCompatability);
            }
        }
    }

    private void resolveBundle(final String componentType, final Bundle currentBundle,
                               final BundleCompatibility bundleCompatability, final Consumer<Bundle> bundleSetter) {
        final List<Bundle> availableBundles = componentBundleLookup.getAvailableBundles(componentType);

        if (availableBundles.contains(currentBundle)) {
            return;
        }

        if (availableBundles.isEmpty()) {
            return;
        }

        switch (bundleCompatability) {
            case RESOLVE_BUNDLE:
                if (availableBundles.size() == 1) {
                    final Bundle resolvedBundle = availableBundles.getFirst();
                    componentLog.debug("Resolved bundle for {} from {} to {}", componentType, currentBundle, resolvedBundle);
                    bundleSetter.accept(resolvedBundle);
                }
                break;
            case RESOLVE_NEWEST_BUNDLE:
                componentBundleLookup.getLatestBundle(componentType).ifPresent(latestBundle -> {
                    componentLog.debug("Resolved bundle for {} from {} to {}", componentType, currentBundle, latestBundle);
                    bundleSetter.accept(latestBundle);
                });
                break;
            default:
                break;
        }
    }

    private void updateParameterContext(final VersionedProcessGroup group, final String parameterContextName) {
        group.setParameterContextName(parameterContextName);
        if (group.getProcessGroups() != null) {
            for (final VersionedProcessGroup childGroup : group.getProcessGroups()) {
                updateParameterContext(childGroup, parameterContextName);
            }
        }
    }

    public static class Builder implements FrameworkConnectorInitializationContextBuilder {
        private String identifier;
        private String name;
        private ComponentLog componentLog;
        private SecretsManager secretsManager;
        private AssetManager assetManager;
        private ComponentBundleLookup componentBundleLookup;

        @Override
        public Builder identifier(final String identifier) {
            this.identifier = identifier;
            return this;
        }

        @Override
        public Builder name(final String name) {
            this.name = name;
            return this;
        }

        @Override
        public Builder componentLog(final ComponentLog componentLog) {
            this.componentLog = componentLog;
            return this;
        }

        @Override
        public Builder secretsManager(final SecretsManager secretsManager) {
            this.secretsManager = secretsManager;
            return this;
        }

        @Override
        public Builder assetManager(final AssetManager assetManager) {
            this.assetManager = assetManager;
            return this;
        }

        @Override
        public Builder componentBundleLookup(final ComponentBundleLookup bundleLookup) {
            this.componentBundleLookup = bundleLookup;
            return this;
        }

        @Override
        public StandardConnectorInitializationContext build() {
            return new StandardConnectorInitializationContext(this);
        }
    }
}
