ragflow/internal/dao/kb.go

//
//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//

package dao

import (
	"errors"
	"fmt"
	"path"
	"ragflow/internal/entity"

	"strconv"
	"strings"

	"gorm.io/gorm"
)

// GetTenantIDByKBID is a convenience function that retrieves the tenant ID
// for a given knowledge base ID. It is a package-level helper so both the
// service and engine layers can use it without circular imports.
func GetTenantIDByKBID(kbID string) (string, error) {
	kbDAO := NewKnowledgebaseDAO()
	kb, err := kbDAO.GetByID(kbID)
	if err != nil {
		return "", fmt.Errorf("knowledgebase not found: %w", err)
	}
	return kb.TenantID, nil
}

// KnowledgebaseDAO knowledge base data access object
type KnowledgebaseDAO struct{}

// IsNotFoundErr returns true if the error indicates a record not found
func IsNotFoundErr(err error) bool {
	return errors.Is(err, gorm.ErrRecordNotFound)
}

// NewKnowledgebaseDAO create knowledge base DAO
func NewKnowledgebaseDAO() *KnowledgebaseDAO {
	return &KnowledgebaseDAO{}
}

// Create creates a new knowledge base record
func (dao *KnowledgebaseDAO) Create(kb *entity.Knowledgebase) error {
	return DB.Create(kb).Error
}

// Update updates a knowledge base record
func (dao *KnowledgebaseDAO) Update(kb *entity.Knowledgebase) error {
	return DB.Save(kb).Error
}

// UpdateByID updates a knowledge base by ID with the given fields
func (dao *KnowledgebaseDAO) UpdateByID(id string, updates map[string]interface{}) error {
	return DB.Model(&entity.Knowledgebase{}).Where("id = ?", id).Updates(updates).Error
}

// Delete soft deletes a knowledge base by setting status to invalid
func (dao *KnowledgebaseDAO) Delete(id string) error {
	return DB.Model(&entity.Knowledgebase{}).Where("id = ?", id).Update("status", string(entity.StatusInvalid)).Error
}

// GetByID retrieves a knowledge base by ID
func (dao *KnowledgebaseDAO) GetByID(id string) (*entity.Knowledgebase, error) {
	var kb entity.Knowledgebase
	err := DB.Where("id = ? AND status = ?", id, string(entity.StatusValid)).First(&kb).Error
	if err != nil {
		return nil, err
	}
	return &kb, nil
}

// GetByIDAndTenantID retrieves a knowledge base by ID and tenant ID
func (dao *KnowledgebaseDAO) GetByIDAndTenantID(id, tenantID string) (*entity.Knowledgebase, error) {
	var kb entity.Knowledgebase
	err := DB.Where("id = ? AND tenant_id = ? AND status = ?", id, tenantID, string(entity.StatusValid)).First(&kb).Error
	if err != nil {
		return nil, err
	}
	return &kb, nil
}

// GetByIDs retrieves multiple knowledge bases by IDs
func (dao *KnowledgebaseDAO) GetByIDs(ids []string) ([]*entity.Knowledgebase, error) {
	var kbs []*entity.Knowledgebase
	err := DB.Where("id IN ? AND status = ?", ids, string(entity.StatusValid)).Find(&kbs).Error
	return kbs, err
}

// GetByName retrieves a knowledge base by name and tenant ID
func (dao *KnowledgebaseDAO) GetByName(name, tenantID string) (*entity.Knowledgebase, error) {
	var kb entity.Knowledgebase
	err := DB.Where("name = ? AND tenant_id = ? AND status = ?", name, tenantID, string(entity.StatusValid)).First(&kb).Error
	if err != nil {
		return nil, err
	}
	return &kb, nil
}

// GetByCreatedBy retrieves knowledge bases created by a specific user
func (dao *KnowledgebaseDAO) GetByCreatedBy(createdBy string) ([]*entity.Knowledgebase, error) {
	var kbs []*entity.Knowledgebase
	err := DB.Where("created_by = ? AND status = ?", createdBy, string(entity.StatusValid)).Find(&kbs).Error
	return kbs, err
}

