package knowledge

import (
	"context"
	"fmt"
	"resume-platform/internal/ai"
	"resume-platform/internal/prompt"
	"strings"
)

type RAGService struct {
	aiProvider            ai.Provider
	knowledgeBaseService *KnowledgeBaseService
}

func NewRAGService(provider ai.Provider, knowledgeBaseService *KnowledgeBaseService) *RAGService {
	return &RAGService{aiProvider: provider, knowledgeBaseService: knowledgeBaseService}
}

func (s *RAGService) Generate(ctx context.Context, query string, userID string) (string, error) {
	contextResults, err := s.knowledgeBaseService.Retrieve(ctx, query, 5, userID)
	if err != nil {
		return "", err
	}

	var contextStr string
	for _, result := range contextResults {
		contextStr += fmt.Sprintf("【来源：%s】\n%s\n\n", result.SourceName, result.Content)
	}

	if contextStr == "" {
		return s.aiProvider.Generate(ctx, query)
	}

	promptText := prompt.BuildRAGAnswerPrompt(contextStr, query)

	return s.aiProvider.Generate(ctx, promptText)
}

func (s *RAGService) GenerateWithContext(ctx context.Context, query string, context []string) (string, error) {
	contextStr := strings.Join(context, "\n\n")

	promptText := prompt.BuildRAGAnswerPrompt(contextStr, query)

	return s.aiProvider.Generate(ctx, promptText)
}