xinmi/internal/application/repository/retriever/weaviate/repository.go

package weaviate

import (
	"context"
	"fmt"
	"maps"
	"slices"
	"strings"
	"unicode/utf8"

	"github.com/Tencent/WeKnora/internal/logger"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/go-openapi/strfmt"
	"github.com/google/uuid"
	"github.com/weaviate/weaviate-go-client/v5/weaviate"
	"github.com/weaviate/weaviate-go-client/v5/weaviate/filters"
	"github.com/weaviate/weaviate-go-client/v5/weaviate/graphql"
	"github.com/weaviate/weaviate/entities/models"
)

const (
	envWeaviateCollection = "WEAVIATE_COLLECTION"
	defaultCollectionName = "Weknora_embeddings"
	fieldContent          = "content"
	fieldSourceID         = "source_id"
	fieldSourceType       = "source_type"
	fieldChunkID          = "chunk_id"
	fieldKnowledgeID      = "knowledge_id"
	fieldKnowledgeBaseID  = "knowledge_base_id"
	fieldTagID            = "tag_id"
	fieldEmbedding        = "embedding"
	fieldIsEnabled        = "is_enabled"
	fieldID               = "id"
)

// NewWeaviateRetrieveEngineRepository creates and initializes a new Weaviate repository.
// indexCfg is optional — pass nil to use env var / default values (env path).
func NewWeaviateRetrieveEngineRepository(client *weaviate.Client, indexCfg *types.IndexConfig) interfaces.RetrieveEngineRepository {
	log := logger.GetLogger(context.Background())
	log.Info("[Weaviate] Initializing Weaviate retriever engine repository")

	collectionBaseName := types.ResolveCollectionName(indexCfg, envWeaviateCollection, defaultCollectionName)

	res := &weaviateRepository{
		client:             client,
		collectionBaseName: collectionBaseName,
		replicationFactor:  indexCfg.GetReplicationFactor(0),
		desiredShardCount:  indexCfg.GetDesiredShardCount(0),
	}

	log.Info("[Weaviate] Successfully initialized repository")
	return res
}

func (w *weaviateRepository) getCollectionName(dimension int) string {
	return fmt.Sprintf("%s_%d", w.collectionBaseName, dimension)
}

func (w *weaviateRepository) ensureCollection(ctx context.Context, dimension int) error {
	collectionName := w.getCollectionName(dimension)

	//Check cache first
	if _, ok := w.initializedCollections.Load(dimension); ok {
		return nil
	}

	log := logger.GetLogger(ctx)

	//Check if collection exists
	exists, err := w.client.Schema().ClassExistenceChecker().WithClassName(collectionName).Do(ctx)
	if err != nil {
		log.Errorf("[Weaviate] Failed to check collection existence: %v", err)
		return fmt.Errorf("failed to check collection existence: %w", err)
	}
	enabled := true
	if !exists {
		log.Infof("[Weaviate] Creating collection %s with dimension %d", collectionName, dimension)

		//定义class结构
		classObj := models.Class{
			Class:       collectionName,
			Description: fmt.Sprintf("WeKnora embeddings collection with dimension %d", dimension),
			VectorConfig: map[string]models.VectorConfig{
				fieldEmbedding: {
					VectorIndexType: "hnsw",
					VectorIndexConfig: map[string]interface{}{
						"distance":       "cosine",
						"efConstruction": 128,
						"maxConnections": 32,
						"ef":             64,
					},
					Vectorizer: map[string]interface{}{
						"none": map[string]interface{}{},
					},
				},
			},
			Properties: []*models.Property{
				{
					Name:         fieldContent,
					DataType:     []string{"text"},
					Tokenization: "gse",
				},
				{
					Name:     fieldSourceID,
					DataType: []string{"text"},
				},
				{
					Name:     fieldSourceType,
					DataType: []string{"int"},
				},
				{
					Name:            fieldChunkID,
					DataType:        []string{"text"},
					IndexFilterable: &enabled,
				},
				{
					Name:            fieldKnowledgeID,
					DataType:        []string{"text"},
					IndexFilterable: &enabled,
				},
				{
					Name:            fieldKnowledgeBaseID,
					DataType:        []string{"text"},
					IndexFilterable: &enabled,
				},
				{
					Name:            fieldTagID,
					DataType:        []string{"text"},
					IndexFilterable: &enabled,
				},
				{
					Name:            fieldIsEnabled,
					DataType:        []string{"boolean"},
					IndexFilterable: &enabled,
				},
			},
		}
		// Set replication factor if explicitly configured (> 0)
		if w.replicationFactor > 0 {
			classObj.ReplicationConfig = &models.ReplicationConfig{
				Factor: int64(w.replicationFactor),
			}
		}
		// Set shard count if explicitly configured (> 0)
		if w.desiredShardCount > 0 {
			classObj.ShardingConfig = map[string]interface{}{
				"desiredCount": w.desiredShardCount,
			}
		}
		//创建collection
		if err = w.client.Schema().ClassCreator().WithClass(&classObj).Do(ctx); err != nil {
			log.Errorf("[Weaviate] Failed to create collection: %v", err)
			return fmt.Errorf("failed to create collection: %w", err)
		}
		log.Infof("[Weaviate] Successfully created collection %s", collectionName)
	}
	w.initializedCollections.Store(dimension, true)
	return nil
}

