/*
 * Copyright 2021 the original author or authors.
 *
 * Licensed 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.gradle.internal.configuration.inputs;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ForwardingMap;
import org.jspecify.annotations.Nullable;

import java.io.Serializable;
import java.util.Collection;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.function.BiConsumer;

/**
 * A wrapper for the {@link System#getenv()} result that notifies a listener about accesses.
 */
public class AccessTrackingEnvMap extends ForwardingMap<String, String> implements Serializable {
    private final Map<String, String> delegate;
    private final BiConsumer<? super String, ? super String> onAccess;

    public AccessTrackingEnvMap(BiConsumer<String, String> onAccess) {
        this(System.getenv(), onAccess);
    }

    @VisibleForTesting
    public AccessTrackingEnvMap(Map<String, String> delegate, BiConsumer<? super String, ? super String> onAccess) {
        this.delegate = delegate;
        this.onAccess = onAccess;
    }

    @Override
    public @Nullable String get(@Nullable Object key) {
        return getAndReport(key);
    }

    @Override
    @SuppressWarnings("SuspiciousMethodCalls")
    @Nullable
    public String getOrDefault(@Nullable Object key, @Nullable String defaultValue) {
        String value = getAndReport(key);
        if (value == null && !delegate.containsKey(key)) {
            return defaultValue;
        }
        return value;
    }

    @Override
    public boolean containsKey(@Nullable Object key) {
        return getAndReport(key) != null;
    }

    @Override
    public Set<String> keySet() {
        return new AccessTrackingSet<>(super.keySet(), keyTrackingListener());
    }

    @Nullable
    @SuppressWarnings("SuspiciousMethodCalls")
    private String getAndReport(@Nullable Object key) {
        String result = delegate.get(key);
        // The delegate will throw if something that isn't a string is used there. Do call delegate.get() first so the exception is thrown form the JDK code to avoid extra blame.
        onAccess.accept((String) key, result);
        return result;
    }

    @Override
    public Collection<String> values() {
        reportAggregatingAccess();
        return delegate.values();
    }

    @Override
    public Set<Entry<String, String>> entrySet() {
        return new AccessTrackingSet<>(delegate.entrySet(), entrySetTrackingListener());
    }

    private void onAccessEntrySetElement(@Nullable Object potentialEntry) {
        Map.Entry<String, String> entry = AccessTrackingUtils.tryConvertingToTrackableEntry(potentialEntry);
        if (entry != null) {
            getAndReport(entry.getKey());
        }
    }

    @Override
    public void forEach(BiConsumer<? super String, ? super String> action) {
        reportAggregatingAccess();
        delegate.forEach(action);
    }

    @Override
    public int size() {
        reportAggregatingAccess();
        return delegate.size();
    }

    @Override
    public boolean isEmpty() {
        reportAggregatingAccess();
        return delegate.isEmpty();
    }

    @Override
    @SuppressWarnings("EqualsWhichDoesntCheckParameterClass")
    public boolean equals(@Nullable Object object) {
        reportAggregatingAccess();
        return delegate.equals(object);
    }

    @Override
    public int hashCode() {
        reportAggregatingAccess();
        return delegate.hashCode();
    }

    @Override
    protected Map<String, String> delegate() {
        return delegate;
    }

    private void reportAggregatingAccess() {
        // Mark all map contents as inputs if some aggregating access is used.
        delegate.forEach(onAccess);
    }

    private AccessTrackingSet.Listener keyTrackingListener() {
        return new AccessTrackingSet.Listener() {
            @Override
            public void onAccess(@Nullable Object o) {
                getAndReport(Objects.requireNonNull(o));
            }

            @Override
            public void onAggregatingAccess() {
                reportAggregatingAccess();
            }

            @Override
            public void onRemove(@Nullable Object object) {
                // Environment variables are immutable.
            }

            @Override
            public void onClear() {
                // Environment variables are immutable.
            }
        };
    }

    private AccessTrackingSet.Listener entrySetTrackingListener() {
        return new AccessTrackingSet.Listener() {
            @Override
            public void onAccess(@Nullable Object o) {
                onAccessEntrySetElement(o);
            }

            @Override
            public void onAggregatingAccess() {
                reportAggregatingAccess();
            }

            @Override
            public void onRemove(@Nullable Object object) {
                // Environment variables are immutable.
            }

            @Override
            public void onClear() {
                // Environment variables are immutable.
            }
        };
    }

    private Object writeReplace() {
        // When we serialize the whole map, it is likely used as a task input.
        // We don't know how it is going to be used afterward, and we cannot affect the fingerprint during execution.
        // Let's be pessimistic and consider everything an input. On a bright side, we don't need access tracking for the deserialized thing.
        reportAggregatingAccess();
        return delegate;
    }
}
