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

// Package v7 implements the Elasticsearch v7 retriever engine repository
// It provides functionality for storing and retrieving embeddings using Elasticsearch v7
// The repository supports both single and batch operations for saving and deleting embeddings
// It also supports retrieving embeddings using different retrieval methods
// The repository is used to store and retrieve embeddings for different tasks
// It is used to store and retrieve embeddings for different tasks
package v7

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"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/v7"
	"github.com/elastic/go-elasticsearch/v7/esapi"
	"github.com/google/uuid"
)

type elasticsearchRepository struct {
	client           *elasticsearch.Client
	index            string
	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 v7 repository.
// indexCfg is optional — pass nil to use env var / default values (env path).
func NewElasticsearchEngineRepository(client *elasticsearch.Client,
	config *config.Config,
	indexCfg *typesLocal.IndexConfig,
) interfaces.RetrieveEngineRepository {
	log := logger.GetLogger(context.Background())
	log.Info("[ElasticsearchV7] Initializing Elasticsearch v7 retriever engine repository")

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

	log.Infof("[ElasticsearchV7] Using index: %s", indexName)
	res := &elasticsearchRepository{
		client:           client,
		index:            indexName,
		numberOfShards:   indexCfg.GetNumberOfShards(0),
		numberOfReplicas: indexCfg.GetNumberOfReplicas(-1),
	}
	if err := res.createIndexIfNotExists(context.Background()); err != nil {
		log.Errorf("[ElasticsearchV7] Failed to create index: %v", err)
	}
	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)

	resp, err := e.client.Indices.GetMapping(
		e.client.Indices.GetMapping.WithIndex(e.index),
		e.client.Indices.GetMapping.WithContext(ctx),
	)
	if err != nil {
		log.Warnf("[ElasticsearchV7] Failed to get index mapping, defaulting to .keyword suffix: %v", err)
		e.useKeywordSuffix = true
		return
	}
	defer resp.Body.Close()

	if resp.IsError() {
		log.Warnf("[ElasticsearchV7] GetMapping returned error: %s, defaulting to .keyword suffix", resp.String())
		e.useKeywordSuffix = true
		return
	}

	var result map[string]interface{}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		log.Warnf("[ElasticsearchV7] Failed to parse mapping response: %v, defaulting to .keyword suffix", err)
		e.useKeywordSuffix = true
		return
	}

	indexData, _ := result[e.index].(map[string]interface{})
	if indexData == nil {
		log.Warnf("[ElasticsearchV7] Index %s not found in mapping response, defaulting to .keyword suffix", e.index)
		e.useKeywordSuffix = true
		return
	}
	mappings, _ := indexData["mappings"].(map[string]interface{})
	if mappings == nil {
		e.useKeywordSuffix = true
		return
	}
	properties, _ := mappings["properties"].(map[string]interface{})
	if properties == nil {
		log.Infof("[ElasticsearchV7] No mapping detected for ID fields (empty index?), defaulting to .keyword suffix")
		e.useKeywordSuffix = true
		return
	}
	chunkIDProp, _ := properties["chunk_id"].(map[string]interface{})
	if chunkIDProp == nil {
		e.useKeywordSuffix = true
		return
	}
	if fieldType, _ := chunkIDProp["type"].(string); fieldType == "keyword" {
		e.useKeywordSuffix = false
		log.Infof("[ElasticsearchV7] Detected keyword type for ID fields, querying without .keyword suffix")
	} else {
		e.useKeywordSuffix = true
		log.Infof("[ElasticsearchV7] Detected %s type for ID fields, querying with .keyword suffix", fieldType)
	}
}

