/*
 * Copyright 2026 Red Hat, Inc. and/or its affiliates
 * and other contributors as indicated by the @author tags.
 *
 * 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.keycloak.models.sessions.infinispan;

import java.util.Objects;
import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
import java.util.function.Supplier;

import org.infinispan.Cache;
import org.infinispan.commons.util.IntSet;
import org.infinispan.commons.util.IntSets;
import org.infinispan.distribution.ch.ConsistentHash;
import org.infinispan.distribution.ch.KeyPartitioner;
import org.infinispan.factories.ComponentRegistry;
import org.infinispan.notifications.Listener;
import org.infinispan.notifications.cachelistener.annotation.TopologyChanged;
import org.infinispan.notifications.cachelistener.event.TopologyChangedEvent;
import org.infinispan.remoting.transport.Address;

/**
 * A key affinity service that generates keys with affinity to the local node in a distributed cache.
 * <p>
 * This service provides similar functionality to {@link org.infinispan.affinity.KeyAffinityService}, but with a more
 * lenient, best-effort approach. Unlike {@code KeyAffinityService}, this implementation never blocks or throws
 * exceptions when unable to generate a key belonging to the local node. Instead, it attempts to generate an affine key
 * up to {@value #MAX_ATTEMPTS} times, and if unsuccessful, returns a random key.
 * <p>
 * The service automatically tracks topology changes and updates its knowledge of which cache segments are owned by the
 * local node. It prefers keys that map to primary segments owned by this node, but falls back to backup segments if no
 * primary segments are available.
 *
 * @param <K> the type of keys generated by this service
 */
@Listener
public class SessionAffinityService<K> implements Supplier<K> {

    public static final int MAX_ATTEMPTS = 32;
    @SuppressWarnings("rawtypes")
    private static final AtomicReferenceFieldUpdater<SessionAffinityService, Topology> UPDATER = AtomicReferenceFieldUpdater.newUpdater(SessionAffinityService.class, Topology.class, "topology");

    private final Supplier<K> generator;
    private final KeyPartitioner keyPartitioner;
    private final Address localAddress;
    private volatile Topology topology = new Topology(IntSets.immutableEmptySet(), -1);

    public SessionAffinityService(Supplier<K> generator, KeyPartitioner keyPartitioner, Address localAddress) {
        this.generator = Objects.requireNonNull(generator);
        this.keyPartitioner = Objects.requireNonNull(keyPartitioner);
        this.localAddress = Objects.requireNonNull(localAddress);
    }

    public static <T> Supplier<T> create(Cache<T, ?> cache, Supplier<T> generator) {
        if (generator == null || !cache.getCacheConfiguration().clustering().cacheMode().isClustered()) {
            return generator;
        }
        var affinityService = new SessionAffinityService<>(
                generator,
                ComponentRegistry.componentOf(cache, KeyPartitioner.class),
                cache.getCacheManager().getAddress());
        cache.addListener(affinityService);
        var cacheTopology = cache.getAdvancedCache().getDistributionManager().getCacheTopology();
        affinityService.computeTopology(cacheTopology.getWriteConsistentHash(), cacheTopology.getTopologyId());
        return affinityService;
    }

    @Override
    public K get() {
        var currentTopology = topology;
        if (currentTopology.segments.isEmpty()) {
            // this node does not have any segment assigned to him, use a random id.
            return generator.get();
        }
        for (int i = 0; i < MAX_ATTEMPTS; i++) {
            var key = generator.get();
            if (currentTopology.segments.contains(keyPartitioner.getSegment(key))) {
                return key;
            }
        }
        return generator.get();
    }

    @TopologyChanged
    public void handleViewChange(TopologyChangedEvent<K, ?> tce) {
        if (tce.isPre()) {
            return;
        }
        computeTopology(tce.getWriteConsistentHashAtEnd(), tce.getNewTopologyId());
    }

    private void computeTopology(ConsistentHash consistentHash, int topologyId) {
        var segments = consistentHash.getPrimarySegmentsForOwner(localAddress);
        if (segments.isEmpty()) {
            segments = consistentHash.getSegmentsForOwner(localAddress);
        }
        var newTopology = new Topology(IntSets.from(segments), topologyId);
        UPDATER.updateAndGet(this, newTopology::chooseLatest);
    }


    private record Topology(IntSet segments, int topologId) {

        private Topology {
            Objects.requireNonNull(segments);
        }

        public Topology chooseLatest(Topology other) {
            return topologId > other.topologId ? this : other;
        }
    }
}