func (w *weaviateRepository) EngineType() types.RetrieverEngineType {
	return types.WeaviateRetrieverEngineType
}

func (w *weaviateRepository) Support() []types.RetrieverType {
	return []types.RetrieverType{types.KeywordsRetrieverType, types.VectorRetrieverType}
}

// EstimateStorageSize calculates the estimated storage size for a list of indices
func (w *weaviateRepository) EstimateStorageSize(ctx context.Context,
	indexInfoList []*types.IndexInfo, params map[string]any,
) int64 {
	var totalStorageSize int64
	for _, embedding := range indexInfoList {
		embeddingDB := toWeaviateVectorEmbedding(embedding, params)
		totalStorageSize += w.calculateStorageSize(embeddingDB)
	}
	logger.GetLogger(ctx).Infof(
		"[Weaviate] Storage size for %d indices: %d bytes", len(indexInfoList), totalStorageSize,
	)
	return totalStorageSize
}

// Save stores a single point in Weaviate
func (w *weaviateRepository) Save(ctx context.Context,
	embedding *types.IndexInfo, additionalParams map[string]any,
) error {
	log := logger.GetLogger(ctx)
	log.Debugf("[Weaviate] Saving index for chunk ID: %s", embedding.ChunkID)

	embeddingDB := toWeaviateVectorEmbedding(embedding, additionalParams)
	if len(embeddingDB.Embedding) == 0 {
		err := fmt.Errorf("empty embedding vector for chunk ID: %s", embedding.ChunkID)
		log.Errorf("[Weaviate] %v", err)
		return err
	}

	dimension := len(embeddingDB.Embedding)
	if err := w.ensureCollection(ctx, dimension); err != nil {
		return err
	}
	collectionName := w.getCollectionName(dimension)
	dataSchema := createPayload(embeddingDB)

	id := embedding.ChunkID
	// Create point in Weaviate
	_, err := w.client.Data().Creator().
		WithClassName(collectionName).
		WithID(id).
		WithProperties(dataSchema).
		WithVector(embeddingDB.Embedding).
		Do(ctx)

	if err != nil {
		log.Errorf("[Weaviate] Failed to save index: %v", err)
		return err
	}
	log.Infof("[Weaviate] Successfully saved index for chunk ID: %s", embedding.ChunkID)
	return nil
}

// BatchSave stores multiple points in Milvus using batch insert
func (w *weaviateRepository) BatchSave(ctx context.Context,
	embeddingList []*types.IndexInfo, additionalParams map[string]any,
) error {
	log := logger.GetLogger(ctx)
	if len(embeddingList) == 0 {
		log.Warn("[Weaviate] Empty list provided to BatchSave, skipping")
		return nil
	}

	log.Infof("[Weaviate] Batch saving %d indices", len(embeddingList))

	// Group points by dimension
	embeddingsByDimension := make(map[int][]*types.IndexInfo)

	for _, embedding := range embeddingList {
		embeddingDB := toWeaviateVectorEmbedding(embedding, additionalParams)
		if len(embeddingDB.Embedding) == 0 {
			log.Warnf("[Weaviate] Skipping empty embedding for chunk ID: %s", embedding.ChunkID)
			continue
		}

		dimension := len(embeddingDB.Embedding)
		embeddingsByDimension[dimension] = append(embeddingsByDimension[dimension], embedding)
		log.Debugf("[Weaviate] Added chunk ID %s to batch request (dimension: %d)", embedding.ChunkID, dimension)
	}

	if len(embeddingsByDimension) == 0 {
		log.Warn("[Weaviate] No valid points to save after filtering")
		return nil
	}

	// Save points to each dimension-specific collection
	totalSaved := 0
	for dimension, embeddings := range embeddingsByDimension {
		batcher := w.client.Batch().ObjectsBatcher()
		if err := w.ensureCollection(ctx, dimension); err != nil {
			return err
		}
		collectionName := w.getCollectionName(dimension)
		for _, embedding := range embeddings {
			embeddingDB := toWeaviateVectorEmbedding(embedding, additionalParams)
			dataSchema := createPayload(embeddingDB)

			obj := &models.Object{
				Class:      collectionName,
				ID:         strfmt.UUID(embeddingDB.ChunkID),
				Properties: dataSchema,
				Vector:     embeddingDB.Embedding,
			}
			batcher.WithObjects(obj)
		}
		// Flush batch
		if _, err := batcher.Do(ctx); err != nil {
			log.Errorf("[Weaviate] Failed to execute batch operation for dimension %d: %v", dimension, err)
			return fmt.Errorf("failed to batch save (dimension %d): %w", dimension, err)
		}
		totalSaved += len(embeddings)
		log.Infof("[Weaviate] Saved %d points to collection %s", len(embeddings), collectionName)
	}

	log.Infof("[Weaviate] Successfully batch saved %d indices", totalSaved)
	return nil
}

