document.go

// document.go
package knowledgesdk

import (
	"context"
	"fmt"

	"gorm.io/gorm"
)

// AddDocument 添加文档并立即分块
func (k *KnowledgeSDK) AddDocument(
	ctx context.Context,
	kbID, name, content string,
	chunkConfig ChunkConfig,
) (*Document, error) {
	return k.AddDocumentWithMetadata(ctx, kbID, name, content, "", nil, "", chunkConfig)
}

// AddDocumentWithCreator 添加文档并指定创建者ID
func (k *KnowledgeSDK) AddDocumentWithCreator(
	ctx context.Context,
	kbID, name, content string,
	creatorID string,
	chunkConfig ChunkConfig,
) (*Document, error) {
	return k.AddDocumentWithMetadata(ctx, kbID, name, content, "", nil, creatorID, chunkConfig)
}

// splitDocumentIntoChunks 将文档内容分割成块并保存到数据库
func (k *KnowledgeSDK) splitDocumentIntoChunks(ctx context.Context, tx *gorm.DB, docID, content string, config ChunkConfig) error {
	contentRunes := []rune(content)
	contentLength := len(contentRunes)

	// 首先删除已存在的分块（如更新文档的情况）
	if err := tx.Where("document_id = ?", docID).Delete(&Chunk{}).Error; err != nil {
		return fmt.Errorf("删除已有分块失败: %w", err)
	}

	// 分块并插入
	pos := 0
	chunkIndex := 0
	for pos < contentLength {
		endPos := pos + config.ChunkSize
		if endPos > contentLength {
			endPos = contentLength
		}

		chunkText := string(contentRunes[pos:endPos])
		chunk := NewChunk(docID, chunkIndex, chunkText)

		if err := tx.Create(chunk).Error; err != nil {
			return fmt.Errorf("创建分块失败: %w", err)
		}

		chunkIndex++
		pos += config.ChunkSize - config.Overlap
		if pos >= contentLength {
			break
		}
	}

	return nil
}

// GetDocument 获取文档
func (k *KnowledgeSDK) GetDocument(ctx context.Context, docID string) (*Document, error) {
	var doc Document
	err := k.db.WithContext(ctx).Where("document_id = ?", docID).First(&doc).Error
	if err != nil {
		return nil, fmt.Errorf("获取文档失败: %w", err)
	}

	return &doc, nil
}

// GetDocumentWithChunks 获取文档及其分块
func (k *KnowledgeSDK) GetDocumentWithChunks(ctx context.Context, docID string) (*Document, error) {
	var doc Document
	err := k.db.WithContext(ctx).Where("document_id = ?", docID).First(&doc).Error
	if err != nil {
		return nil, fmt.Errorf("获取文档失败: %w", err)
	}

	var chunks []Chunk
	err = k.db.WithContext(ctx).Where("document_id = ?", docID).Find(&chunks).Error
	if err != nil {
		return nil, fmt.Errorf("获取文档分块失败: %w", err)
	}

	doc.Chunks = chunks
	return &doc, nil
}

// DeleteDocument 删除文档
func (k *KnowledgeSDK) DeleteDocument(ctx context.Context, docID string) error {
	return k.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		// 删除文档会级联删除分块
		if err := tx.Where("document_id = ?", docID).Delete(&Document{}).Error; err != nil {
			return fmt.Errorf("删除文档失败: %w", err)
		}
		return nil
	})
}

// AddDocumentWithMetadata 添加文档并立即分块，支持元数据和创建者ID
func (k *KnowledgeSDK) AddDocumentWithMetadata(
	ctx context.Context,
	kbID, name, content string,
	contentType string,
	metadata MetadataMap,
	creatorID string,
	chunkConfig ChunkConfig,
) (*Document, error) {
	// 检查知识库是否存在
	var kb KnowledgeBase
	if err := k.db.WithContext(ctx).Where("kb_id = ?", kbID).First(&kb).Error; err != nil {
		return nil, fmt.Errorf("找不到知识库: %w", err)
	}

	// 创建文档
	doc := NewDocument(kbID, name, content, creatorID)
	doc.ContentType = String(contentType)
	doc.Metadata = metadata
	doc.Status = String(DocStatusExtractSuccess) // 直接设置为抽取成功状态，准备切分

	// 开始事务
	err := k.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		// 创建文档
		if err := tx.Create(doc).Error; err != nil {
			return fmt.Errorf("创建文档失败: %w", err)
		}

		// 进行文档分块
		if err := k.splitDocumentIntoChunks(ctx, tx, StringValue(doc.ID), StringValue(doc.OriginalContent), chunkConfig); err != nil {
			// 如果分块失败，更新状态
			doc.Status = String(DocStatusSplitFailed)
			tx.Save(doc)
			return fmt.Errorf("文档分块失败: %w", err)
		}

		// 更新文档状态
		doc.Status = String(DocStatusSplitSuccess) // 切分成功，等待索引
		if err := tx.Save(doc).Error; err != nil {
			return fmt.Errorf("更新文档状态失败: %w", err)
		}

		return nil
	})

	if err != nil {
		return nil, err
	}

	return doc, nil
}