xinmi/internal/application/repository/retriever/elasticsearch/v8/repository.go

package v8

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

	elasticsearchRetriever "github.com/Tencent/WeKnora/internal/application/repository/retriever/elasticsearch"
	"github.com/Tencent/WeKnora/internal/config"
	"github.com/Tencent/WeKnora/internal/logger"
	typesLocal "github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/elastic/go-elasticsearch/v8"
	"github.com/elastic/go-elasticsearch/v8/typedapi/core/search"
	"github.com/elastic/go-elasticsearch/v8/typedapi/types"
	"github.com/elastic/go-elasticsearch/v8/typedapi/types/enums/scriptlanguage"
	"github.com/google/uuid"
)

// elasticsearchRepository implements the RetrieveEngineRepository interface for Elasticsearch v8
type elasticsearchRepository struct {
	client           *elasticsearch.TypedClient // Elasticsearch client instance
	index            string                     // Name of the Elasticsearch index to use
	useKeywordSuffix bool                       // Whether to append .keyword suffix to ID field names in queries
	numberOfShards   int                        // Shard count for index creation (0 = ES default)
	numberOfReplicas int                        // Replica count for index creation (-1 = unset, use ES default)
}

// NewElasticsearchEngineRepository creates and initializes a new Elasticsearch v8 repository.
// indexCfg is optional — pass nil to use env var / default values (env path).
func NewElasticsearchEngineRepository(client *elasticsearch.TypedClient,
	config *config.Config,
	indexCfg *typesLocal.IndexConfig,
) interfaces.RetrieveEngineRepository {
	log := logger.GetLogger(context.Background())
	log.Info("[Elasticsearch] Initializing Elasticsearch v8 retriever engine repository")

	indexName := typesLocal.ResolveIndexName(indexCfg, "ELASTICSEARCH_INDEX", "xwrag_default")

	// Create repository instance and ensure index exists
	res := &elasticsearchRepository{
		client:           client,
		index:            indexName,
		numberOfShards:   indexCfg.GetNumberOfShards(0),
		numberOfReplicas: indexCfg.GetNumberOfReplicas(-1),
	}
	if err := res.createIndexIfNotExists(context.Background()); err != nil {
		log.Errorf("[Elasticsearch] Failed to create index: %v", err)
	} else {
		log.Info("[Elasticsearch] Successfully initialized repository")
	}
	res.detectFieldTypes(context.Background())
	return res
}

// idField returns the query field name for an ID field, appending ".keyword"
// suffix when the index uses text-type mappings with keyword sub-fields.
func (e *elasticsearchRepository) idField(name string) string {
	if e.useKeywordSuffix {
		return name + ".keyword"
	}
	return name
}

// detectFieldTypes inspects the index mapping to determine whether ID fields
// are mapped as "keyword" (no suffix needed) or "text" (needs ".keyword" suffix).
func (e *elasticsearchRepository) detectFieldTypes(ctx context.Context) {
	log := logger.GetLogger(ctx)

	mappingResp, err := e.client.Indices.GetMapping().Index(e.index).Do(ctx)
	if err != nil {
		log.Warnf("[Elasticsearch] Failed to get index mapping, defaulting to .keyword suffix: %v", err)
		e.useKeywordSuffix = true
		return
	}

	indexMapping, ok := mappingResp[e.index]
	if !ok {
		log.Warnf("[Elasticsearch] Index %s not found in mapping response, defaulting to .keyword suffix", e.index)
		e.useKeywordSuffix = true
		return
	}

	if prop, ok := indexMapping.Mappings.Properties["chunk_id"]; ok {
		propBytes, err := json.Marshal(prop)
		if err == nil {
			var propInfo struct {
				Type string `json:"type"`
			}
			if json.Unmarshal(propBytes, &propInfo) == nil && propInfo.Type == "keyword" {
				e.useKeywordSuffix = false
				log.Infof("[Elasticsearch] Detected keyword type for ID fields, querying without .keyword suffix")
				return
			}
		}
		e.useKeywordSuffix = true
		log.Infof("[Elasticsearch] ID fields are not keyword type, querying with .keyword suffix")
		return
	}

	log.Infof("[Elasticsearch] No mapping detected for chunk_id (empty index?), defaulting to .keyword suffix")
	e.useKeywordSuffix = true
}