// DeleteByChunkIDList removes points from the collection based on chunk IDs
func (w *weaviateRepository) DeleteByChunkIDList(ctx context.Context, chunkIDList []string, dimension int, knowledgeType string) error {
	log := logger.GetLogger(ctx)
	if len(chunkIDList) == 0 {
		log.Warn("[Weaviate] Empty chunk ID list provided for deletion, skipping")
		return nil
	}

	collectionName := w.getCollectionName(dimension)
	log.Infof("[Weaviate] Deleting indices by chunk IDs from %s, count: %d", collectionName, len(chunkIDList))

	//define filter
	filter := w.client.Batch().ObjectsBatchDeleter().
		WithClassName(collectionName).
		WithWhere(filters.Where().
			WithPath([]string{fieldChunkID}).
			WithOperator(filters.ContainsAny).
			WithValueText(chunkIDList...)).
		WithOutput("minimal")

	// Execute deletion
	if _, err := filter.Do(ctx); err != nil {
		log.Errorf("[Weaviate] Failed to delete by chunk IDs: %v", err)
		return fmt.Errorf("failed to delete by chunk IDs: %w", err)
	}
	log.Infof("[Weaviate] Successfully deleted documents by chunk IDs")
	return nil
}

// DeleteByKnowledgeIDList removes points from the collection based on knowledge IDs
func (w *weaviateRepository) DeleteByKnowledgeIDList(ctx context.Context,
	knowledgeIDList []string, dimension int, knowledgeType string,
) error {
	log := logger.GetLogger(ctx)
	if len(knowledgeIDList) == 0 {
		log.Warn("[Weaviate] Empty knowledge ID list provided for deletion, skipping")
		return nil
	}

	collectionName := w.getCollectionName(dimension)
	log.Infof("[Weaviate] Deleting indices by knowledge IDs from %s, count: %d", collectionName, len(knowledgeIDList))

	//define filter
	filter := w.client.Batch().ObjectsBatchDeleter().
		WithClassName(collectionName).
		WithWhere(filters.Where().
			WithPath([]string{fieldKnowledgeID}).
			WithOperator(filters.ContainsAny).
			WithValueText(knowledgeIDList...)).
		WithOutput("minimal")

	// Execute deletion
	if _, err := filter.Do(ctx); err != nil {
		log.Errorf("[Weaviate] Failed to delete by knowledge IDs: %v", err)
		return fmt.Errorf("failed to delete by knowledge IDs: %w", err)
	}
	log.Infof("[Weaviate] Successfully deleted documents by knowledge IDs")
	return nil
}

// DeleteBySourceIDList removes points from the collection based on source IDs
func (w *weaviateRepository) DeleteBySourceIDList(ctx context.Context,
	sourceIDList []string, dimension int, knowledgeType string,
) error {
	log := logger.GetLogger(ctx)
	if len(sourceIDList) == 0 {
		log.Warn("[Weaviate] Empty Source ID list provided for deletion, skipping")
		return nil
	}
	collectionName := w.getCollectionName(dimension)
	log.Infof("[Weaviate] Deleting indices by source IDs from %s, count: %d", collectionName, len(sourceIDList))

	//define filter
	filter := w.client.Batch().ObjectsBatchDeleter().
		WithClassName(collectionName).
		WithWhere(filters.Where().
			WithPath([]string{fieldSourceID}).
			WithOperator(filters.ContainsAny).
			WithValueText(sourceIDList...)).
		WithOutput("minimal")

	// Execute deletion
	if _, err := filter.Do(ctx); err != nil {
		log.Errorf("[Weaviate] Failed to delete by source IDs: %v", err)
		return fmt.Errorf("failed to delete by source IDs: %w", err)
	}
	log.Infof("[Weaviate] Successfully deleted documents by source IDs")
	return nil
}