// Query retrieves knowledge bases with filters
func (dao *KnowledgebaseDAO) Query(filters map[string]interface{}) ([]*entity.Knowledgebase, error) {
	var kbs []*entity.Knowledgebase
	query := DB.Where("status = ?", string(entity.StatusValid))

	for key, value := range filters {
		if value != nil && value != "" {
			query = query.Where(key+" = ?", value)
		}
	}

	err := query.Find(&kbs).Error
	return kbs, err
}

// QueryOne retrieves a single knowledge base with filters
func (dao *KnowledgebaseDAO) QueryOne(filters map[string]interface{}) (*entity.Knowledgebase, error) {
	var kb entity.Knowledgebase
	query := DB.Where("status = ?", string(entity.StatusValid))

	for key, value := range filters {
		if value != nil && value != "" {
			query = query.Where(key+" = ?", value)
		}
	}

	err := query.First(&kb).Error
	if err != nil {
		return nil, err
	}
	return &kb, nil
}

// Count returns the count of knowledge bases matching the filters
func (dao *KnowledgebaseDAO) Count(filters map[string]interface{}) (int64, error) {
	var count int64
	query := DB.Model(&entity.Knowledgebase{}).Where("status = ?", string(entity.StatusValid))

	for key, value := range filters {
		if value != nil && value != "" {
			query = query.Where(key+" = ?", value)
		}
	}

	err := query.Count(&count).Error
	return count, err
}

// GetByTenantIDs retrieves knowledge bases by tenant IDs with pagination
// This matches the Python get_by_tenant_ids method
func (dao *KnowledgebaseDAO) GetByTenantIDs(tenantIDs []string, userID string, pageNumber, itemsPerPage int, orderby string, desc bool, keywords, parserID string) ([]*entity.KnowledgebaseListItem, int64, error) {
	var kbs []*entity.KnowledgebaseListItem
	var total int64

	query := DB.Model(&entity.Knowledgebase{}).
		Select(`knowledgebase.id, knowledgebase.avatar, knowledgebase.name,
			knowledgebase.language, knowledgebase.description, knowledgebase.tenant_id,
			knowledgebase.permission, knowledgebase.doc_num, knowledgebase.token_num,
			knowledgebase.chunk_num, knowledgebase.parser_id, knowledgebase.embd_id,
			user.nickname, user.avatar as tenant_avatar, knowledgebase.update_time`).
		Joins("LEFT JOIN user ON knowledgebase.tenant_id = user.id").
		Where("((knowledgebase.tenant_id IN ? AND knowledgebase.permission = ?) OR knowledgebase.tenant_id = ?) AND knowledgebase.status = ?",
			tenantIDs, string(entity.TenantPermissionTeam), userID, string(entity.StatusValid))

	if keywords != "" {
		query = query.Where("LOWER(knowledgebase.name) LIKE ?", "%"+strings.ToLower(keywords)+"%")
	}

	if parserID != "" {
		query = query.Where("knowledgebase.parser_id = ?", parserID)
	}

	if desc {
		query = query.Order("knowledgebase." + orderby + " DESC")
	} else {
		query = query.Order("knowledgebase." + orderby + " ASC")
	}

	if err := query.Count(&total).Error; err != nil {
		return nil, 0, err
	}

	if pageNumber > 0 && itemsPerPage > 0 {
		offset := (pageNumber - 1) * itemsPerPage
		if err := query.Offset(offset).Limit(itemsPerPage).Scan(&kbs).Error; err != nil {
			return nil, 0, err
		}
	} else {
		if err := query.Scan(&kbs).Error; err != nil {
			return nil, 0, err
		}
	}

	return kbs, total, nil
}

// GetAllByTenantIDs retrieves all permitted knowledge bases by tenant IDs
// This matches the Python get_all_kb_by_tenant_ids method
func (dao *KnowledgebaseDAO) GetAllByTenantIDs(tenantIDs []string, userID string) ([]*entity.Knowledgebase, error) {
	var kbs []*entity.Knowledgebase

	err := DB.Where(
		"(tenant_id IN ? AND permission = ?) OR tenant_id = ?",
		tenantIDs, string(entity.TenantPermissionTeam), userID,
	).Order("create_time ASC").Find(&kbs).Error

	return kbs, err
}