// EngineType returns the type of retriever engine (Elasticsearch)
func (e *elasticsearchRepository) EngineType() typesLocal.RetrieverEngineType {
	return typesLocal.ElasticsearchRetrieverEngineType
}

// Support returns the retrieval types supported by this repository (Keywords and Vector)
func (e *elasticsearchRepository) Support() []typesLocal.RetrieverType {
	return []typesLocal.RetrieverType{typesLocal.KeywordsRetrieverType, typesLocal.VectorRetrieverType}
}

// calculateStorageSize estimates the storage size in bytes for a single index document
func (e *elasticsearchRepository) calculateStorageSize(embedding *elasticsearchRetriever.VectorEmbedding) int64 {
	// 1. Content text size
	contentSizeBytes := int64(len(embedding.Content))

	// 2. Vector embedding size
	var vectorSizeBytes int64 = 0
	if embedding.Embedding != nil {
		// 4 bytes per dimension (full precision float)
		vectorSizeBytes = int64(len(embedding.Embedding) * 4)
	}

	// 3. Metadata size (IDs, timestamps, and other fixed overhead)
	metadataSizeBytes := int64(250) // Approximately 250 bytes of metadata

	// 4. Index overhead (Elasticsearch index expansion factor ~1.5)
	indexOverheadBytes := (contentSizeBytes + vectorSizeBytes) * 5 / 10

	// Total size in bytes
	totalSizeBytes := contentSizeBytes + vectorSizeBytes + metadataSizeBytes + indexOverheadBytes
	return totalSizeBytes
}

// EstimateStorageSize calculates the estimated storage size for a list of indices
// Returns the total size in bytes
func (e *elasticsearchRepository) EstimateStorageSize(ctx context.Context,
	indexInfoList []*typesLocal.IndexInfo, params map[string]any,
) int64 {
	var totalStorageSize int64 = 0
	for _, embedding := range indexInfoList {
		embeddingDB := elasticsearchRetriever.ToDBVectorEmbedding(embedding, params)
		totalStorageSize += e.calculateStorageSize(embeddingDB)
	}
	logger.GetLogger(ctx).Infof(
		"[Elasticsearch] Storage size for %d indices: %d bytes", len(indexInfoList), totalStorageSize,
	)
	return totalStorageSize
}

// Save stores a single index document in Elasticsearch
// Returns an error if the operation fails
func (e *elasticsearchRepository) Save(ctx context.Context,
	embedding *typesLocal.IndexInfo,
	additionalParams map[string]any,
) error {
	log := logger.GetLogger(ctx)
	log.Debugf("[Elasticsearch] Saving index for chunk ID: %s", embedding.ChunkID)

	// Convert to database format
	embeddingDB := elasticsearchRetriever.ToDBVectorEmbedding(embedding, additionalParams)
	if len(embeddingDB.Embedding) == 0 {
		err := fmt.Errorf("empty embedding vector for chunk ID: %s", embedding.ChunkID)
		log.Errorf("[Elasticsearch] %v", err)
		return err
	}

	// Index the document
	resp, err := e.client.Index(e.index).Request(embeddingDB).Do(ctx)
	if err != nil {
		log.Errorf("[Elasticsearch] Failed to save index: %v", err)
		return err
	}

	log.Infof("[Elasticsearch] Successfully saved index for chunk ID: %s, document ID: %s", embedding.ChunkID, resp.Id_)
	return nil
}