// BatchUpdateChunkEnabledStatus updates the enabled status of chunks in batch
func (w *weaviateRepository) BatchUpdateChunkEnabledStatus(ctx context.Context, chunkStatusMap map[string]bool) error {
	log := logger.GetLogger(ctx)
	if len(chunkStatusMap) == 0 {
		log.Warn("[Weaviate] Empty chunk status map provided, skipping")
		return nil
	}

	log.Infof("[Weaviate] Batch updating chunk enabled status, count: %d", len(chunkStatusMap))

	// Get all collections
	collections, err := w.ListCollections(ctx)
	if err != nil {
		log.Errorf("[Weaviate] Failed to list collections: %v", err)
		return fmt.Errorf("failed to list collections: %w", err)
	}

	// Update in all matching collections
	for _, collectionName := range collections {
		// Only process collections that start with our base name
		if len(collectionName) <= len(w.collectionBaseName) ||
			collectionName[:len(w.collectionBaseName)] != w.collectionBaseName {
			continue
		}
		for chunkID, enabled := range chunkStatusMap {
			if err != nil {
				log.Errorf("[Weaviate] Failed to search ID by chunk ID %s in %s: %v", chunkID, collectionName, err)
				continue
			}
			err = w.client.Data().Updater().
				WithClassName(collectionName).
				WithID(chunkID).
				WithProperties(map[string]interface{}{
					fieldIsEnabled: enabled,
				}).
				Do(ctx)

			isEnabled := "enabled"
			if !enabled {
				isEnabled = "disabled"
			}

			if err != nil {
				log.Errorf("[Weaviate] Failed to update chunk %s status in %s: %v", isEnabled, collectionName, err)
				continue
			}
		}
	}
	log.Infof("[Weaviate] Batch update chunk enabled status completed")
	return nil
}

// BatchUpdateChunkTagID updates the tag ID of chunks in batch
func (w *weaviateRepository) BatchUpdateChunkTagID(ctx context.Context, chunkTagMap map[string]string) error {
	log := logger.GetLogger(ctx)
	if len(chunkTagMap) == 0 {
		log.Warn("[Weaviate] Empty chunk tag map provided, skipping")
		return nil
	}

	log.Infof("[Weaviate] Batch updating chunk tag ID, count: %d", len(chunkTagMap))

	// Get all collections
	collections, err := w.ListCollections(ctx)
	if err != nil {
		log.Errorf("[Weaviate] Failed to list collections: %v", err)
		return fmt.Errorf("failed to list collections: %w", err)
	}

	for _, collectionName := range collections {
		// Only process collections that start with our base name
		if len(collectionName) <= len(w.collectionBaseName) ||
			collectionName[:len(w.collectionBaseName)] != w.collectionBaseName {
			continue
		}

		for chunkID, tagID := range chunkTagMap {
			if err != nil {
				log.Warnf("[Weaviate] Failed to search ID by chunk ID %s in %s: %v", chunkID, collectionName, err)
				continue
			}
			err = w.client.Data().Updater().
				WithClassName(collectionName).
				WithID(chunkID).
				WithProperties(map[string]interface{}{
					fieldTagID: tagID,
				}).
				Do(ctx)
			if err != nil {
				log.Warnf("[Weaviate] Failed to update chunk %s tag ID in %s: %v", chunkID, collectionName, err)
				continue
			}
		}
	}
	log.Infof("[Weaviate] Batch update chunk tag ID completed")
	return nil

}

func (w *weaviateRepository) getBaseFilter(params types.RetrieveParams) *filters.WhereBuilder {
	var operands []*filters.WhereBuilder
	operands = append(operands, filters.Where().
		WithPath([]string{fieldIsEnabled}).
		WithOperator(filters.Equal).
		WithValueBoolean(true))

	if len(params.KnowledgeBaseIDs) > 0 {
		operands = append(operands, filters.Where().
			WithPath([]string{fieldKnowledgeBaseID}).
			WithOperator(filters.ContainsAny).
			WithValueText(params.KnowledgeBaseIDs...))
	}
	if len(params.KnowledgeIDs) > 0 {
		operands = append(operands, filters.Where().
			WithPath([]string{fieldKnowledgeID}).
			WithOperator(filters.ContainsAny).
			WithValueText(params.KnowledgeIDs...))
	}

	if len(params.TagIDs) > 0 {
		operands = append(operands, filters.Where().
			WithPath([]string{fieldTagID}).
			WithOperator(filters.ContainsAny).
			WithValueText(params.TagIDs...))
	}
	if len(params.ExcludeKnowledgeIDs) > 0 {
		operands = append(operands, filters.Where().
			WithPath([]string{fieldKnowledgeID}).
			WithOperator(filters.NotEqual).
			WithValueText(params.ExcludeKnowledgeIDs...))
	}
	if len(params.ExcludeChunkIDs) > 0 {
		operands = append(operands, filters.Where().
			WithPath([]string{fieldChunkID}).
			WithOperator(filters.NotEqual).
			WithValueText(params.ExcludeChunkIDs...))
	}

	return filters.Where().
		WithOperator(filters.And).
		WithOperands(operands)
}

