//
//  VectorIndexTest.h
//  CouchbaseLite
//
//  Copyright (c) 2024 Couchbase, Inc All rights reserved.
//
//  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.
//

#import "CBLTestCase.h"

#ifdef COUCHBASE_ENTERPRISE

NS_ASSUME_NONNULL_BEGIN

#define kWordsDatabaseName @"words_db"

#define kWordsCollectionName @"words"

#define kExtWordsCollectionName @"extwords"

#define kWordsIndexName @"words_index"

#define kWordPredictiveModelName @"WordEmbedding"

#define VECTOR_INDEX_CONFIG(E, D, C) [[CBLVectorIndexConfiguration alloc] initWithExpression: (E) dimensions: (D) centroids: (C)]

@interface VectorSearchTest : CBLTestCase

@property (nonatomic, readonly) CBLDatabase* wordDB;

@property (nonatomic, readonly) CBLCollection* wordsCollection;

@property (nonatomic, readonly) CBLCollection* extWordsCollection;

- (void) createVectorIndexInCollection: (CBLCollection*)collection
                                  name: (NSString*)name
                                config: (CBLVectorIndexConfiguration*)config;

- (void) createWordsIndexWithConfig: (CBLVectorIndexConfiguration*)config;

- (void) deleteWordsIndex;

/** For the test subclasses to override the default vector expression. */
- (NSString*) wordsQueryDefaultExpression;

- (NSString*) wordsQueryStringWithLimit: (NSUInteger)limit
                                 metric: (nullable NSString*)metric
                       vectorExpression: (nullable NSString*)vectorExpression
                            whereClause: (nullable NSString*)whereClause;

- (NSString*) wordsQueryStringWithLimit: (NSUInteger)limit;

- (CBLQueryResultSet*) executeWordsQueryWithLimit: (NSUInteger)limit
                                           metric: (nullable NSString*)metric
                                 vectorExpression: (nullable NSString*)vectorExpression
                                      whereClause: (nullable NSString*)whereClause
                                    checkTraining: (BOOL) checkTraining;

- (CBLQueryResultSet*) executeWordsQueryWithLimit: (NSUInteger)limit;

- (CBLQueryResultSet*) executeWordsQueryNoTrainingCheckWithLimit: (NSUInteger)limit;

- (NSDictionary<NSString*, NSString*>*) toDocIDWordMap: (CBLQueryResultSet*)resultSet;

@end

NS_ASSUME_NONNULL_END

#endif