// BatchSave stores multiple index documents in Elasticsearch using bulk API
// Returns an error if the operation fails
func (e *elasticsearchRepository) BatchSave(ctx context.Context,
	embeddingList []*typesLocal.IndexInfo, additionalParams map[string]any,
) error {
	log := logger.GetLogger(ctx)
	if len(embeddingList) == 0 {
		log.Warn("[Elasticsearch] Empty list provided to BatchSave, skipping")
		return nil
	}

	log.Infof("[Elasticsearch] Batch saving %d indices", len(embeddingList))
	indexRequest := e.client.Bulk().Index(e.index)

	// Add each document to the bulk request
	for _, embedding := range embeddingList {
		embeddingDB := elasticsearchRetriever.ToDBVectorEmbedding(embedding, additionalParams)
		err := indexRequest.CreateOp(types.CreateOperation{Index_: &e.index}, embeddingDB)
		if err != nil {
			log.Errorf("[Elasticsearch] Failed to create bulk operation: %v", err)
			return fmt.Errorf("failed to create op: %w", err)
		}
		log.Debugf("[Elasticsearch] Added chunk ID %s to bulk request", embedding.ChunkID)
	}

	// Execute the bulk request
	_, err := indexRequest.Do(ctx)
	if err != nil {
		log.Errorf("[Elasticsearch] Failed to execute bulk operation: %v", err)
		return fmt.Errorf("failed to do bulk: %w", err)
	}

	log.Infof("[Elasticsearch] Successfully batch saved %d indices", len(embeddingList))
	return nil
}

// DeleteByChunkIDList removes documents from the index based on chunk IDs
// Returns an error if the delete operation fails
func (e *elasticsearchRepository) DeleteByChunkIDList(ctx context.Context, chunkIDList []string, dimension int, knowledgeType string) error {
	log := logger.GetLogger(ctx)
	if len(chunkIDList) == 0 {
		log.Warn("[Elasticsearch] Empty chunk ID list provided for deletion, skipping")
		return nil
	}

	log.Infof("[Elasticsearch] Deleting indices by chunk IDs, count: %d", len(chunkIDList))
	// Use DeleteByQuery to delete all documents matching the chunk IDs
	_, err := e.client.DeleteByQuery(e.index).Query(&types.Query{
		Terms: &types.TermsQuery{TermsQuery: map[string]types.TermsQueryField{e.idField("chunk_id"): chunkIDList}},
	}).Do(ctx)
	if err != nil {
		log.Errorf("[Elasticsearch] Failed to delete by chunk IDs: %v", err)
		return fmt.Errorf("failed to delete by query: %w", err)
	}

	log.Infof("[Elasticsearch] Successfully deleted documents by chunk IDs")
	return nil
}

// DeleteBySourceIDList removes documents from the index based on source IDs
// Returns an error if the delete operation fails
func (e *elasticsearchRepository) DeleteBySourceIDList(ctx context.Context, sourceIDList []string, dimension int, knowledgeType string) error {
	log := logger.GetLogger(ctx)
	if len(sourceIDList) == 0 {
		log.Warn("[Elasticsearch] Empty source ID list provided for deletion, skipping")
		return nil
	}

	log.Infof("[Elasticsearch] Deleting indices by source IDs, count: %d", len(sourceIDList))
	// Use DeleteByQuery to delete all documents matching the source IDs
	_, err := e.client.DeleteByQuery(e.index).Query(&types.Query{
		Terms: &types.TermsQuery{TermsQuery: map[string]types.TermsQueryField{e.idField("source_id"): sourceIDList}},
	}).Do(ctx)
	if err != nil {
		log.Errorf("[Elasticsearch] Failed to delete by source IDs: %v", err)
		return fmt.Errorf("failed to delete by query: %w", err)
	}

	log.Infof("[Elasticsearch] Successfully deleted documents by source IDs")
	return nil
}