// GetDetail retrieves detailed knowledge base information with joined pipeline data
// This matches the Python get_detail method
func (dao *KnowledgebaseDAO) GetDetail(kbID string) (*entity.KnowledgebaseDetail, error) {
	var detail entity.KnowledgebaseDetail

	err := DB.Table("knowledgebase").
		Select(`knowledgebase.id, knowledgebase.embd_id, knowledgebase.avatar, knowledgebase.name,
			knowledgebase.language, knowledgebase.description, knowledgebase.permission,
			knowledgebase.doc_num, knowledgebase.token_num, knowledgebase.chunk_num,
			knowledgebase.parser_id, knowledgebase.pipeline_id,
			user_canvas.title as pipeline_name, user_canvas.avatar as pipeline_avatar,
			knowledgebase.parser_config, knowledgebase.pagerank,
			knowledgebase.graphrag_task_id, knowledgebase.graphrag_task_finish_at,
			knowledgebase.raptor_task_id, knowledgebase.raptor_task_finish_at,
			knowledgebase.mindmap_task_id, knowledgebase.mindmap_task_finish_at,
			knowledgebase.create_time, knowledgebase.update_time`).
		Joins("LEFT JOIN user_canvas ON knowledgebase.pipeline_id = user_canvas.id").
		Where("knowledgebase.id = ? AND knowledgebase.status = ?", kbID, string(entity.StatusValid)).
		Scan(&detail).Error

	if err != nil {
		return nil, err
	}

	return &detail, nil
}

// Accessible checks if a knowledge base is accessible by a user.
// This matches the Python accessible method:
// 1. KB must exist and be VALID
// 2. If user is the owner tenant, return true
// 3. If permission is "me", only owner tenant can access
// 4. If permission is "team", user must be a member of the tenant
func (dao *KnowledgebaseDAO) Accessible(kbID, userID string) bool {
	var kb entity.Knowledgebase
	err := DB.Where("id = ? AND status = ?", kbID, string(entity.StatusValid)).First(&kb).Error
	if err != nil {
		return false
	}

	// User is the owner tenant itself
	if kb.TenantID == userID {
		return true
	}

	// If permission is "me", only the owner can access
	if kb.Permission == string(entity.TenantPermissionMe) {
		return false
	}

	var count int64
	err = DB.Table("user_tenant").
		Where("tenant_id = ? AND user_id = ?", kb.TenantID, userID).
		Count(&count).Error

	if err != nil {
		return false
	}
	return count > 0
}

// Accessible4Deletion checks if a knowledge base can be deleted by a user
// This matches the Python accessible4deletion method
func (dao *KnowledgebaseDAO) Accessible4Deletion(kbID, userID string) bool {
	var count int64
	err := DB.Model(&entity.Knowledgebase{}).
		Where("id = ? AND created_by = ? AND status = ?", kbID, userID, string(entity.StatusValid)).
		Count(&count).Error

	if err != nil {
		return false
	}
	return count > 0
}

// DuplicateName generates a unique name by appending parentheses if name already exists
// This matches the Python duplicate_name function behavior
func (dao *KnowledgebaseDAO) DuplicateName(name, tenantID string) string {
	const maxRetries = 1000

	currentName := name
	for retries := 0; retries < maxRetries; retries++ {
		var count int64
		err := DB.Model(&entity.Knowledgebase{}).
			Where("LOWER(name) = ? AND tenant_id = ? AND status = ?", strings.ToLower(currentName), tenantID, string(entity.StatusValid)).
			Count(&count).Error
		if err != nil || count == 0 {
			return currentName
		}

		suffix := path.Ext(currentName)
		stem := strings.TrimSuffix(currentName, suffix)
		mainPart, counter := splitNameCounter(stem)
		nextCounter := 1
		if counter > 0 {
			nextCounter = counter + 1
		}

		currentName = mainPart + "(" + strconv.Itoa(nextCounter) + ")" + suffix
	}

	return currentName
}

func splitNameCounter(name string) (string, int) {
	if !strings.HasSuffix(name, ")") {
		return name, 0
	}

	leftBracketIndex := strings.LastIndex(name, "(")
	if leftBracketIndex < 0 || leftBracketIndex >= len(name)-1 {
		return name, 0
	}

	counterValue := name[leftBracketIndex+1 : len(name)-1]
	counter, err := strconv.Atoi(counterValue)
	if err != nil {
		return name, 0
	}

	return strings.TrimRight(name[:leftBracketIndex], " "), counter
}