// createIndexIfNotExists checks if the specified index exists and creates it if not.
// Uses esapi low-level client since v7 SDK does not have typed API.
func (e *elasticsearchRepository) createIndexIfNotExists(ctx context.Context) error {
	log := logger.GetLogger(ctx)
	log.Debugf("[ElasticsearchV7] Checking if index exists: %s", e.index)

	res, err := e.client.Indices.Exists([]string{e.index}, e.client.Indices.Exists.WithContext(ctx))
	if err != nil {
		return fmt.Errorf("check index existence: %w", err)
	}
	defer res.Body.Close()

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

	// Build settings body with optional shards/replicas
	var body string
	if e.numberOfShards > 0 || e.numberOfReplicas >= 0 {
		settings := make(map[string]interface{})
		if e.numberOfShards > 0 {
			settings["number_of_shards"] = e.numberOfShards
		}
		if e.numberOfReplicas >= 0 {
			settings["number_of_replicas"] = e.numberOfReplicas
		}
		bodyBytes, err := json.Marshal(map[string]interface{}{"settings": settings})
		if err != nil {
			return fmt.Errorf("marshal index settings: %w", err)
		}
		body = string(bodyBytes)
	}

	log.Infof("[ElasticsearchV7] Creating index: %s", e.index)
	var opts []func(*esapi.IndicesCreateRequest)
	opts = append(opts, e.client.Indices.Create.WithContext(ctx))
	if body != "" {
		opts = append(opts, e.client.Indices.Create.WithBody(strings.NewReader(body)))
	}

	createRes, err := e.client.Indices.Create(e.index, opts...)
	if err != nil {
		return fmt.Errorf("create index: %w", err)
	}
	defer createRes.Body.Close()

	if createRes.IsError() {
		// Log detailed response server-side; return generic message to avoid leaking cluster info
		log.Errorf("[ElasticsearchV7] Create index response: %s", createRes.String())
		return fmt.Errorf("failed to create index %s", e.index)
	}

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

func (e *elasticsearchRepository) EngineType() typesLocal.RetrieverEngineType {
	return typesLocal.ElasticsearchRetrieverEngineType
}

func (e *elasticsearchRepository) Support() []typesLocal.RetrieverType {
	return []typesLocal.RetrieverType{typesLocal.KeywordsRetrieverType}
}

// EstimateStorageSize 估算存储空间大小
func (e *elasticsearchRepository) EstimateStorageSize(ctx context.Context,
	indexInfoList []*typesLocal.IndexInfo, params map[string]any,
) int64 {
	log := logger.GetLogger(ctx)
	log.Infof("[ElasticsearchV7] Estimating storage size for %d indices", len(indexInfoList))

	// 计算总存储大小
	var totalStorageSize int64 = 0
	for _, indexInfo := range indexInfoList {
		embeddingDB := elasticsearchRetriever.ToDBVectorEmbedding(indexInfo, params)
		// 计算单个文档的存储大小并累加
		totalStorageSize += e.calculateStorageSize(embeddingDB)
	}

	// 记录存储大小
	log.Infof("[ElasticsearchV7] Estimated storage size: %d bytes (%d MB) for %d indices",
		totalStorageSize, totalStorageSize/(1024*1024), len(indexInfoList))
	return totalStorageSize
}

// 计算单个索引的存储占用大小(Bytes)
func (e *elasticsearchRepository) calculateStorageSize(embedding *elasticsearchRetriever.VectorEmbedding) int64 {
	// 1. 文本内容大小
	contentSizeBytes := int64(len(embedding.Content))

	// 2. 向量存储大小
	var vectorSizeBytes int64 = 0
	if embedding.Embedding != nil {
		// 4字节/维度 (全精度浮点数)
		vectorSizeBytes = int64(len(embedding.Embedding) * 4)
	}

	// 3. 元数据大小 (ID、时间戳等固定开销)
	metadataSizeBytes := int64(250) // 约250字节的元数据

	// 4. 索引开销 (ES索引的放大系数约为1.5)
	indexOverheadBytes := (contentSizeBytes + vectorSizeBytes) * 5 / 10

	// 总大小 (字节)
	totalSizeBytes := contentSizeBytes + vectorSizeBytes + metadataSizeBytes + indexOverheadBytes

	return totalSizeBytes
}

// Save 保存索引
func (e *elasticsearchRepository) Save(ctx context.Context,
	embedding *typesLocal.IndexInfo, additionalParams map[string]any,
) error {
	log := logger.GetLogger(ctx)
	log.Debugf("[ElasticsearchV7] Saving index for chunk ID: %s", embedding.ChunkID)

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

	docBytes, err := json.Marshal(embeddingDB)
	if err != nil {
		log.Errorf("[ElasticsearchV7] Failed to marshal embedding: %v", err)
		return err
	}

	docID := uuid.New().String()
	log.Debugf("[ElasticsearchV7] Creating document with ID: %s for chunk ID: %s", docID, embedding.ChunkID)

	resp, err := e.client.Create(
		e.index,
		docID,
		bytes.NewReader(docBytes),
		e.client.Create.WithContext(ctx),
	)
	if err != nil {
		log.Errorf("[ElasticsearchV7] Failed to create document: %v", err)
		return err
	}
	defer resp.Body.Close()

	if resp.IsError() {
		log.Errorf("[ElasticsearchV7] Failed to index document: %s", resp.String())
		return fmt.Errorf("failed to index document: %s", resp.String())
	}

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

// BatchSave performs bulk indexing of multiple embeddings in Elasticsearch
// It constructs a bulk request with index operations for each embedding
// Returns error if any operation fails during the bulk indexing process
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("[ElasticsearchV7] Empty list provided to BatchSave, skipping")
		return nil
	}

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

	// Prepare bulk request body
	body, processedCount, err := e.prepareBulkRequestBody(ctx, embeddingList, additionalParams)
	if err != nil {
		return err
	}

	if processedCount == 0 {
		log.Warn("[ElasticsearchV7] No valid documents to index after filtering, skipping bulk request")
		return nil
	}

	// Execute bulk request
	log.Debugf("[ElasticsearchV7] Executing bulk request with %d documents", processedCount)
	resp, err := e.client.Bulk(
		body,
		e.client.Bulk.WithIndex(e.index),
		e.client.Bulk.WithContext(ctx),
	)
	if err != nil {
		log.Errorf("[ElasticsearchV7] Failed to execute bulk index operation: %v", err)
		return err
	}
	defer resp.Body.Close()

	// Process bulk response
	err = e.processBulkResponse(ctx, resp, len(embeddingList))
	if err != nil {
		return err
	}

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

// prepareBulkRequestBody prepares the bulk request body for batch indexing
func (e *elasticsearchRepository) prepareBulkRequestBody(ctx context.Context,
	embeddingList []*typesLocal.IndexInfo, additionalParams map[string]any,
) (*bytes.Buffer, int, error) {
	log := logger.GetLogger(ctx)
	body := &bytes.Buffer{}
	processedCount := 0

	// Prepare bulk request body
	for _, embedding := range embeddingList {
		// Convert to Elasticsearch document format
		embeddingDB := elasticsearchRetriever.ToDBVectorEmbedding(embedding, additionalParams)

		// Generate document ID and metadata line
		docID := uuid.New().String()
		meta := []byte(fmt.Sprintf(`{ "index" : { "_id" : "%s" } }%s`, docID, "\n"))

		// Marshal document to JSON
		data, err := json.Marshal(embeddingDB)
		if err != nil {
			log.Errorf("[ElasticsearchV7] Failed to marshal embedding for chunk ID %s: %v", embedding.ChunkID, err)
			return nil, 0, err
		}

		// Append newline and add to bulk request body
		data = append(data, "\n"...)
		body.Grow(len(meta) + len(data))
		body.Write(meta)
		body.Write(data)
		processedCount++

		log.Debugf("[ElasticsearchV7] Added document for chunk ID: %s to bulk request", embedding.ChunkID)
	}

	return body, processedCount, nil
}

// processBulkResponse processes the response from a bulk indexing operation
func (e *elasticsearchRepository) processBulkResponse(ctx context.Context,
	resp *esapi.Response, totalDocuments int,
) error {
	log := logger.GetLogger(ctx)

	// Check for bulk operation errors
	if resp.IsError() {
		log.Errorf("[ElasticsearchV7] Bulk operation failed: %s", resp.String())
		return fmt.Errorf("failed to index documents: %s", resp.String())
	}

	// Parse bulk response to check for individual document errors
	var bulkResponse map[string]interface{}
	if err := json.NewDecoder(resp.Body).Decode(&bulkResponse); err != nil {
		log.Warnf("[ElasticsearchV7] Could not parse bulk response: %v", err)
		return nil
	}

	// Check for errors in individual operations
	if hasErrors, ok := bulkResponse["errors"].(bool); ok && hasErrors {
		errorCount := e.countBulkErrors(ctx, bulkResponse, totalDocuments)
		if errorCount > 0 {
			log.Warnf("[ElasticsearchV7] %d/%d documents failed to index", errorCount, totalDocuments)
		}
	}

	return nil
}

// countBulkErrors counts the number of errors in a bulk response
func (e *elasticsearchRepository) countBulkErrors(ctx context.Context,
	bulkResponse map[string]interface{}, totalDocuments int,
) int {
	log := logger.GetLogger(ctx)
	log.Warn("[ElasticsearchV7] Bulk operation completed with some errors")

	errorCount := 0
	if items, ok := bulkResponse["items"].([]interface{}); ok {
		for _, item := range items {
			if itemMap, ok := item.(map[string]interface{}); ok {
				if indexResp, ok := itemMap["index"].(map[string]interface{}); ok {
					if indexResp["error"] != nil {
						errorCount++
						log.Errorf("[ElasticsearchV7] Item error: %v", indexResp["error"])
					}
				}
			}
		}
	}

	return errorCount
}

// DeleteByChunkIDList Delete indices by chunk ID list
func (e *elasticsearchRepository) DeleteByChunkIDList(ctx context.Context, chunkIDList []string, dimension int, knowledgeType string) error {
	return e.deleteByFieldList(ctx, e.idField("chunk_id"), chunkIDList)
}

// DeleteBySourceIDList Delete indices by source ID list
func (e *elasticsearchRepository) DeleteBySourceIDList(ctx context.Context, sourceIDList []string, dimension int, knowledgeType string) error {
	return e.deleteByFieldList(ctx, e.idField("source_id"), sourceIDList)
}

// DeleteByKnowledgeIDList Delete indices by knowledge ID list
func (e *elasticsearchRepository) DeleteByKnowledgeIDList(ctx context.Context,
	knowledgeIDList []string, dimension int, knowledgeType string,
) error {
	return e.deleteByFieldList(ctx, e.idField("knowledge_id"), knowledgeIDList)
}

// deleteByFieldList Delete documents by field value list
func (e *elasticsearchRepository) deleteByFieldList(ctx context.Context, field string, valueList []string) error {
	log := logger.GetLogger(ctx)
	if len(valueList) == 0 {
		log.Warnf("[ElasticsearchV7] Empty %s list provided for deletion, skipping", field)
		return nil
	}

	log.Infof("[ElasticsearchV7] Deleting indices by %s, count: %d", field, len(valueList))
	ids, err := json.Marshal(valueList)
	if err != nil {
		log.Errorf("[ElasticsearchV7] Failed to marshal %s list: %v", field, err)
		return err
	}

	query := fmt.Sprintf(`{"query": {"terms": {"%s": %s}}}`, field, ids)
	log.Debugf("[ElasticsearchV7] Executing delete by query: %s", query)

	resp, err := e.client.DeleteByQuery(
		[]string{e.index},
		bytes.NewReader([]byte(query)),
		e.client.DeleteByQuery.WithContext(ctx),
	)
	if err != nil {
		log.Errorf("[ElasticsearchV7] Failed to execute delete by query: %v", err)
		return err
	}
	defer resp.Body.Close()

	if resp.IsError() {
		errMsg := fmt.Sprintf("failed to delete by query: %s", resp.String())
		log.Errorf("[ElasticsearchV7] %s", errMsg)
		return fmt.Errorf("failed to delete by query: %s", resp.String())
	}

	// Try to extract deletion count from response
	var deleteResponse map[string]interface{}
	if err := json.NewDecoder(resp.Body).Decode(&deleteResponse); err != nil {
		log.Warnf("[ElasticsearchV7] Could not parse delete response: %v", err)
	} else {
		if deleted, ok := deleteResponse["deleted"].(float64); ok {
			log.Infof("[ElasticsearchV7] Successfully deleted %d documents by %s", int(deleted), field)
		} else {
			log.Infof("[ElasticsearchV7] Successfully deleted documents by %s", field)
		}
	}

	return nil
}

// getBaseConds Construct base Elasticsearch query conditions based on retrieval parameters
// It creates MUST conditions for required fields and MUST_NOT conditions for excluded fields
// KnowledgeBaseIDs and KnowledgeIDs use AND logic (search specific documents within knowledge bases)
// Returns a JSON string representing the query conditions
func (e *elasticsearchRepository) getBaseConds(params typesLocal.RetrieveParams) string {
	// Build MUST conditions (positive filters)
	must := make([]map[string]interface{}, 0)

	// 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, map[string]interface{}{
			"terms": map[string]interface{}{
				e.idField("knowledge_base_id"): params.KnowledgeBaseIDs,
			},
		})
	}
	if len(params.KnowledgeIDs) > 0 {
		must = append(must, map[string]interface{}{
			"terms": map[string]interface{}{
				e.idField("knowledge_id"): params.KnowledgeIDs,
			},
		})
	}
	// Filter by tag IDs if specified
	if len(params.TagIDs) > 0 {
		must = append(must, map[string]interface{}{
			"terms": map[string]interface{}{
				e.idField("tag_id"): params.TagIDs,
			},
		})
	}

	// Build MUST_NOT conditions (negative filters)
	mustNot := make([]map[string]interface{}, 0)
	// Exclude disabled chunks (is_enabled = false)
	// Note: Historical data without is_enabled field will be included (not matching must_not)
	mustNot = append(mustNot, map[string]interface{}{
		"term": map[string]interface{}{
			"is_enabled": false,
		},
	})
	if len(params.ExcludeKnowledgeIDs) > 0 {
		mustNot = append(mustNot, map[string]interface{}{
			"terms": map[string]interface{}{
				e.idField("knowledge_id"): params.ExcludeKnowledgeIDs,
			},
		})
	}
	if len(params.ExcludeChunkIDs) > 0 {
		mustNot = append(mustNot, map[string]interface{}{
			"terms": map[string]interface{}{
				e.idField("chunk_id"): params.ExcludeChunkIDs,
			},
		})
	}

	// Combine conditions based on presence
	var query map[string]interface{}
	if len(must) == 0 && len(mustNot) == 0 {
		query = map[string]interface{}{}
	} else if len(must) == 0 {
		query = map[string]interface{}{
			"bool": map[string]interface{}{
				"must_not": mustNot,
			},
		}
	} else if len(mustNot) == 0 {
		query = map[string]interface{}{
			"bool": map[string]interface{}{
				"must": must,
			},
		}
	} else {
		query = map[string]interface{}{
			"bool": map[string]interface{}{
				"must":     must,
				"must_not": mustNot,
			},
		}
	}

	// Marshal to JSON string
	jsonBytes, err := json.Marshal(query)
	if err != nil {
		return "{}"
	}
	return string(jsonBytes)
}