// DeleteByKnowledgeIDList removes documents from the index based on knowledge IDs
// Returns an error if the delete operation fails
func (e *elasticsearchRepository) DeleteByKnowledgeIDList(ctx context.Context,
	knowledgeIDList []string, dimension int, knowledgeType string,
) error {
	log := logger.GetLogger(ctx)
	if len(knowledgeIDList) == 0 {
		log.Warn("[Elasticsearch] Empty knowledge ID list provided for deletion, skipping")
		return nil
	}

	log.Infof("[Elasticsearch] Deleting indices by knowledge IDs, count: %d", len(knowledgeIDList))
	// Use DeleteByQuery to delete all documents matching the knowledge IDs
	_, err := e.client.DeleteByQuery(e.index).Query(&types.Query{
		Terms: &types.TermsQuery{TermsQuery: map[string]types.TermsQueryField{e.idField("knowledge_id"): knowledgeIDList}},
	}).Do(ctx)
	if err != nil {
		log.Errorf("[Elasticsearch] Failed to delete by knowledge IDs: %v", err)
		return fmt.Errorf("failed to delete by query: %w", err)
	}

	log.Infof("[Elasticsearch] Successfully deleted documents by knowledge IDs")
	return nil
}

// getBaseConds creates the base query conditions for retrieval operations
// Returns a slice of Query objects with must and must_not conditions
// KnowledgeBaseIDs and KnowledgeIDs use AND logic (search specific documents within knowledge bases)
func (e *elasticsearchRepository) getBaseConds(params typesLocal.RetrieveParams) []types.Query {
	must := []types.Query{}

	// KnowledgeBaseIDs and KnowledgeIDs use AND logic
	// - If only KnowledgeBaseIDs: search entire knowledge bases
	// - If only KnowledgeIDs: search specific documents
	// - If both: search specific documents within the knowledge bases (AND)
	if len(params.KnowledgeBaseIDs) > 0 {
		must = append(must, types.Query{Terms: &types.TermsQuery{
			TermsQuery: map[string]types.TermsQueryField{
				e.idField("knowledge_base_id"): params.KnowledgeBaseIDs,
			},
		}})
	}
	if len(params.KnowledgeIDs) > 0 {
		must = append(must, types.Query{Terms: &types.TermsQuery{
			TermsQuery: map[string]types.TermsQueryField{
				e.idField("knowledge_id"): params.KnowledgeIDs,
			},
		}})
	}
	// Filter by tag IDs if specified
	if len(params.TagIDs) > 0 {
		must = append(must, types.Query{Terms: &types.TermsQuery{
			TermsQuery: map[string]types.TermsQueryField{
				e.idField("tag_id"): params.TagIDs,
			},
		}})
	}

	mustNot := make([]types.Query, 0)
	// Exclude disabled chunks (is_enabled = false)
	// Note: Historical data without is_enabled field will be included (not matching must_not)
	mustNot = append(mustNot, types.Query{Term: map[string]types.TermQuery{
		"is_enabled": {Value: false},
	}})
	if len(params.ExcludeKnowledgeIDs) > 0 {
		mustNot = append(mustNot, types.Query{Terms: &types.TermsQuery{
			TermsQuery: map[string]types.TermsQueryField{e.idField("knowledge_id"): params.ExcludeKnowledgeIDs},
		}})
	}
	if len(params.ExcludeChunkIDs) > 0 {
		mustNot = append(mustNot, types.Query{Terms: &types.TermsQuery{
			TermsQuery: map[string]types.TermsQueryField{e.idField("chunk_id"): params.ExcludeChunkIDs},
		}})
	}
	return []types.Query{{Bool: &types.BoolQuery{Must: must, MustNot: mustNot}}}
}

