knowledge_base.go

// Package knowledgesdk knowledge_base.go
package knowledgesdk

import (
	"context"
	"errors"
	"fmt"

	"gorm.io/gorm"
)

// CreateKnowledgeBase 创建新的知识库
func (k *KnowledgeSDK) CreateKnowledgeBase(ctx context.Context, kb *KnowledgeBase) (*KnowledgeBase, error) {
	err := k.db.WithContext(ctx).Create(kb).Error
	if err != nil {
		return nil, fmt.Errorf("创建知识库失败: %w", err)
	}

	return kb, nil
}

// GetKnowledgeBase 获取知识库
func (k *KnowledgeSDK) GetKnowledgeBase(ctx context.Context, kbID string) (*KnowledgeBase, error) {
	var kb KnowledgeBase
	err := k.db.WithContext(ctx).Where("kb_id = ?", kbID).First(&kb).Error
	if err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return nil, NewError(ErrKnowledgeNotExist, "知识库不存在", err)
		}
		return nil, NewError(ErrDatabaseQueryFailed, "获取知识库失败", err)
	}

	return &kb, nil
}

// ListKnowledgeBases 列出所有知识库
func (k *KnowledgeSDK) ListKnowledgeBases(ctx context.Context) ([]KnowledgeBase, error) {
	var kbs []KnowledgeBase
	err := k.db.WithContext(ctx).Find(&kbs).Error
	if err != nil {
		return nil, fmt.Errorf("列出知识库失败: %w", err)
	}

	return kbs, nil
}

// ListKnowledgeBasesByCreatorID 根据创建者ID列出知识库
func (k *KnowledgeSDK) ListKnowledgeBasesByCreatorID(ctx context.Context, creatorID string) ([]KnowledgeBase, error) {
	var kbs []KnowledgeBase
	err := k.db.WithContext(ctx).Where("creator_id = ?", creatorID).Find(&kbs).Error
	if err != nil {
		return nil, fmt.Errorf("列出创建者的知识库失败: %w", err)
	}

	return kbs, nil
}

// ListKnowledgeBasesByIDs 根据知识库ID列表获取多个知识库
func (k *KnowledgeSDK) ListKnowledgeBasesByIDs(ctx context.Context, kbIDs []string) ([]KnowledgeBase, error) {
	if len(kbIDs) == 0 {
		return []KnowledgeBase{}, nil
	}

	var kbs []KnowledgeBase
	err := k.db.WithContext(ctx).Where("kb_id IN ?", kbIDs).Find(&kbs).Error
	if err != nil {
		return nil, fmt.Errorf("获取指定知识库列表失败: %w", err)
	}

	return kbs, nil
}

// UpdateKnowledgeBase 更新知识库
func (k *KnowledgeSDK) UpdateKnowledgeBase(ctx context.Context, kb *KnowledgeBase) (*KnowledgeBase, error) {
	result := k.db.WithContext(ctx).Model(&KnowledgeBase{}).Where("kb_id = ?", StringValue(kb.ID)).Updates(kb)
	if result.Error != nil {
		return nil, fmt.Errorf("更新知识库失败: %w", result.Error)
	}

	if result.RowsAffected == 0 {
		return nil, fmt.Errorf("找不到知识库")
	}

	// 获取更新后的记录
	var updatedKB KnowledgeBase
	err := k.db.WithContext(ctx).Where("kb_id = ?", StringValue(kb.ID)).First(&updatedKB).Error
	if err != nil {
		return nil, fmt.Errorf("获取更新后的知识库失败: %w", err)
	}

	return &updatedKB, nil
}

// DeleteKnowledgeBase 删除知识库
func (k *KnowledgeSDK) DeleteKnowledgeBase(ctx context.Context, kbID string) error {
	// 开始事务
	return k.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		// 检查知识库是否存在
		var kb KnowledgeBase
		if err := tx.Where("kb_id = ?", kbID).First(&kb).Error; err != nil {
			return fmt.Errorf("找不到知识库: %w", err)
		}

		// 删除知识库中的所有文档（级联删除会自动删除chunks）
		if err := tx.Where("kb_id = ?", kbID).Delete(&Document{}).Error; err != nil {
			return fmt.Errorf("删除知识库中的文档失败: %w", err)
		}

		// 删除知识库
		if err := tx.Delete(&kb).Error; err != nil {
			return fmt.Errorf("删除知识库失败: %w", err)
		}

		return nil
	})
}

// ListKnowledgeBaseDocuments 添加获取知识库文档列表的方法
func (k *KnowledgeSDK) ListKnowledgeBaseDocuments(ctx context.Context, kbID string) ([]Document, error) {
	var docs []Document
	err := k.db.WithContext(ctx).Where("kb_id = ?", kbID).Find(&docs).Error
	if err != nil {
		return nil, fmt.Errorf("获取知识库文档列表失败: %w", err)
	}

	return docs, nil
}