// Retrieve dispatches the retrieval operation to the appropriate method based on retriever type
func (w *weaviateRepository) Retrieve(ctx context.Context,
	params types.RetrieveParams,
) ([]*types.RetrieveResult, error) {
	log := logger.GetLogger(ctx)
	log.Debugf("[Weaviate] Processing retrieval request of type: %s", params.RetrieverType)

	switch params.RetrieverType {
	case types.VectorRetrieverType:
		return w.VectorRetrieve(ctx, params)
	case types.KeywordsRetrieverType:
		return w.KeywordsRetrieve(ctx, params)
	}

	err := fmt.Errorf("invalid retriever type: %v", params.RetrieverType)
	log.Errorf("[Weaviate] %v", err)
	return nil, err
}

// VectorRetrieve performs vector similarity search
func (w *weaviateRepository) VectorRetrieve(ctx context.Context,
	params types.RetrieveParams,
) ([]*types.RetrieveResult, error) {
	log := logger.GetLogger(ctx)
	dimension := len(params.Embedding)
	log.Infof("[Weaviate] Vector retrieval: dim=%d, topK=%d, threshold=%.4f",
		dimension, params.TopK, params.Threshold)

	// Get collection name based on embedding dimension
	collectionName := w.getCollectionName(dimension)

	// Check if collection exists
	hasCollection, err := w.client.Schema().ClassExistenceChecker().WithClassName(collectionName).Do(ctx)
	if err != nil {
		log.Errorf("[Weaviate] Failed to check collection existence: %v", err)
		return nil, fmt.Errorf("failed to check collection: %w", err)
	}
	if !hasCollection {
		log.Warnf("[Weaviate] Collection %s does not exist, returning empty results", collectionName)
		return buildRetrieveResult(nil, types.VectorRetrieverType), nil
	}

	where := w.getBaseFilter(params)
	limit := params.TopK
	scoreThreshold := float32(params.Threshold)
	fields := getEmbeddingFields()
	result, err := w.client.GraphQL().Get().WithClassName(collectionName).
		WithWhere(where).
		WithLimit(limit).
		WithFields(fields...).
		WithNearVector(w.client.GraphQL().NearVectorArgBuilder().
			WithVector(params.Embedding).
			WithCertainty(scoreThreshold)).
		Do(ctx)

	if err != nil {
		log.Errorf("[Weaviate] Vector search failed: %v", err)
		return nil, fmt.Errorf("failed to search: %w", err)
	}
	if len(result.Errors) > 0 {
		log.Errorf("[Weaviate] Vector search failed: %v", result.Errors)
		return nil, fmt.Errorf("graphql search failed: %s", result.Errors[0].Message)
	}

	data, ok := result.Data["Get"].(map[string]interface{})
	if !ok || data[collectionName] == nil {
		log.Warnf("[Weaviate] No vector matches found that meet threshold %.4f", params.Threshold)
		return buildRetrieveResult(nil, types.VectorRetrieverType), nil
	}
	items := data[collectionName].([]interface{})
	results := parseGraphQLResponse(items, collectionName, types.MatchTypeEmbedding)

	if len(results) == 0 {
		log.Warnf("[Weaviate] No vector matches found that meet threshold %.4f", params.Threshold)
	} else {
		log.Infof("[Weaviate] Vector retrieval found %d results", len(results))
		log.Debugf("[Weaviate] Top result score: %.4f", results[0].Score)
	}

	return buildRetrieveResult(results, types.VectorRetrieverType), nil
}

// KeywordsRetrieve performs keyword-based search in document content
// This searches across all collections since keyword search doesn't depend on dimension
func (w *weaviateRepository) KeywordsRetrieve(ctx context.Context,
	params types.RetrieveParams,
) ([]*types.RetrieveResult, error) {
	log := logger.GetLogger(ctx)
	log.Infof("[Weaviate] Performing keywords retrieval with query: %s, topK: %d", params.Query, params.TopK)

	// Get all collections that match our base name pattern
	collections, err := w.ListCollections(ctx)
	if err != nil {
		log.Errorf("[Weaviate] Failed to list collections: %v", err)
		return nil, fmt.Errorf("failed to list collections: %w", err)
	}

	var allResults []*types.IndexWithScore

	for _, collectionName := range collections {
		log.Debugf("[Weaviate] Checking collection: %s", collectionName)
		// Only process collections that start with our base name
		if len(collectionName) <= len(w.collectionBaseName) ||
			collectionName[:len(w.collectionBaseName)] != w.collectionBaseName {
			log.Debugf("[Weaviate] Skipping collection %s (doesn't match base name %s)", collectionName, w.collectionBaseName)
			continue
		}

		filter := w.getBaseFilter(params)

		//bm25 search
		bm25 := w.client.GraphQL().Bm25ArgBuilder().
			WithQuery(params.Query).
			WithProperties([]string{fieldContent}...)

		fields := getKeywordsFields()

		result, err := w.client.GraphQL().Get().WithClassName(collectionName).
			WithWhere(filter).
			WithLimit(params.TopK).
			WithFields(fields...).
			WithBM25(bm25).
			Do(ctx)

		if err != nil {
			log.Errorf("[Weaviate] keywords search failed: %v", err)
			return nil, fmt.Errorf("failed to search: %w", err)
		}
		if len(result.Errors) > 0 {
			log.Errorf("[Weaviate] keywords search failed: %v", result.Errors)
			return nil, fmt.Errorf("graphql search failed: %s", result.Errors[0].Message)
		}
		data, ok := result.Data["Get"].(map[string]interface{})
		if !ok || data[collectionName] == nil {
			log.Warnf("[Weaviate] No keywords matches found that meet threshold %.4f", params.Threshold)
			continue
		}
		items := data[collectionName].([]interface{})
		results := parseGraphQLResponse(items, collectionName, types.MatchTypeKeywords)
		allResults = append(allResults, results...)
	}

	// Limit results to topK
	if len(allResults) > params.TopK {
		allResults = allResults[:params.TopK]
	}

	if len(allResults) == 0 {
		log.Warnf("[Weaviate] No keyword matches found for query: %s", params.Query)
	} else {
		log.Infof("[Weaviate] Keywords retrieval found %d results", len(allResults))
	}
	return buildRetrieveResult(allResults, types.KeywordsRetrieverType), nil
}