// createIndexIfNotExists checks if the specified index exists and creates it if not
// Returns an error if the operation fails
func (e *elasticsearchRepository) createIndexIfNotExists(ctx context.Context) error {
	log := logger.GetLogger(ctx)
	log.Debugf("[Elasticsearch] Checking if index exists: %s", e.index)

	// Check if index exists
	exists, err := e.client.Indices.Exists(e.index).Do(ctx)
	if err != nil {
		log.Errorf("[Elasticsearch] Failed to check if index exists: %v", err)
		return err
	}

	if exists {
		log.Debugf("[Elasticsearch] Index already exists: %s", e.index)
		return nil
	}

	// Create index if it doesn't exist, with optional shards/replicas settings
	log.Infof("[Elasticsearch] Creating index: %s", e.index)
	createReq := e.client.Indices.Create(e.index)
	if e.numberOfShards > 0 || e.numberOfReplicas >= 0 {
		settings := &types.IndexSettings{}
		if e.numberOfShards > 0 {
			shards := fmt.Sprintf("%d", e.numberOfShards)
			settings.NumberOfShards = &shards
		}
		if e.numberOfReplicas >= 0 {
			replicas := fmt.Sprintf("%d", e.numberOfReplicas)
			settings.NumberOfReplicas = &replicas
		}
		createReq = createReq.Settings(settings)
	}
	_, err = createReq.Do(ctx)
	if err != nil {
		log.Errorf("[Elasticsearch] Failed to create index: %v", err)
		return err
	}

	log.Infof("[Elasticsearch] Index created successfully: %s", e.index)
	return nil
}

