/*
 * 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.cassandra.index.sai;

import java.io.IOException;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.NavigableSet;
import java.util.Set;
import java.util.TreeSet;

import io.github.jbellis.jvector.util.Bits;
import org.apache.cassandra.db.ReadCommand;
import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
import org.apache.cassandra.index.sai.disk.v1.vector.DiskAnn;
import org.apache.cassandra.index.sai.disk.v1.vector.OnHeapGraph;
import org.apache.cassandra.index.sai.utils.PrimaryKey;


/**
 * This represents the state of a vector query. It is repsonsible for maintaining a list of any {@link PrimaryKey}s
 * that have been updated or deleted during a search of the indexes.
 * <p>
 * The number of {@link #shadowedPrimaryKeys} is compared before and after a search is performed. If it changes, it
 * means that a {@link PrimaryKey} was found to have been changed. In this case the whole search is repeated until the
 * counts match.
 * <p>
 * When this process has completed, a {@link Bits} array is generated. This is used by the vector graph search to
 * identify which nodes in the graph to include in the results.
 */
public class VectorQueryContext
{
    private final int limit;
    // Holds primary keys that are shadowed by expired TTL or row tombstone or range tombstone.
    // They are populated by the StorageAttachedIndexSearcher during filtering. They are used to generate
    // a bitset for the graph search to indicate graph nodes to ignore.
    private TreeSet<PrimaryKey> shadowedPrimaryKeys;

    public VectorQueryContext(ReadCommand readCommand)
    {
        this.limit = readCommand.limits().count();
    }

    public int limit()
    {
        return limit;
    }

    public void recordShadowedPrimaryKeys(Set<PrimaryKey> keys)
    {
        if (shadowedPrimaryKeys == null)
            shadowedPrimaryKeys = new TreeSet<>();
        shadowedPrimaryKeys.addAll(keys);
    }

    // Returns true if the row ID will be included or false if the row ID will be shadowed
    public boolean shouldInclude(long sstableRowId, PrimaryKeyMap primaryKeyMap)
    {
        return shadowedPrimaryKeys == null || !shadowedPrimaryKeys.contains(primaryKeyMap.primaryKeyFromRowId(sstableRowId));
    }

    public boolean shouldInclude(PrimaryKey pk)
    {
        return shadowedPrimaryKeys == null || !shadowedPrimaryKeys.contains(pk);
    }

    public boolean containsShadowedPrimaryKey(PrimaryKey primaryKey)
    {
        return shadowedPrimaryKeys != null && shadowedPrimaryKeys.contains(primaryKey);
    }

    /**
     * @return shadowed primary keys, in ascending order
     */
    public NavigableSet<PrimaryKey> getShadowedPrimaryKeys()
    {
        if (shadowedPrimaryKeys == null)
            return Collections.emptyNavigableSet();
        return shadowedPrimaryKeys;
    }

    public Bits bitsetForShadowedPrimaryKeys(OnHeapGraph<PrimaryKey> graph)
    {
        if (shadowedPrimaryKeys == null)
            return null;

        return new IgnoredKeysBits(graph, shadowedPrimaryKeys);
    }

    public Bits bitsetForShadowedPrimaryKeys(SegmentMetadata metadata, PrimaryKeyMap primaryKeyMap, DiskAnn graph) throws IOException
    {
        Set<Integer> ignoredOrdinals = null;
        try (var ordinalsView = graph.getOrdinalsView())
        {
            for (PrimaryKey primaryKey : getShadowedPrimaryKeys())
            {
                // not in current segment
                if (primaryKey.compareTo(metadata.minKey) < 0 || primaryKey.compareTo(metadata.maxKey) > 0)
                    continue;

                long sstableRowId = primaryKeyMap.rowIdFromPrimaryKey(primaryKey);
                if (sstableRowId == Long.MAX_VALUE) // not found
                    continue;

                int segmentRowId = Math.toIntExact(sstableRowId - metadata.rowIdOffset);
                // not in segment yet
                if (segmentRowId < 0)
                    continue;
                // end of segment
                if (segmentRowId > metadata.maxSSTableRowId)
                    break;

                int ordinal = ordinalsView.getOrdinalForRowId(segmentRowId);
                if (ordinal >= 0)
                {
                    if (ignoredOrdinals == null)
                        ignoredOrdinals = new HashSet<>();
                    ignoredOrdinals.add(ordinal);
                }
            }
        }

        if (ignoredOrdinals == null)
            return null;

        return new IgnoringBits(ignoredOrdinals, metadata);
    }

    private static class IgnoringBits implements Bits
    {
        private final Set<Integer> ignoredOrdinals;
        private final int length;

        public IgnoringBits(Set<Integer> ignoredOrdinals, SegmentMetadata metadata)
        {
            this.ignoredOrdinals = ignoredOrdinals;
            this.length = 1 + Math.toIntExact(metadata.maxSSTableRowId - metadata.rowIdOffset);
        }

        @Override
        public boolean get(int index)
        {
            return !ignoredOrdinals.contains(index);
        }

        @Override
        public int length()
        {
            return length;
        }
    }

    private static class IgnoredKeysBits implements Bits
    {
        private final OnHeapGraph<PrimaryKey> graph;
        private final NavigableSet<PrimaryKey> ignored;

        public IgnoredKeysBits(OnHeapGraph<PrimaryKey> graph, NavigableSet<PrimaryKey> ignored)
        {
            this.graph = graph;
            this.ignored = ignored;
        }

        @Override
        public boolean get(int ordinal)
        {
            Collection<PrimaryKey> keys = graph.keysFromOrdinal(ordinal);
            return keys.stream().anyMatch(k -> !ignored.contains(k));
        }

        @Override
        public int length()
        {
            return graph.size();
        }
    }
}