// AtomicIncreaseDocNumByID atomically increments the document count
// This matches the Python atomic_increase_doc_num_by_id method
func (dao *KnowledgebaseDAO) AtomicIncreaseDocNumByID(kbID string) error {
	return DB.Model(&entity.Knowledgebase{}).
		Where("id = ?", kbID).
		Updates(map[string]interface{}{
			"doc_num": DB.Raw("doc_num + 1"),
		}).Error
}

// DecreaseDocumentNum decreases document, chunk, and token counts
// This matches the Python decrease_document_num_in_delete method
func (dao *KnowledgebaseDAO) DecreaseDocumentNum(kbID string, docNum, chunkNum, tokenNum int64) error {
	return DB.Model(&entity.Knowledgebase{}).
		Where("id = ?", kbID).
		Updates(map[string]interface{}{
			"doc_num":   DB.Raw("doc_num - ?", docNum),
			"chunk_num": DB.Raw("chunk_num - ?", chunkNum),
			"token_num": DB.Raw("token_num - ?", tokenNum),
		}).Error
}

// GetKBIDsByTenantID retrieves all knowledge base IDs for a tenant
// This matches the Python get_kb_ids method
func (dao *KnowledgebaseDAO) GetKBIDsByTenantID(tenantID string) ([]string, error) {
	var kbIDs []string
	err := DB.Model(&entity.Knowledgebase{}).
		Where("tenant_id = ? AND status = ?", tenantID, string(entity.StatusValid)).
		Pluck("id", &kbIDs).Error
	return kbIDs, err
}

// GetAllIDs retrieves all knowledge base IDs
// This matches the Python get_all_ids method
func (dao *KnowledgebaseDAO) GetAllIDs() ([]string, error) {
	var kbIDs []string
	err := DB.Model(&entity.Knowledgebase{}).
		Where("status = ?", string(entity.StatusValid)).
		Pluck("id", &kbIDs).Error
	return kbIDs, err
}

// UpdateParserConfig updates the parser configuration with deep merge
// This matches the Python update_parser_config method
func (dao *KnowledgebaseDAO) UpdateParserConfig(id string, config map[string]interface{}) error {
	var kb entity.Knowledgebase
	if err := DB.Where("id = ? AND status = ?", id, string(entity.StatusValid)).First(&kb).Error; err != nil {
		return err
	}

	mergedConfig := mergeConfig(kb.ParserConfig, config)
	return DB.Model(&entity.Knowledgebase{}).
		Where("id = ?", id).
		Update("parser_config", mergedConfig).Error
}

// DeleteFieldMap removes the field_map from parser_config
// This matches the Python delete_field_map method
func (dao *KnowledgebaseDAO) DeleteFieldMap(id string) error {
	var kb entity.Knowledgebase
	if err := DB.Where("id = ? AND status = ?", id, string(entity.StatusValid)).First(&kb).Error; err != nil {
		return err
	}

	if kb.ParserConfig != nil {
		delete(kb.ParserConfig, "field_map")
		return DB.Model(&entity.Knowledgebase{}).
			Where("id = ?", id).
			Update("parser_config", kb.ParserConfig).Error
	}
	return nil
}

// GetFieldMap retrieves field mappings from multiple knowledge bases
// This matches the Python get_field_map method
func (dao *KnowledgebaseDAO) GetFieldMap(ids []string) (map[string]interface{}, error) {
	conf := make(map[string]interface{})
	kbs, err := dao.GetByIDs(ids)
	if err != nil {
		return nil, err
	}

	for _, kb := range kbs {
		if kb.ParserConfig != nil {
			if fieldMap, ok := kb.ParserConfig["field_map"]; ok {
				if fm, ok := fieldMap.(map[string]interface{}); ok {
					for k, v := range fm {
						conf[k] = v
					}
				}
			}
		}
	}
	return conf, nil
}