// CopyIndices copies index data from source knowledge base to target knowledge base
func (w *weaviateRepository) CopyIndices(ctx context.Context,
	sourceKnowledgeBaseID string,
	sourceToTargetKBIDMap map[string]string,
	sourceToTargetChunkIDMap map[string]string,
	targetKnowledgeBaseID string,
	dimension int,
	knowledgeType string,
) error {
	log := logger.GetLogger(ctx)
	log.Infof("[Weaviate] Copying indices from %s to %s, count: %d",
		sourceKnowledgeBaseID, targetKnowledgeBaseID, len(sourceToTargetChunkIDMap))

	if len(sourceToTargetChunkIDMap) == 0 {
		return nil
	}

	collectionName := w.getCollectionName(dimension)
	batchSize := 64
	var lastID string
	totalCopied := 0
	fields := getVectorFields()

	for {
		result, err := w.client.GraphQL().Get().
			WithClassName(collectionName).
			WithWhere(filters.Where().
				WithPath([]string{fieldKnowledgeBaseID}).
				WithOperator(filters.Equal).
				WithValueString(sourceKnowledgeBaseID)).
			WithLimit(batchSize).
			WithFields(fields...).
			WithAfter(lastID).
			Do(ctx)
		if err != nil {
			log.Errorf("[Weaviate] Failed to query source points: %v", err)
			return err
		}

		objects, ok := result.Data["Get"].(map[string]interface{})[collectionName].([]interface{})
		if !ok || len(objects) == 0 {
			break
		}
		log.Infof("[Weaviate] Found %d source points in batch", len(objects))

		batcher := w.client.Batch().ObjectsBatcher()
		currentBatchCount := 0

		targetObjects := make([]*models.Object, 0, len(objects))
		for _, obj := range objects {
			data, ok := obj.(map[string]interface{})
			if !ok {
				continue
			}
			additional, ok := data["_additional"].(map[string]interface{})
			if !ok {
				continue
			}

			lastID = additional["id"].(string)

			sourceChunkID, ok := data[fieldChunkID].(string)
			if !ok {
				continue
			}
			sourceKnowledgeID, ok := data[fieldKnowledgeID].(string)
			if !ok {
				continue
			}
			originalSourceID, ok := data[fieldSourceID].(string)
			if !ok {
				continue
			}

			targetChunkID, ok1 := sourceToTargetChunkIDMap[sourceChunkID]
			targetKnowledgeID, ok2 := sourceToTargetKBIDMap[sourceKnowledgeID]
			if !ok1 || !ok2 {
				continue
			}

			// Handle SourceID transformation for generated questions
			// Generated questions have SourceID format: {chunkID}-{questionID}
			// Regular chunks have SourceID == ChunkID
			var targetSourceID string
			if originalSourceID == sourceChunkID {
				// Regular chunk, use targetChunkID as SourceID
				targetSourceID = targetChunkID
			} else if strings.HasPrefix(originalSourceID, sourceChunkID+"-") {
				// This is a generated question, preserve the questionID part
				questionID := strings.TrimPrefix(originalSourceID, sourceChunkID+"-")
				targetSourceID = fmt.Sprintf("%s-%s", targetChunkID, questionID)
			} else {
				// For other complex scenarios, generate new unique SourceID
				targetSourceID = uuid.New().String()
			}

			vectorRaw, ok := additional["vector"].([]interface{})
			if !ok {
				continue
			}
			vector := make([]float32, len(vectorRaw))
			for i, v := range vectorRaw {
				vector[i] = float32(v.(float64))
			}

			isEnabled := true
			newObj := &models.Object{
				Class: collectionName,
				ID:    strfmt.UUID(uuid.New().String()),
				Properties: map[string]interface{}{
					fieldContent:         data[fieldContent],
					fieldSourceID:        targetSourceID,
					fieldSourceType:      data[fieldSourceType],
					fieldChunkID:         targetChunkID,
					fieldKnowledgeID:     targetKnowledgeID,
					fieldKnowledgeBaseID: targetKnowledgeBaseID,
					fieldTagID:           data[fieldTagID],
					fieldIsEnabled:       isEnabled,
				},
				Vector: vector,
			}
			targetObjects = append(targetObjects, newObj)
			currentBatchCount++
		}
		if len(targetObjects) > 0 {
			resp, err := batcher.WithObjects(targetObjects...).Do(ctx)
			if err != nil {
				return fmt.Errorf("batch upsert failed: %w", err)
			}
			for _, r := range resp {
				if r.Result.Errors != nil {
					log.Errorf("[Weaviate] Object error: %v", r.Result.Errors.Error[0].Message)
				}
			}
			totalCopied += len(targetObjects)
			log.Infof("[Weaviate] Successfully copied batch, total: %d", totalCopied)
		}
	}
	log.Infof("[Weaviate] Index copy completed, total copied: %d", totalCopied)
	return nil
}

