package rag

import (
	"context"
	"fmt"
	"strings"
)

// Retriever provides a high-level knowledge retrieval interface.
type Retriever struct {
	store *KnowledgeStore
}

// NewRetriever creates a new knowledge retriever.
func NewRetriever(store *KnowledgeStore) *Retriever {
	return &Retriever{store: store}
}

// Retrieve searches the knowledge base and returns formatted results.
func (r *Retriever) Retrieve(ctx context.Context, query string, topK int) (*RetrievalResult, error) {
	results, err := r.store.Search(ctx, query, topK)
	if err != nil {
		return nil, fmt.Errorf("knowledge search: %w", err)
	}

	ret := &RetrievalResult{
		Query:   query,
		Results: results,
		Summary: r.buildSummary(results),
	}
	return ret, nil
}

// RetrievalResult holds knowledge retrieval output.
type RetrievalResult struct {
	Query   string         `json:"query"`
	Results []SearchResult `json:"results"`
	Summary string         `json:"summary"`
}

func (r *Retriever) buildSummary(results []SearchResult) string {
	if len(results) == 0 {
		return "知识库中未找到相关信息。"
	}
	var sb strings.Builder
	sb.WriteString(fmt.Sprintf("从知识库中找到 %d 条相关信息:\n\n", len(results)))
	for i, result := range results {
		sb.WriteString(fmt.Sprintf("--- 来源: %s (段落 %d, 相关度 %.0f%%) ---\n",
			result.Chunk.DocTitle, result.Chunk.Index+1, result.Score*100))
		sb.WriteString(result.Chunk.Content)
		if i < len(results)-1 {
			sb.WriteString("\n\n")
		}
	}
	return sb.String()
}

// Stats returns knowledge base statistics.
func (r *Retriever) Stats() map[string]interface{} {
	return r.store.Stats()
}