// Retrieve dispatches the retrieval operation to the appropriate method based on retriever type
// Returns a slice of RetrieveResult and an error if the operation fails
func (e *elasticsearchRepository) Retrieve(ctx context.Context,
	params typesLocal.RetrieveParams,
) ([]*typesLocal.RetrieveResult, error) {
	log := logger.GetLogger(ctx)
	log.Debugf("[Elasticsearch] Processing retrieval request of type: %s", params.RetrieverType)

	// Route to appropriate retrieval method
	switch params.RetrieverType {
	case typesLocal.VectorRetrieverType:
		return e.VectorRetrieve(ctx, params)
	case typesLocal.KeywordsRetrieverType:
		return e.KeywordsRetrieve(ctx, params)
	}

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

// VectorRetrieve performs vector similarity search using cosine similarity
// Returns a slice of RetrieveResult containing matching documents
func (e *elasticsearchRepository) VectorRetrieve(ctx context.Context,
	params typesLocal.RetrieveParams,
) ([]*typesLocal.RetrieveResult, error) {
	log := logger.GetLogger(ctx)
	log.Infof("[Elasticsearch] Vector retrieval: dim=%d, topK=%d, threshold=%.4f",
		len(params.Embedding), params.TopK, params.Threshold)

	filter := e.getBaseConds(params)

	// Build script scoring query with cosine similarity
	queryVectorJSON, err := json.Marshal(params.Embedding)
	if err != nil {
		log.Errorf("[Elasticsearch] Failed to marshal query vector: %v", err)
		return nil, fmt.Errorf("failed to marshal query embedding: %w", err)
	}

	scoreSource := "cosineSimilarity(params.query_vector, 'embedding')"
	minScore := float32(params.Threshold)
	scriptScore := &types.ScriptScoreQuery{
		Query: types.Query{Bool: &types.BoolQuery{Filter: filter}},
		Script: types.Script{
			Source: &scoreSource,
			Params: map[string]json.RawMessage{
				"query_vector": json.RawMessage(queryVectorJSON),
			},
		},
		MinScore: &minScore,
	}
	// Exclude embedding field from source to reduce response size
	sourceFilter := &types.SourceFilter{
		Excludes: []string{"embedding"},
	}

	log.Debugf("[Elasticsearch] Executing vector search in index: %s", e.index)
	// Execute search with minimum score threshold
	response, err := e.client.Search().Index(e.index).Request(&search.Request{
		Query:   &types.Query{ScriptScore: scriptScore},
		Size:    &params.TopK,
		Source_: sourceFilter,
	}).Do(ctx)
	if err != nil {
		log.Errorf("[Elasticsearch] Vector search failed: %v", err)
		return nil, err
	}

	// Process search results
	var results []*typesLocal.IndexWithScore
	for _, hit := range response.Hits.Hits {
		var embedding *elasticsearchRetriever.VectorEmbeddingWithScore
		if err := json.Unmarshal(hit.Source_, &embedding); err != nil {
			log.Errorf("[Elasticsearch] Failed to unmarshal search result: %v", err)
			return nil, err
		}
		embedding.Score = float64(*hit.Score_)
		results = append(results,
			elasticsearchRetriever.FromDBVectorEmbeddingWithScore(*hit.Id_, embedding, typesLocal.MatchTypeEmbedding))
	}

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

	return []*typesLocal.RetrieveResult{
		{
			Results:             results,
			RetrieverEngineType: typesLocal.ElasticsearchRetrieverEngineType,
			RetrieverType:       typesLocal.VectorRetrieverType,
			Error:               nil,
		},
	}, nil
}

// KeywordsRetrieve performs keyword-based search in document content
// Returns a slice of RetrieveResult containing matching documents
func (e *elasticsearchRepository) KeywordsRetrieve(ctx context.Context,
	params typesLocal.RetrieveParams,
) ([]*typesLocal.RetrieveResult, error) {
	log := logger.GetLogger(ctx)
	log.Infof("[Elasticsearch] Performing keywords retrieval with query: %s, topK: %d", params.Query, params.TopK)

	filter := e.getBaseConds(params)
	// Build must conditions for content matching
	must := []types.Query{
		{Match: map[string]types.MatchQuery{"content": {Query: params.Query}}},
	}
	// Exclude embedding field from source to reduce response size
	sourceFilter := &types.SourceFilter{
		Excludes: []string{"embedding"},
	}

	log.Debugf("[Elasticsearch] Executing keyword search in index: %s", e.index)
	response, err := e.client.Search().Index(e.index).Request(&search.Request{
		Query:   &types.Query{Bool: &types.BoolQuery{Filter: filter, Must: must}},
		Size:    &params.TopK,
		Source_: sourceFilter,
	}).Do(ctx)
	if err != nil {
		log.Errorf("[Elasticsearch] Keywords search failed: %v", err)
		return nil, err
	}

	// Process search results
	var results []*typesLocal.IndexWithScore
	for _, hit := range response.Hits.Hits {
		var embedding *elasticsearchRetriever.VectorEmbeddingWithScore
		if err := json.Unmarshal(hit.Source_, &embedding); err != nil {
			log.Errorf("[Elasticsearch] Failed to unmarshal search result: %v", err)
			return nil, err
		}
		embedding.Score = float64(*hit.Score_)
		results = append(results,
			elasticsearchRetriever.FromDBVectorEmbeddingWithScore(*hit.Id_, embedding, typesLocal.MatchTypeKeywords),
		)
	}

	if len(results) == 0 {
		log.Warnf("[Elasticsearch] No keyword matches found for query: %s", params.Query)
	} else {
		log.Infof("[Elasticsearch] Keywords retrieval found %d results", len(results))
		log.Debugf("[Elasticsearch] Top result score: %.4f", results[0].Score)
	}

	return []*typesLocal.RetrieveResult{
		{
			Results:             results,
			RetrieverEngineType: typesLocal.ElasticsearchRetrieverEngineType,
			RetrieverType:       typesLocal.KeywordsRetrieverType,
			Error:               nil,
		},
	}, nil
}

// CopyIndices 复制索引数据
func (e *elasticsearchRepository) 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(
		"[Elasticsearch] Copying indices from source knowledge base %s to target knowledge base %s, count: %d",
		sourceKnowledgeBaseID, targetKnowledgeBaseID, len(sourceToTargetChunkIDMap),
	)

	if len(sourceToTargetChunkIDMap) == 0 {
		log.Warn("[Elasticsearch] Empty mapping, skipping copy")
		return nil
	}

	// Build query parameters
	params := typesLocal.RetrieveParams{
		KnowledgeBaseIDs: []string{sourceKnowledgeBaseID},
	}

	// Build base query conditions
	filter := e.getBaseConds(params)

	// Set batch processing parameters
	batchSize := 500
	from := 0
	totalCopied := 0

	for {
		// Execute pagination query
		searchResponse, err := e.client.Search().Index(e.index).
			Query(&types.Query{Bool: &types.BoolQuery{Filter: filter}}).
			From(from).
			Size(batchSize).
			Do(ctx)
		if err != nil {
			log.Errorf("[Elasticsearch] Failed to query source index data: %v", err)
			return err
		}

		hitsCount := len(searchResponse.Hits.Hits)
		if hitsCount == 0 {
			break
		}

		log.Infof("[Elasticsearch] Found %d source index data, batch start position: %d", hitsCount, from)

		// Prepare index list for BatchSave
		var indexInfoList []*typesLocal.IndexInfo

		// Collect all embedding vector data for additionalParams
		embeddingMap := make(map[string][]float32)

		for _, hit := range searchResponse.Hits.Hits {
			// Parse source document
			var sourceDoc elasticsearchRetriever.VectorEmbedding
			if err := json.Unmarshal(hit.Source_, &sourceDoc); err != nil {
				log.Errorf("[Elasticsearch] Failed to parse source index data: %v", err)
				continue
			}

			// Get mapped target chunk ID and knowledge ID
			targetChunkID, ok := sourceToTargetChunkIDMap[sourceDoc.ChunkID]
			if !ok {
				log.Warnf("[Elasticsearch] Source chunk %s not found in target mapping, skipping", sourceDoc.ChunkID)
				continue
			}

			targetKnowledgeID, ok := sourceToTargetKBIDMap[sourceDoc.KnowledgeID]
			if !ok {
				log.Warnf(
					"[Elasticsearch] Source knowledge %s not found in target mapping, skipping",
					sourceDoc.KnowledgeID,
				)
				continue
			}

			// Save embedding vector to embeddingMap
			if len(sourceDoc.Embedding) > 0 {
				embeddingMap[targetChunkID] = sourceDoc.Embedding
			}

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

			// Create new index information
			indexInfo := &typesLocal.IndexInfo{
				Content:         sourceDoc.Content,
				SourceID:        targetSourceID,
				SourceType:      typesLocal.SourceType(sourceDoc.SourceType),
				ChunkID:         targetChunkID,
				KnowledgeID:     targetKnowledgeID,
				KnowledgeBaseID: targetKnowledgeBaseID,
			}

			indexInfoList = append(indexInfoList, indexInfo)
			totalCopied++
		}

		// Use BatchSave function to save index
		if len(indexInfoList) > 0 {
			// Add embedding vector to additional parameters
			additionalParams := map[string]any{
				"embedding": embeddingMap,
			}

			if err := e.BatchSave(ctx, indexInfoList, additionalParams); err != nil {
				log.Errorf("[Elasticsearch] Failed to batch save index: %v", err)
				return err
			}

			log.Infof("[Elasticsearch] Successfully copied batch data, batch size: %d, total copied: %d",
				len(indexInfoList), totalCopied)
		}

		// Move to next batch
		from += hitsCount

		// If the number of returned records is less than the request size, it means the last page has been reached
		if hitsCount < batchSize {
			break
		}
	}

	log.Infof("[Elasticsearch] Index copy completed, total copied: %d", totalCopied)
	return nil
}