// ListKnowledgeBaseDocumentsPaginated 分页获取知识库文档列表并排序，支持文件名关键字匹配
func (k *KnowledgeSDK) ListKnowledgeBaseDocumentsPaginated(ctx context.Context, kbID string, keyword string, page, pageSize int, orderBy string, creatorID string) ([]Document, int64, error) {
	var docs []Document
	var total int64

	// 构建查询条件
	query := k.db.WithContext(ctx).Model(&Document{}).Where("kb_id = ?", kbID)

	// 如果提供了创建者ID，添加创建者ID过滤条件
	if creatorID != "" {
		query = query.Where("creator_id = ?", creatorID)
	}

	// 如果提供了关键字，添加文件名模糊匹配条件
	if keyword != "" {
		query = query.Where("document_name LIKE ?", "%"+keyword+"%")
	}

	// 先获取满足条件的总数
	err := query.Count(&total).Error
	if err != nil {
		return nil, 0, fmt.Errorf("获取知识库文档总数失败: %w", err)
	}

	// 计算偏移量
	offset := (page - 1) * pageSize

	// 设置默认排序
	if orderBy == "" {
		orderBy = "uploaded_at DESC" // 默认按上传时间降序排列
	}

	// 分页查询并排序
	err = query.Order(orderBy).Offset(offset).Limit(pageSize).Find(&docs).Error
	if err != nil {
		return nil, 0, fmt.Errorf("获取知识库文档列表失败: %w", err)
	}

	return docs, total, nil
}

// SearchKnowledgeBasesByName 根据名称模糊搜索知识库
func (k *KnowledgeSDK) SearchKnowledgeBasesByName(ctx context.Context, name string) ([]KnowledgeBase, error) {
	var kbs []KnowledgeBase
	err := k.db.WithContext(ctx).Where("kb_name LIKE ?", "%"+name+"%").Find(&kbs).Error
	if err != nil {
		return nil, fmt.Errorf("根据名称搜索知识库失败: %w", err)
	}

	return kbs, nil
}

// SearchKnowledgeBasesByDescription 根据描述模糊搜索知识库
func (k *KnowledgeSDK) SearchKnowledgeBasesByDescription(ctx context.Context, description string) ([]KnowledgeBase, error) {
	var kbs []KnowledgeBase
	err := k.db.WithContext(ctx).Where("kb_description LIKE ?", "%"+description+"%").Find(&kbs).Error
	if err != nil {
		return nil, fmt.Errorf("根据描述搜索知识库失败: %w", err)
	}

	return kbs, nil
}

// SearchKnowledgeBasesByKeyword 根据关键字搜索知识库（同时搜索名称和描述）
func (k *KnowledgeSDK) SearchKnowledgeBasesByKeyword(ctx context.Context, keyword string) ([]KnowledgeBase, error) {
	var kbs []KnowledgeBase
	err := k.db.WithContext(ctx).Where("kb_name LIKE ? OR kb_description LIKE ?", "%"+keyword+"%", "%"+keyword+"%").Find(&kbs).Error
	if err != nil {
		return nil, fmt.Errorf("根据关键字搜索知识库失败: %w", err)
	}

	return kbs, nil
}

// SearchKnowledgeBasesAdvanced 高级搜索知识库，支持多种条件组合
func (k *KnowledgeSDK) SearchKnowledgeBasesAdvanced(ctx context.Context, params KnowledgeBaseSearchParams) ([]KnowledgeBase, int64, error) {
	var kbs []KnowledgeBase
	var total int64

	// 构建查询条件
	query := k.db.WithContext(ctx).Model(&KnowledgeBase{})

	// 如果提供了创建者ID，添加创建者ID过滤条件
	if params.CreatorID != "" {
		query = query.Where("creator_id = ?", params.CreatorID)
	}

	// 如果提供了关键字，添加名称或描述模糊匹配条件
	if params.Keyword != "" {
		query = query.Where("kb_name LIKE ? OR kb_description LIKE ?", "%"+params.Keyword+"%", "%"+params.Keyword+"%")
	}

	// 如果提供了名称关键字，单独匹配名称
	if params.Name != "" {
		query = query.Where("kb_name LIKE ?", "%"+params.Name+"%")
	}

	// 如果提供了描述关键字，单独匹配描述
	if params.Description != "" {
		query = query.Where("kb_description LIKE ?", "%"+params.Description+"%")
	}

	// 如果提供了模型ID，精确匹配模型ID
	if params.ModelID != "" {
		query = query.Where("model_id = ?", params.ModelID)
	}

	// 先获取满足条件的总数
	err := query.Count(&total).Error
	if err != nil {
		return nil, 0, fmt.Errorf("获取知识库总数失败: %w", err)
	}

	// 计算偏移量
	offset := (params.Page - 1) * params.PageSize

	// 设置默认排序
	orderBy := params.OrderBy
	if orderBy == "" {
		orderBy = "created_at DESC" // 默认按创建时间降序排列
	}

	// 分页查询并排序
	err = query.Order(orderBy).Offset(offset).Limit(params.PageSize).Find(&kbs).Error
	if err != nil {
		return nil, 0, fmt.Errorf("搜索知识库失败: %w", err)
	}

	return kbs, total, nil
}

// ReplaceKnowledgeBaseModel 替换所有知识库的模型ID
func (k *KnowledgeSDK) ReplaceKnowledgeBaseModel(ctx context.Context, oldModelID, newModelID string) error {
	// 开始事务
	return k.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		// 更新所有知识库的模型ID
		result := tx.Model(&KnowledgeBase{}).Where("model_id = ?", oldModelID).Update("model_id", newModelID)
		if result.Error != nil {
			return fmt.Errorf("替换知识库模型ID失败: %w", result.Error)
		}

		if result.RowsAffected == 0 {
			return nil
		}

		return nil
	})
}