func (e *elasticsearchRepository) Retrieve(ctx context.Context,
	params typesLocal.RetrieveParams,
) ([]*typesLocal.RetrieveResult, error) {
	log := logger.GetLogger(ctx)
	log.Debugf("[ElasticsearchV7] Processing retrieval request of type: %s", params.RetrieverType)

	switch params.RetrieverType {
	case typesLocal.KeywordsRetrieverType:
		return e.KeywordsRetrieve(ctx, params)
	}

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

// VectorRetrieve Implement vector similarity retrieval
func (e *elasticsearchRepository) VectorRetrieve(ctx context.Context,
	params typesLocal.RetrieveParams,
) ([]*typesLocal.RetrieveResult, error) {
	log := logger.GetLogger(ctx)
	log.Infof("[ElasticsearchV7] Vector retrieval: dim=%d, topK=%d, threshold=%.4f",
		len(params.Embedding), params.TopK, params.Threshold)

	// Build search query
	query, err := e.buildVectorSearchQuery(ctx, params)
	if err != nil {
		return nil, err
	}

	// Execute search
	results, err := e.executeVectorSearch(ctx, query)
	if err != nil {
		return nil, err
	}

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

// buildVectorSearchQuery builds the vector search query JSON
func (e *elasticsearchRepository) buildVectorSearchQuery(ctx context.Context,
	params typesLocal.RetrieveParams,
) (string, error) {
	log := logger.GetLogger(ctx)

	// Parse filter conditions
	var filterQuery map[string]interface{}
	filterJSON := e.getBaseConds(params)
	if err := json.Unmarshal([]byte(filterJSON), &filterQuery); err != nil {
		log.Errorf("[ElasticsearchV7] Failed to unmarshal filter: %v", err)
		filterQuery = map[string]interface{}{}
	}

	// Construct the script_score query using structured objects
	queryObj := map[string]interface{}{
		"query": map[string]interface{}{
			"script_score": map[string]interface{}{
				"query": map[string]interface{}{
					"bool": map[string]interface{}{
						"filter": []interface{}{filterQuery},
					},
				},
				"script": map[string]interface{}{
					"source": "cosineSimilarity(params.query_vector,'embedding')",
					"params": map[string]interface{}{
						"query_vector": params.Embedding,
					},
				},
				"min_score": params.Threshold,
			},
		},
		"size": params.TopK,
	}

	// Marshal to JSON string
	queryBytes, err := json.Marshal(queryObj)
	if err != nil {
		log.Errorf("[ElasticsearchV7] Failed to marshal query: %v", err)
		return "", fmt.Errorf("failed to marshal query: %w", err)
	}

	query := string(queryBytes)
	log.Debugf("[ElasticsearchV7] Executing vector search with query: %s", query)
	return query, nil
}

// executeVectorSearch executes the vector search query
func (e *elasticsearchRepository) executeVectorSearch(
	ctx context.Context,
	query string,
) ([]*typesLocal.IndexWithScore, error) {
	log := logger.GetLogger(ctx)

	response, err := e.client.Search(
		e.client.Search.WithIndex(e.index),
		e.client.Search.WithBody(strings.NewReader(query)),
		e.client.Search.WithContext(ctx),
	)
	if err != nil {
		log.Errorf("[ElasticsearchV7] Vector search failed: %v", err)
		return nil, err
	}
	defer response.Body.Close()

	results, err := e.processSearchResponse(ctx, response, typesLocal.VectorRetrieverType)
	if err != nil {
		return nil, err
	}

	return results, nil
}

// KeywordsRetrieve Implement keyword retrieval
func (e *elasticsearchRepository) KeywordsRetrieve(ctx context.Context,
	params typesLocal.RetrieveParams,
) ([]*typesLocal.RetrieveResult, error) {
	log := logger.GetLogger(ctx)
	log.Infof("[ElasticsearchV7] Keywords retrieval: query=%s, topK=%d", params.Query, params.TopK)

	// Build search query
	query, err := e.buildKeywordSearchQuery(ctx, params)
	if err != nil {
		return nil, err
	}

	// Execute search
	results, err := e.executeKeywordSearch(ctx, query)
	if err != nil {
		return nil, err
	}

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

// buildKeywordSearchQuery builds the keyword search query JSON
func (e *elasticsearchRepository) buildKeywordSearchQuery(ctx context.Context,
	params typesLocal.RetrieveParams,
) (string, error) {
	log := logger.GetLogger(ctx)
	content, err := json.Marshal(params.Query)
	if err != nil {
		log.Errorf("[ElasticsearchV7] Failed to marshal query: %v", err)
		return "", err
	}

	filter := e.getBaseConds(params)
	query := fmt.Sprintf(
		`{"query": {"bool": {"must": [{"match": {"content": %s}}], "filter": [%s]}}}`,
		string(content), filter,
	)

	log.Debugf("[ElasticsearchV7] Executing keyword search with query: %s", query)
	return query, nil
}

// executeKeywordSearch executes the keyword search query
func (e *elasticsearchRepository) executeKeywordSearch(
	ctx context.Context, query string,
) ([]*typesLocal.IndexWithScore, error) {
	log := logger.GetLogger(ctx)

	response, err := e.client.Search(
		e.client.Search.WithIndex(e.index),
		e.client.Search.WithBody(
			strings.NewReader(query),
		),
		e.client.Search.WithContext(ctx),
	)
	if err != nil {
		log.Errorf("[ElasticsearchV7] Keywords search failed: %v", err)
		return nil, err
	}
	defer response.Body.Close()

	results, err := e.processSearchResponse(ctx, response, typesLocal.KeywordsRetrieverType)
	if err != nil {
		return nil, err
	}

	return results, nil
}

// processSearchResponse Process search response
func (e *elasticsearchRepository) processSearchResponse(ctx context.Context,
	response *esapi.Response, retrieverType typesLocal.RetrieverType,
) ([]*typesLocal.IndexWithScore, error) {
	log := logger.GetLogger(ctx)

	if response.IsError() {
		errMsg := fmt.Sprintf("failed to retrieve: %s", response.String())
		log.Errorf("[ElasticsearchV7] %s", errMsg)
		return nil, errors.New(errMsg)
	}

	// Decode response body
	rJson, err := e.decodeSearchResponse(ctx, response)
	if err != nil {
		return nil, err
	}

	// Extract hits from response
	hitsList, err := e.extractHitsFromResponse(ctx, rJson)
	if err != nil {
		return nil, err
	}

	// Process hits into results
	results, err := e.processHits(ctx, hitsList, retrieverType)
	if err != nil {
		return nil, err
	}

	// Log results summary
	e.logResultsSummary(ctx, results, retrieverType)

	return results, nil
}

// decodeSearchResponse decodes the search response body
func (e *elasticsearchRepository) decodeSearchResponse(ctx context.Context,
	response *esapi.Response,
) (map[string]any, error) {
	log := logger.GetLogger(ctx)
	var rJson map[string]any

	if err := json.NewDecoder(response.Body).Decode(&rJson); err != nil {
		log.Errorf("[ElasticsearchV7] Failed to decode search response: %v", err)
		return nil, err
	}

	return rJson, nil
}

// extractHitsFromResponse extracts the hits list from the response JSON
func (e *elasticsearchRepository) extractHitsFromResponse(ctx context.Context,
	rJson map[string]any,
) ([]interface{}, error) {
	log := logger.GetLogger(ctx)

	// Extract hits from response
	hitsObj, ok := rJson["hits"].(map[string]interface{})
	if !ok {
		log.Errorf("[ElasticsearchV7] Invalid search response format: 'hits' object missing")
		return nil, fmt.Errorf("invalid search response format")
	}

	hitsList, ok := hitsObj["hits"].([]interface{})
	if !ok {
		log.Warnf("[ElasticsearchV7] No hits found in search response")
		return []interface{}{}, nil
	}

	return hitsList, nil
}

// processHits processes the hits into IndexWithScore results
func (e *elasticsearchRepository) processHits(ctx context.Context,
	hitsList []interface{}, retrieverType typesLocal.RetrieverType,
) ([]*typesLocal.IndexWithScore, error) {
	log := logger.GetLogger(ctx)
	results := make([]*typesLocal.IndexWithScore, 0, len(hitsList))

	for _, hit := range hitsList {
		indexWithScore, err := e.processHit(ctx, hit, retrieverType)
		if err != nil {
			// Log error but continue processing other hits
			log.Warnf("[ElasticsearchV7] Error processing hit: %v", err)
			continue
		}

		if indexWithScore != nil {
			results = append(results, indexWithScore)
		}
	}

	return results, nil
}

// processHit processes a single hit into an IndexWithScore
func (e *elasticsearchRepository) processHit(ctx context.Context,
	hit interface{}, retrieverType typesLocal.RetrieverType,
) (*typesLocal.IndexWithScore, error) {
	log := logger.GetLogger(ctx)

	hitMap, ok := hit.(map[string]interface{})
	if !ok {
		return nil, fmt.Errorf("invalid hit object format")
	}

	// Get document ID
	docID, ok := hitMap["_id"].(string)
	if !ok {
		return nil, fmt.Errorf("hit missing document ID")
	}

	// Get document source
	sourceObj, ok := hitMap["_source"]
	if !ok {
		return nil, fmt.Errorf("hit %s missing _source", docID)
	}

	// Get score
	score, ok := hitMap["_score"].(float64)
	if !ok {
		return nil, fmt.Errorf("hit %s missing score", docID)
	}

	// Convert source to embedding
	embedding, err := e.convertSourceToEmbedding(ctx, sourceObj, docID, score)
	if err != nil {
		return nil, err
	}

	result := elasticsearchRetriever.FromDBVectorEmbeddingWithScore(
		docID, embedding, typesLocal.MatchTypeKeywords,
	)

	matchType := "keyword"
	if retrieverType == typesLocal.VectorRetrieverType {
		matchType = "vector"
	}
	log.Debugf("[ElasticsearchV7] %s search result: id=%s, score=%.4f", matchType, docID, score)

	return result, nil
}

// convertSourceToEmbedding converts the source object to an embedding
func (e *elasticsearchRepository) convertSourceToEmbedding(ctx context.Context,
	sourceObj interface{}, docID string, score float64,
) (*elasticsearchRetriever.VectorEmbeddingWithScore, error) {
	log := logger.GetLogger(ctx)

	// Convert source to embedding
	var embedding *elasticsearchRetriever.VectorEmbeddingWithScore
	sourceBytes, err := json.Marshal(sourceObj)
	if err != nil {
		log.Warnf("[ElasticsearchV7] Failed to marshal source for hit %s: %v", docID, err)
		return nil, fmt.Errorf("failed to marshal source for hit %s: %v", docID, err)
	}

	if err := json.Unmarshal(sourceBytes, &embedding); err != nil {
		log.Warnf("[ElasticsearchV7] Failed to unmarshal source for hit %s: %v", docID, err)
		return nil, fmt.Errorf("failed to unmarshal source for hit %s: %v", docID, err)
	}

	embedding.Score = score
	return embedding, nil
}

// logResultsSummary logs a summary of the results
func (e *elasticsearchRepository) logResultsSummary(ctx context.Context,
	results []*typesLocal.IndexWithScore, retrieverType typesLocal.RetrieverType,
) {
	log := logger.GetLogger(ctx)

	if len(results) == 0 {
		if retrieverType == typesLocal.KeywordsRetrieverType {
			log.Warnf("[ElasticsearchV7] No keyword matches found")
		} else {
			log.Warnf("[ElasticsearchV7] No vector matches found that meet threshold")
		}
	} else {
		retrievalType := "Keywords"
		if retrieverType == typesLocal.VectorRetrieverType {
			retrievalType = "Vector"
		}
		log.Infof("[ElasticsearchV7] %s retrieval found %d results", retrievalType, len(results))
		log.Debugf("[ElasticsearchV7] Top result score: %.4f", results[0].Score)
	}
}

// CopyIndices Copy index data
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(
		"[ElasticsearchV7] Copying indices from source knowledge base %s to target knowledge base %s, count: %d",
		sourceKnowledgeBaseID, targetKnowledgeBaseID, len(sourceToTargetChunkIDMap),
	)

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

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

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

	for {
		// Query source data batch
		hitsList, err := e.querySourceBatch(ctx, retrieveParams, from, batchSize)
		if err != nil {
			return err
		}

		// If no more data, break the loop
		if len(hitsList) == 0 {
			break
		}

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

		// Process the batch and create index information
		indexInfoList, err := e.processSourceBatch(ctx, hitsList, sourceToTargetKBIDMap,
			sourceToTargetChunkIDMap, targetKnowledgeBaseID)
		if err != nil {
			return err
		}

		// Save processed indices
		if len(indexInfoList) > 0 {
			err := e.saveCopiedIndices(ctx, indexInfoList)
			if err != nil {
				return err
			}

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

		// Move to next batch
		from += len(hitsList)

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

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

// querySourceBatch queries a batch of source data
func (e *elasticsearchRepository) querySourceBatch(ctx context.Context,
	retrieveParams typesLocal.RetrieveParams, from int, batchSize int,
) ([]interface{}, error) {
	log := logger.GetLogger(ctx)

	// Build query request safely
	filterJSON := e.getBaseConds(retrieveParams)
	var filter map[string]interface{}
	if err := json.Unmarshal([]byte(filterJSON), &filter); err != nil {
		log.Errorf("[ElasticsearchV7] Failed to parse base conditions: %v", err)
		filter = map[string]interface{}{}
	}

	queryBody := map[string]interface{}{
		"query": filter,
		"from":  from,
		"size":  batchSize,
	}

	queryBytes, err := json.Marshal(queryBody)
	if err != nil {
		log.Errorf("[ElasticsearchV7] Failed to marshal query body: %v", err)
		return nil, err
	}

	// Execute query
	response, err := e.client.Search(
		e.client.Search.WithIndex(e.index),
		e.client.Search.WithBody(strings.NewReader(string(queryBytes))),
		e.client.Search.WithContext(ctx),
	)
	if err != nil {
		log.Errorf("[ElasticsearchV7] Failed to query source index data: %v", err)
		return nil, err
	}
	defer response.Body.Close()

	if response.IsError() {
		log.Errorf("[ElasticsearchV7] Failed to query source index data: %s", response.String())
		return nil, fmt.Errorf("failed to query source index data: %s", response.String())
	}

	// 解析搜索结果
	var searchResult map[string]interface{}
	if err := json.NewDecoder(response.Body).Decode(&searchResult); err != nil {
		log.Errorf("[ElasticsearchV7] Failed to parse query result: %v", err)
		return nil, err
	}

	// 提取结果列表
	hitsObj, ok := searchResult["hits"].(map[string]interface{})
	if !ok {
		log.Errorf("[ElasticsearchV7] Invalid search result format: 'hits' object missing")
		return nil, fmt.Errorf("invalid search result format")
	}

	hitsList, ok := hitsObj["hits"].([]interface{})
	if !ok || len(hitsList) == 0 {
		if from == 0 {
			log.Warnf("[ElasticsearchV7] No source index data found")
		}
		return []interface{}{}, nil
	}

	return hitsList, nil
}

// processSourceBatch processes a batch of source data and creates index information
func (e *elasticsearchRepository) processSourceBatch(ctx context.Context,
	hitsList []interface{},
	sourceToTargetKBIDMap map[string]string,
	sourceToTargetChunkIDMap map[string]string,
	targetKnowledgeBaseID string,
) ([]*typesLocal.IndexInfo, error) {
	log := logger.GetLogger(ctx)

	// Prepare index information for batch save
	indexInfoList := make([]*typesLocal.IndexInfo, 0, len(hitsList))
	embeddingMap := make(map[string][]float32)

	// Process each hit result
	for _, hit := range hitsList {
		indexInfo, embeddingVector, err := e.processSingleHit(ctx, hit,
			sourceToTargetKBIDMap, sourceToTargetChunkIDMap, targetKnowledgeBaseID)
		if err != nil {
			log.Warnf("[ElasticsearchV7] Error processing hit: %v", err)
			continue
		}

		if indexInfo != nil {
			indexInfoList = append(indexInfoList, indexInfo)
			if embeddingVector != nil {
				embeddingMap[indexInfo.ChunkID] = embeddingVector
			}
		}
	}

	return indexInfoList, nil
}

// processSingleHit processes a single hit and creates index information
func (e *elasticsearchRepository) processSingleHit(ctx context.Context,
	hit interface{},
	sourceToTargetKBIDMap map[string]string,
	sourceToTargetChunkIDMap map[string]string,
	targetKnowledgeBaseID string,
) (*typesLocal.IndexInfo, []float32, error) {
	log := logger.GetLogger(ctx)

	hitMap, ok := hit.(map[string]interface{})
	if !ok {
		log.Warnf("[ElasticsearchV7] Invalid hit object format")
		return nil, nil, fmt.Errorf("invalid hit object format")
	}

	// Get document source
	sourceObj, ok := hitMap["_source"].(map[string]interface{})
	if !ok {
		log.Warnf("[ElasticsearchV7] Hit missing _source field")
		return nil, nil, fmt.Errorf("hit missing _source field")
	}

	// Get source ChunkID and corresponding target ChunkID
	sourceChunkID, ok := sourceObj["chunk_id"].(string)
	if !ok {
		log.Warnf("[ElasticsearchV7] Source index data missing chunk_id field")
		return nil, nil, fmt.Errorf("source index data missing chunk_id field")
	}

	targetChunkID, ok := sourceToTargetChunkIDMap[sourceChunkID]
	if !ok {
		log.Warnf("[ElasticsearchV7] Source chunk ID %s not found in mapping", sourceChunkID)
		return nil, nil, fmt.Errorf("source chunk ID %s not found in mapping", sourceChunkID)
	}

	// Get mapped target knowledge ID
	sourceKnowledgeID, ok := sourceObj["knowledge_id"].(string)
	if !ok {
		log.Warnf("[ElasticsearchV7] Source index data missing knowledge_id field")
		return nil, nil, fmt.Errorf("source index data missing knowledge_id field")
	}

	targetKnowledgeID, ok := sourceToTargetKBIDMap[sourceKnowledgeID]
	if !ok {
		log.Warnf("[ElasticsearchV7] Source knowledge ID %s not found in mapping", sourceKnowledgeID)
		return nil, nil, fmt.Errorf("source knowledge ID %s not found in mapping", sourceKnowledgeID)
	}

	// Extract basic content
	content, _ := sourceObj["content"].(string)
	originalSourceID, _ := sourceObj["source_id"].(string)
	sourceType := 0
	if st, ok := sourceObj["source_type"].(float64); ok {
		sourceType = int(st)
	}

	// Extract is_enabled (default true for backward compatibility)
	isEnabled := true
	if v, ok := sourceObj["is_enabled"].(bool); ok {
		isEnabled = v
	}

	// Extract is_recommended
	isRecommended := false
	if v, ok := sourceObj["is_recommended"].(bool); ok {
		isRecommended = v
	}

	// Extract tag_id
	tagID, _ := sourceObj["tag_id"].(string)

	// 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()
	}

	// Extract embedding vector (if exists)
	var embedding []float32
	if embeddingInterface, ok := sourceObj["embedding"].([]interface{}); ok {
		embedding = make([]float32, len(embeddingInterface))
		for i, v := range embeddingInterface {
			if f, ok := v.(float64); ok {
				embedding[i] = float32(f)
			}
		}
		log.Debugf("[ElasticsearchV7] Extracted embedding vector with %d dimensions for chunk %s",
			len(embedding), targetChunkID)
	}

	// Create IndexInfo object
	indexInfo := &typesLocal.IndexInfo{
		ChunkID:         targetChunkID,
		SourceID:        targetSourceID,
		KnowledgeID:     targetKnowledgeID,
		KnowledgeBaseID: targetKnowledgeBaseID,
		Content:         content,
		SourceType:      typesLocal.SourceType(sourceType),
		IsEnabled:       isEnabled,
		IsRecommended:   isRecommended,
		TagID:           tagID,
	}

	return indexInfo, embedding, nil
}

// saveCopiedIndices saves the copied indices
func (e *elasticsearchRepository) saveCopiedIndices(ctx context.Context, indexInfoList []*typesLocal.IndexInfo) error {
	log := logger.GetLogger(ctx)

	if len(indexInfoList) == 0 {
		log.Info("[ElasticsearchV7] No indices to save, skipping")
		return nil
	}

	// Prepare additional params with embedding map
	additionalParams := make(map[string]any)
	embeddingMap := make(map[string][]float32)

	// No need to extract embeddings from metadata as they're not stored there
	// We'll use the embeddings directly from the embedding map created in processSourceBatch

	if len(embeddingMap) > 0 {
		additionalParams["embedding"] = embeddingMap
		log.Infof("[ElasticsearchV7] Found %d embeddings to save", len(embeddingMap))
	}

	// Perform batch save
	err := e.BatchSave(ctx, indexInfoList, additionalParams)
	if err != nil {
		log.Errorf("[ElasticsearchV7] Failed to batch save copied indices: %v", err)
		return err
	}

	log.Infof("[ElasticsearchV7] Successfully saved %d indices", len(indexInfoList))
	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("[ElasticsearchV7] Chunk status map is empty, skipping update")
		return nil
	}

	log.Infof("[ElasticsearchV7] 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 := map[string]interface{}{
			"query": map[string]interface{}{
				"terms": map[string]interface{}{
					e.idField("chunk_id"): enabledChunkIDs,
				},
			},
			"script": map[string]interface{}{
				"source": "ctx._source.is_enabled = true",
				"lang":   "painless",
			},
		}
		queryJSON, _ := json.Marshal(query)
		res, err := esapi.UpdateByQueryRequest{
			Index: []string{e.index},
			Body:  strings.NewReader(string(queryJSON)),
		}.Do(ctx, e.client)
		if err != nil {
			log.Errorf("[ElasticsearchV7] Failed to update enabled chunks: %v", err)
			return err
		}
		defer res.Body.Close()
		if res.IsError() {
			var e map[string]interface{}
			if err := json.NewDecoder(res.Body).Decode(&e); err != nil {
				log.Errorf("[ElasticsearchV7] Error parsing the response body: %v", err)
			} else {
				log.Errorf("[ElasticsearchV7] Error updating enabled chunks: %v", e["error"])
			}
			return fmt.Errorf("elasticsearch update_by_query failed with status: %d", res.StatusCode)
		}
		log.Infof("[ElasticsearchV7] Updated %d chunks to enabled", len(enabledChunkIDs))
	}

	// Batch update disabled chunks using update_by_query
	if len(disabledChunkIDs) > 0 {
		query := map[string]interface{}{
			"query": map[string]interface{}{
				"terms": map[string]interface{}{
					e.idField("chunk_id"): disabledChunkIDs,
				},
			},
			"script": map[string]interface{}{
				"source": "ctx._source.is_enabled = false",
				"lang":   "painless",
			},
		}
		queryJSON, _ := json.Marshal(query)
		res, err := esapi.UpdateByQueryRequest{
			Index: []string{e.index},
			Body:  strings.NewReader(string(queryJSON)),
		}.Do(ctx, e.client)
		if err != nil {
			log.Errorf("[ElasticsearchV7] Failed to update disabled chunks: %v", err)
			return err
		}
		defer res.Body.Close()
		if res.IsError() {
			var e map[string]interface{}
			if err := json.NewDecoder(res.Body).Decode(&e); err != nil {
				log.Errorf("[ElasticsearchV7] Error parsing the response body: %v", err)
			} else {
				log.Errorf("[ElasticsearchV7] Error updating disabled chunks: %v", e["error"])
			}
			return fmt.Errorf("elasticsearch update_by_query failed with status: %d", res.StatusCode)
		}
		log.Infof("[ElasticsearchV7] Updated %d chunks to disabled", len(disabledChunkIDs))
	}

	log.Infof("[ElasticsearchV7] 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("[ElasticsearchV7] Chunk tag map is empty, skipping update")
		return nil
	}

	log.Infof("[ElasticsearchV7] 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 := map[string]interface{}{
			"query": map[string]interface{}{
				"terms": map[string]interface{}{
					e.idField("chunk_id"): chunkIDs,
				},
			},
			"script": map[string]interface{}{
				"source": "ctx._source.tag_id = params.tag_id",
				"lang":   "painless",
				"params": map[string]interface{}{
					"tag_id": tagID,
				},
			},
		}
		queryJSON, _ := json.Marshal(query)
		res, err := esapi.UpdateByQueryRequest{
			Index: []string{e.index},
			Body:  strings.NewReader(string(queryJSON)),
		}.Do(ctx, e.client)
		if err != nil {
			log.Errorf("[ElasticsearchV7] Failed to update chunks with tag_id %s: %v", tagID, err)
			return err
		}
		defer res.Body.Close()
		if res.IsError() {
			var e map[string]interface{}
			if err := json.NewDecoder(res.Body).Decode(&e); err != nil {
				log.Errorf("[ElasticsearchV7] Error parsing the response body: %v", err)
			} else {
				log.Errorf("[ElasticsearchV7] Error updating chunks with tag_id: %v", e["error"])
			}
			return fmt.Errorf("elasticsearch update_by_query failed with status: %d", res.StatusCode)
		}
		log.Infof("[ElasticsearchV7] Updated %d chunks to tag_id=%s", len(chunkIDs), tagID)
	}

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