// BatchUpdateChunkEnabledStatus updates the enabled status of chunks in batch
func (e *elasticsearchRepository) BatchUpdateChunkEnabledStatus(
	ctx context.Context,
	chunkStatusMap map[string]bool,
) error {
	log := logger.GetLogger(ctx)
	if len(chunkStatusMap) == 0 {
		log.Warnf("[Elasticsearch] Chunk status map is empty, skipping update")
		return nil
	}

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

	// Group chunks by enabled status for batch updates
	enabledChunkIDs := make([]string, 0)
	disabledChunkIDs := make([]string, 0)

	for chunkID, enabled := range chunkStatusMap {
		if enabled {
			enabledChunkIDs = append(enabledChunkIDs, chunkID)
		} else {
			disabledChunkIDs = append(disabledChunkIDs, chunkID)
		}
	}

	// Batch update enabled chunks using update_by_query
	if len(enabledChunkIDs) > 0 {
		query := types.NewQuery()
		query.Bool = &types.BoolQuery{
			Must: []types.Query{
				{Terms: &types.TermsQuery{
					TermsQuery: map[string]types.TermsQueryField{
						e.idField("chunk_id"): enabledChunkIDs,
					},
				}},
			},
		}
		source := "ctx._source.is_enabled = true"
		lang := scriptlanguage.Painless
		script := types.Script{
			Source: &source,
			Lang:   &lang,
		}
		_, err := e.client.UpdateByQuery(e.index).Query(query).Script(&script).Do(ctx)
		if err != nil {
			log.Errorf("[Elasticsearch] Failed to update enabled chunks: %v", err)
			return err
		}
		log.Infof("[Elasticsearch] Updated %d chunks to enabled", len(enabledChunkIDs))
	}

	// Batch update disabled chunks using update_by_query
	if len(disabledChunkIDs) > 0 {
		query := types.NewQuery()
		query.Bool = &types.BoolQuery{
			Must: []types.Query{
				{Terms: &types.TermsQuery{
					TermsQuery: map[string]types.TermsQueryField{
						e.idField("chunk_id"): disabledChunkIDs,
					},
				}},
			},
		}
		source := "ctx._source.is_enabled = false"
		lang := scriptlanguage.Painless
		script := types.Script{
			Source: &source,
			Lang:   &lang,
		}
		_, err := e.client.UpdateByQuery(e.index).Query(query).Script(&script).Do(ctx)
		if err != nil {
			log.Errorf("[Elasticsearch] Failed to update disabled chunks: %v", err)
			return err
		}
		log.Infof("[Elasticsearch] Updated %d chunks to disabled", len(disabledChunkIDs))
	}

	log.Infof("[Elasticsearch] Successfully batch updated chunk enabled status")
	return nil
}