func (w *weaviateRepository) ListCollections(ctx context.Context) ([]string, error) {
	schema, err := w.client.Schema().Getter().Do(ctx)
	if err != nil {
		return nil, fmt.Errorf("weaviate 获取 schema 失败: %w", err)
	}
	var collectionNames []string
	for _, class := range schema.Classes {
		collectionNames = append(collectionNames, class.Class)
	}

	return collectionNames, nil
}

func createPayload(embedding *WeaviateVectorEmbedding) map[string]interface{} {
	payload := map[string]any{
		fieldContent:         embedding.Content,
		fieldSourceID:        embedding.SourceID,
		fieldSourceType:      int64(embedding.SourceType),
		fieldChunkID:         embedding.ChunkID,
		fieldKnowledgeID:     embedding.KnowledgeID,
		fieldKnowledgeBaseID: embedding.KnowledgeBaseID,
		fieldTagID:           embedding.TagID,
		fieldIsEnabled:       embedding.IsEnabled,
	}
	return payload
}

func buildRetrieveResult(results []*types.IndexWithScore, retrieverType types.RetrieverType) []*types.RetrieveResult {
	return []*types.RetrieveResult{
		{
			Results:             results,
			RetrieverEngineType: types.WeaviateRetrieverEngineType,
			RetrieverType:       retrieverType,
			Error:               nil,
		},
	}
}

func getKeywordsFields() []graphql.Field {
	return []graphql.Field{
		{Name: fieldContent},
		{Name: fieldSourceID},
		{Name: fieldSourceType},
		{Name: fieldChunkID},
		{Name: fieldKnowledgeID},
		{Name: fieldKnowledgeBaseID},
		{Name: fieldTagID},
		{
			Name: "_additional",
			Fields: []graphql.Field{
				{Name: "id"},
				{Name: "score"},
			},
		},
	}
}

func getEmbeddingFields() []graphql.Field {
	return []graphql.Field{
		{Name: fieldContent},
		{Name: fieldSourceID},
		{Name: fieldSourceType},
		{Name: fieldChunkID},
		{Name: fieldKnowledgeID},
		{Name: fieldKnowledgeBaseID},
		{Name: fieldTagID},
		{
			Name: "_additional",
			Fields: []graphql.Field{
				{Name: "id"},
				{Name: "certainty"},
			},
		},
	}
}

func getVectorFields() []graphql.Field {
	return []graphql.Field{
		{Name: fieldContent},
		{Name: fieldSourceID},
		{Name: fieldSourceType},
		{Name: fieldChunkID},
		{Name: fieldKnowledgeID},
		{Name: fieldKnowledgeBaseID},
		{Name: fieldTagID},
		{
			Name: "_additional",
			Fields: []graphql.Field{
				{Name: "id"},
				{Name: "vector"},
			},
		},
	}
}

// parseGraphQLResponse parses the GraphQL response for vector embeddings or keyword search results.
func parseGraphQLResponse(items []interface{}, collectionName string, matchType types.MatchType) []*types.IndexWithScore {
	var results []*types.IndexWithScore
	var additionalName string
	if matchType == types.MatchTypeEmbedding {
		additionalName = "certainty"
	} else {
		additionalName = "score"
	}
	for _, item := range items {
		obj := item.(map[string]interface{})
		additional := obj["_additional"].(map[string]interface{})

		pointID := additional["id"].(string)
		score := 0.0
		if s, ok := additional[additionalName].(float64); ok {
			if matchType == types.MatchTypeKeywords {
				score = 1.0
			} else {
				score = s
			}
		}

		getString := func(key string) string {
			if v, ok := obj[key].(string); ok {
				return v
			}
			return ""
		}
		getInt := func(key string) int {
			if v, ok := obj[key].(float64); ok {
				return int(v)
			}
			return 0
		}

		embedding := &WeaviateVectorEmbeddingWithScore{
			WeaviateVectorEmbedding{
				Content:         getString(fieldContent),
				SourceID:        getString(fieldSourceID),
				SourceType:      getInt(fieldSourceType),
				ChunkID:         getString(fieldChunkID),
				KnowledgeID:     getString(fieldKnowledgeID),
				KnowledgeBaseID: getString(fieldKnowledgeBaseID),
				TagID:           getString(fieldTagID),
			},
			float64(score),
		}

		results = append(results, fromWeaviateVectorEmbedding(pointID, embedding, matchType))
	}
	return results
}