// GetKBByIDAndUserID retrieves a knowledge base by ID and user ID with tenant join
// This matches the Python get_kb_by_id method
func (dao *KnowledgebaseDAO) GetKBByIDAndUserID(kbID, userID string) ([]*entity.Knowledgebase, error) {
	var kbs []*entity.Knowledgebase
	err := DB.Model(&entity.Knowledgebase{}).
		Joins("JOIN user_tenant ON user_tenant.tenant_id = knowledgebase.tenant_id").
		Where("knowledgebase.id = ? AND user_tenant.user_id = ?", kbID, userID).
		Limit(1).
		Find(&kbs).Error
	return kbs, err
}

// GetKBByNameAndUserID retrieves a knowledge base by name and user ID with tenant join
// This matches the Python get_kb_by_name method
func (dao *KnowledgebaseDAO) GetKBByNameAndUserID(kbName, userID string) ([]*entity.Knowledgebase, error) {
	var kbs []*entity.Knowledgebase
	err := DB.Model(&entity.Knowledgebase{}).
		Joins("JOIN user_tenant ON user_tenant.tenant_id = knowledgebase.tenant_id").
		Where("knowledgebase.name = ? AND user_tenant.user_id = ?", kbName, userID).
		Limit(1).
		Find(&kbs).Error
	return kbs, err
}

// GetList retrieves knowledge bases with filtering by ID and name
// This matches the Python get_list method
func (dao *KnowledgebaseDAO) GetList(tenantIDs []string, userID string, pageNumber, itemsPerPage int, orderby string, desc bool, id, name string) ([]*entity.Knowledgebase, int64, error) {
	var kbs []*entity.Knowledgebase
	var total int64

	query := DB.Model(&entity.Knowledgebase{}).
		Where("((tenant_id IN ? AND permission = ?) OR tenant_id = ?) AND status = ?",
			tenantIDs, string(entity.TenantPermissionTeam), userID, string(entity.StatusValid))

	if id != "" {
		query = query.Where("id = ?", id)
	}
	if name != "" {
		query = query.Where("name = ?", name)
	}

	if desc {
		query = query.Order(orderby + " DESC")
	} else {
		query = query.Order(orderby + " ASC")
	}

	if err := query.Count(&total).Error; err != nil {
		return nil, 0, err
	}

	if pageNumber > 0 && itemsPerPage > 0 {
		offset := (pageNumber - 1) * itemsPerPage
		if err := query.Offset(offset).Limit(itemsPerPage).Find(&kbs).Error; err != nil {
			return nil, 0, err
		}
	} else {
		if err := query.Find(&kbs).Error; err != nil {
			return nil, 0, err
		}
	}

	return kbs, total, nil
}

// mergeConfig performs a deep merge of configuration maps
func mergeConfig(old, new map[string]interface{}) map[string]interface{} {
	result := make(map[string]interface{})
	for k, v := range old {
		result[k] = v
	}

	for k, v := range new {
		if existing, ok := result[k]; ok {
			if existingMap, ok := existing.(map[string]interface{}); ok {
				if newMap, ok := v.(map[string]interface{}); ok {
					result[k] = mergeConfig(existingMap, newMap)
					continue
				}
			}
			if existingSlice, ok := existing.([]interface{}); ok {
				if newSlice, ok := v.([]interface{}); ok {
					merged := append(existingSlice, newSlice...)
					seen := make(map[interface{}]bool)
					unique := make([]interface{}, 0)
					for _, item := range merged {
						if !seen[item] {
							seen[item] = true
							unique = append(unique, item)
						}
					}
					result[k] = unique
					continue
				}
			}
		}
		result[k] = v
	}

	return result
}

// DeleteByTenantID deletes all knowledge bases by tenant ID (hard delete)
func (dao *KnowledgebaseDAO) DeleteByTenantID(tenantID string) (int64, error) {
	result := DB.Unscoped().Where("tenant_id = ?", tenantID).Delete(&entity.Knowledgebase{})
	return result.RowsAffected, result.Error
}

// GetKBIDsByTenantID gets all knowledge base IDs by tenant ID
func (dao *KnowledgebaseDAO) GetKBIDsByTenantIDSimple(tenantID string) ([]string, error) {
	var kbIDs []string
	err := DB.Model(&entity.Knowledgebase{}).
		Where("tenant_id = ?", tenantID).
		Pluck("id", &kbIDs).Error
	return kbIDs, err
}