// BatchUpdateChunkTagID updates the tag ID of chunks in batch
func (e *elasticsearchRepository) BatchUpdateChunkTagID(
	ctx context.Context,
	chunkTagMap map[string]string,
) error {
	log := logger.GetLogger(ctx)
	if len(chunkTagMap) == 0 {
		log.Warnf("[Elasticsearch] Chunk tag map is empty, skipping update")
		return nil
	}

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

	// Group chunks by tag ID for batch updates
	tagGroups := make(map[string][]string)
	for chunkID, tagID := range chunkTagMap {
		tagGroups[tagID] = append(tagGroups[tagID], chunkID)
	}

	// Batch update chunks for each tag ID using update_by_query
	for tagID, chunkIDs := range tagGroups {
		query := types.NewQuery()
		query.Bool = &types.BoolQuery{
			Must: []types.Query{
				{Terms: &types.TermsQuery{
					TermsQuery: map[string]types.TermsQueryField{
						e.idField("chunk_id"): chunkIDs,
					},
				}},
			},
		}
		source := "ctx._source.tag_id = params.tag_id"
		lang := scriptlanguage.Painless
		script := types.Script{
			Source: &source,
			Lang:   &lang,
			Params: map[string]json.RawMessage{
				"tag_id": json.RawMessage(`"` + tagID + `"`),
			},
		}
		_, err := e.client.UpdateByQuery(e.index).Query(query).Script(&script).Do(ctx)
		if err != nil {
			log.Errorf("[Elasticsearch] Failed to update chunks with tag_id %s: %v", tagID, err)
			return err
		}
		log.Infof("[Elasticsearch] Updated %d chunks to tag_id=%s", len(chunkIDs), tagID)
	}

	log.Infof("[Elasticsearch] Successfully batch updated chunk tag ID")
	return nil
}