// Ref: https://github.com/weaviate/weaviate/blob/b4aec91c6fe464df50e9fa1e2d643322fbb85679/entities/vectorindex/hnsw/config.go#L27
func (w *weaviateRepository) calculateStorageSize(embedding *WeaviateVectorEmbedding) int64 {
	// Payload fields
	payloadSizeBytes := int64(0)
	payloadSizeBytes += int64(len(embedding.Content))         // content string
	payloadSizeBytes += int64(len(embedding.SourceID))        // source_id string
	payloadSizeBytes += int64(len(embedding.ChunkID))         // chunk_id string
	payloadSizeBytes += int64(len(embedding.KnowledgeID))     // knowledge_id string
	payloadSizeBytes += int64(len(embedding.KnowledgeBaseID)) // knowledge_base_id string
	payloadSizeBytes += 8                                     // source_type int64

	// Vector storage and index
	var vectorSizeBytes int64 = 0
	var hnswIndexBytes int64 = 0
	if embedding.Embedding != nil {
		dimensions := int64(len(embedding.Embedding))
		vectorSizeBytes = dimensions * 4

		// HNSW graph links per vector: M×2 neighbors in layer 0, ~8 bytes per link
		// (4 bytes for neighbor ID + multi-layer amortization).
		// Graph link count depends on M, NOT on vector dimensions.
		const hnswM = 32
		hnswIndexBytes = hnswM * 2 * 8
	}

	// ID tracker metadata: 24 bytes per vector
	// (forward refs + backward refs + version tracking = 8 + 8 + 8)
	const idTrackerBytes int64 = 24

	totalSizeBytes := payloadSizeBytes + vectorSizeBytes + hnswIndexBytes + idTrackerBytes
	return totalSizeBytes
}

// toWeaviateVectorEmbedding converts IndexInfo to Weaviate payload format
func toWeaviateVectorEmbedding(embedding *types.IndexInfo, additionalParams map[string]interface{}) *WeaviateVectorEmbedding {
	vector := &WeaviateVectorEmbedding{
		Content:         embedding.Content,
		SourceID:        embedding.SourceID,
		SourceType:      int(embedding.SourceType),
		ChunkID:         embedding.ChunkID,
		KnowledgeID:     embedding.KnowledgeID,
		KnowledgeBaseID: embedding.KnowledgeBaseID,
		TagID:           embedding.TagID,
		IsEnabled:       embedding.IsEnabled,
	}
	if additionalParams != nil && slices.Contains(slices.Collect(maps.Keys(additionalParams)), fieldEmbedding) {
		if embeddingMap, ok := additionalParams[fieldEmbedding].(map[string][]float32); ok {
			vector.Embedding = embeddingMap[embedding.SourceID]
		}
	}
	return vector
}

// fromWeaviateVectorEmbedding converts Weaviate point to IndexWithScore domain model
func fromWeaviateVectorEmbedding(id string,
	embedding *WeaviateVectorEmbeddingWithScore,
	matchType types.MatchType,
) *types.IndexWithScore {
	return &types.IndexWithScore{
		ID:              id,
		SourceID:        embedding.SourceID,
		SourceType:      types.SourceType(embedding.SourceType),
		ChunkID:         embedding.ChunkID,
		KnowledgeID:     embedding.KnowledgeID,
		KnowledgeBaseID: embedding.KnowledgeBaseID,
		TagID:           embedding.TagID,
		Content:         embedding.Content,
		Score:           embedding.Score,
		MatchType:       matchType,
	}
}

// tokenizeQuery splits a query string into tokens for OR-based full-text search.
// It uses jieba for professional Chinese word segmentation.
func tokenizeQuery(query string) []string {
	query = strings.TrimSpace(query)
	if query == "" {
		return nil
	}

	// Use jieba for segmentation (search mode for better recall)
	words := types.Jieba.CutForSearch(query, true)

	// Filter and deduplicate
	seen := make(map[string]bool)
	result := make([]string, 0, len(words))
	for _, word := range words {
		word = strings.TrimSpace(strings.ToLower(word))
		// Skip empty, single-char, and already seen words
		if utf8.RuneCountInString(word) < 2 || seen[word] {
			continue
		}
		seen[word] = true
		result = append(result, word)
	}

	return result
}