// Modified in the independent AgentDesk fork: localized knowledge workflows and safe display errors.
package rag

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/mlogclub/simple/sqls"

	"agent-desk/internal/ai"
	"agent-desk/internal/models"
	"agent-desk/internal/pkg/dto"
	"agent-desk/internal/pkg/dto/request"
	"agent-desk/internal/pkg/dto/response"
	"agent-desk/internal/pkg/enums"
	"agent-desk/internal/pkg/errorsx"
	"agent-desk/internal/pkg/i18nx"
	"agent-desk/internal/repositories"
)

type answer struct {
}

var Answer = &answer{}

func (s *answer) DebugSearch(ctx context.Context, req request.KnowledgeSearchRequest) (*response.KnowledgeSearchResponse, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if strings.TrimSpace(req.Question) == "" {
		return nil, errorsx.InvalidParamI18n("error.e0340")
	}
	startedAt := time.Now()
	results, err := s.retrieve(req, ctx)
	if err != nil {
		return nil, err
	}

	respResults := make([]response.KnowledgeSearchResult, 0, len(results))
	for _, item := range results {
		respResults = append(respResults, response.KnowledgeSearchResult{
			KnowledgeBaseID: item.KnowledgeBaseID,
			ChunkID:         item.ChunkID,
			DocumentID:      item.DocumentID,
			DocumentTitle:   item.DocumentTitle,
			FaqID:           item.FaqID,
			FaqQuestion:     item.FaqQuestion,
			ChunkNo:         item.ChunkNo,
			Title:           item.Title,
			SectionPath:     item.SectionPath,
			Content:         item.Content,
			Score:           float64(item.Score),
		})
	}

	return &response.KnowledgeSearchResponse{
		Question:  req.Question,
		Results:   respResults,
		HitCount:  len(respResults),
		LatencyMs: time.Since(startedAt).Milliseconds(),
	}, nil
}

func (s *answer) DebugAnswer(ctx context.Context, req request.KnowledgeAnswerRequest, operator *dto.AuthPrincipal) (*response.KnowledgeAnswerResponse, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if strings.TrimSpace(req.Question) == "" {
		return nil, errorsx.InvalidParamI18n("error.e0340")
	}
	startedAt := time.Now()

	retrieveStartedAt := time.Now()
	results, err := s.retrieve(request.KnowledgeSearchRequest{
		KnowledgeBaseIDs: req.KnowledgeBaseIDs,
		Question:         req.Question,
		TopK:             req.TopK,
		ScoreThreshold:   req.ScoreThreshold,
		RerankLimit:      req.RerankLimit,
	}, ctx)
	if err != nil {
		return nil, err
	}
	retrieveMs := time.Since(retrieveStartedAt).Milliseconds()
	knowledgeBase := s.resolveAnswerKnowledgeBase(req.KnowledgeBaseIDs, results)
	contextResults := buildContextHits(Retrieve.SelectContextResultsWithContext(ctx, results, 4000))

	hits := make([]response.KnowledgeSearchResult, 0, len(results))
	topScore := 0.0
	for i, item := range results {
		score := float64(item.Score)
		if i == 0 {
			topScore = score
		}
		hits = append(hits, response.KnowledgeSearchResult{
			KnowledgeBaseID: item.KnowledgeBaseID,
			ChunkID:         item.ChunkID,
			DocumentID:      item.DocumentID,
			DocumentTitle:   item.DocumentTitle,
			FaqID:           item.FaqID,
			FaqQuestion:     item.FaqQuestion,
			ChunkNo:         item.ChunkNo,
			Title:           item.Title,
			SectionPath:     item.SectionPath,
			Content:         item.Content,
			Score:           score,
		})
	}
	citations := buildKnowledgeCitations(hits, 3)

	answerMode := enums.KnowledgeAnswerMode(req.AnswerMode)
	if answerMode == 0 {
		if knowledgeBase != nil {
			answerMode = enums.KnowledgeAnswerMode(knowledgeBase.AnswerMode)
		}
		if answerMode == 0 {
			answerMode = enums.KnowledgeAnswerModeStrict
		}
	}

	fallbackMode := enums.AIAgentFallbackModeNoAnswer

	generateStartedAt := time.Now()
	generated, err := generateKnowledgeAnswer(ctx, req.Question, results, answerMode, fallbackMode, ai.LLM.Chat)
	if err != nil {
		return nil, err
	}
	answerStatus := generated.Status
	answer := generated.Content
	modelName := generated.ModelName
	promptTokens := generated.PromptTokens
	completionTokens := generated.CompletionTokens
	generateMs := time.Since(generateStartedAt).Milliseconds()

	rerankLimit := 0
	chunkProvider := ""
	chunkTargetTokens := 0
	chunkMaxTokens := 0
	chunkOverlapTokens := 0
	if knowledgeBase != nil {
		rerankLimit = resolveRerankLimit(req.RerankLimit, knowledgeBase.DefaultRerankLimit)
		chunkProvider = knowledgeBase.ChunkProvider
		chunkTargetTokens = knowledgeBase.ChunkTargetTokens
		chunkMaxTokens = knowledgeBase.ChunkMaxTokens
		chunkOverlapTokens = knowledgeBase.ChunkOverlapTokens
	}

	logItem, err := RetrieveLog.CreateRetrieveLog(&CreateRetrieveLogRequest{
		KnowledgeBaseID:    firstKnowledgeBaseID(req.KnowledgeBaseIDs),
		Channel:            defaultRetrieveChannel(req.Channel),
		Scene:              defaultRetrieveScene(req.Scene),
		SessionID:          req.SessionID,
		ConversationID:     req.ConversationID,
		Question:           req.Question,
		RewriteQuestion:    "",
		Answer:             answer,
		AnswerStatus:       int(answerStatus),
		ChunkProvider:      chunkProvider,
		ChunkTargetTokens:  chunkTargetTokens,
		ChunkMaxTokens:     chunkMaxTokens,
		ChunkOverlapTokens: chunkOverlapTokens,
		RerankEnabled:      rerankLimit > 0,
		RerankLimit:        rerankLimit,
		Hits:               hits,
		UsedHits:           contextResults,
		Citations:          citations,
		LatencyMs:          time.Since(startedAt).Milliseconds(),
		RetrieveMs:         retrieveMs,
		GenerateMs:         generateMs,
		PromptTokens:       promptTokens,
		CompletionTokens:   completionTokens,
		ModelName:          modelName,
	}, operator)
	if err != nil {
		return nil, err
	}

	return &response.KnowledgeAnswerResponse{
		Question:         req.Question,
		Answer:           answer,
		AnswerStatus:     int(answerStatus),
		AnswerStatusName: getAnswerStatusName(ctx, answerStatus),
		Citations:        citations,
		Hits:             hits,
		HitCount:         len(hits),
		TopScore:         topScore,
		LatencyMs:        time.Since(startedAt).Milliseconds(),
		RetrieveMs:       retrieveMs,
		GenerateMs:       generateMs,
		PromptTokens:     promptTokens,
		CompletionTokens: completionTokens,
		ModelName:        modelName,
		RetrieveLogID:    logItem.ID,
	}, nil
}

func buildContextHits(results []RetrieveResult) []response.KnowledgeSearchResult {
	if len(results) == 0 {
		return nil
	}
	hits := make([]response.KnowledgeSearchResult, 0, len(results))
	for _, item := range results {
		hits = append(hits, response.KnowledgeSearchResult{
			KnowledgeBaseID: item.KnowledgeBaseID,
			ChunkID:         item.ChunkID,
			DocumentID:      item.DocumentID,
			DocumentTitle:   item.DocumentTitle,
			FaqID:           item.FaqID,
			FaqQuestion:     item.FaqQuestion,
			ChunkNo:         item.ChunkNo,
			Title:           item.Title,
			SectionPath:     item.SectionPath,
			Content:         item.Content,
			Score:           float64(item.Score),
		})
	}
	return hits
}

func buildKnowledgeCitations(hits []response.KnowledgeSearchResult, limit int) []response.KnowledgeCitation {
	if len(hits) == 0 || limit <= 0 {
		return nil
	}
	citations := make([]response.KnowledgeCitation, 0, limit)
	seen := make(map[string]struct{})
	for _, item := range hits {
		key := fmt.Sprintf("%d|%d|%s|%d", item.DocumentID, item.FaqID, item.SectionPath, item.ChunkNo)
		if _, ok := seen[key]; ok {
			continue
		}
		seen[key] = struct{}{}
		citations = append(citations, response.KnowledgeCitation{
			DocumentID:    item.DocumentID,
			DocumentTitle: item.DocumentTitle,
			FaqID:         item.FaqID,
			FaqQuestion:   item.FaqQuestion,
			ChunkNo:       item.ChunkNo,
			Title:         item.Title,
			SectionPath:   item.SectionPath,
			Snippet:       truncateCitationSnippet(item.Content, 120),
			Score:         item.Score,
		})
		if len(citations) >= limit {
			break
		}
	}
	return citations
}

func truncateCitationSnippet(text string, limit int) string {
	if limit <= 0 {
		return ""
	}
	runes := []rune(strings.TrimSpace(text))
	if len(runes) <= limit {
		return string(runes)
	}
	return string(runes[:limit]) + "..."
}

func (s *answer) BuildDocumentIndex(ctx context.Context, documentID int64) error {
	return Index.IndexDocumentByID(ctx, documentID)
}

func (s *answer) retrieve(req request.KnowledgeSearchRequest, ctx context.Context) ([]RetrieveResult, error) {
	if len(normalizeKnowledgeBaseIDs(req.KnowledgeBaseIDs)) == 0 {
		return nil, errorsx.InvalidParamI18n("error.e0284")
	}
	knowledgeBases := s.loadKnowledgeBases(req.KnowledgeBaseIDs)

	results, err := Retrieve.Retrieve(ctx, RetrieveRequest{
		KnowledgeBaseIDs: req.KnowledgeBaseIDs,
		Query:            req.Question,
		TopK:             req.TopK,
		ScoreThreshold:   req.ScoreThreshold,
	})
	if err != nil {
		return nil, err
	}

	defaultRerankLimit := resolveDefaultRerankLimit(knowledgeBases)
	rerankLimit := resolveRerankLimit(req.RerankLimit, defaultRerankLimit)
	if rerankLimit > 0 && len(results) > rerankLimit {
		return Retrieve.RetrieveWithRerank(ctx, RetrieveRequest{
			KnowledgeBaseIDs: req.KnowledgeBaseIDs,
			Query:            req.Question,
			TopK:             req.TopK,
			ScoreThreshold:   req.ScoreThreshold,
		}, rerankLimit)
	}
	return results, nil
}

func (s *answer) loadKnowledgeBases(knowledgeBaseIDs []int64) []models.KnowledgeBase {
	normalized := normalizeKnowledgeBaseIDs(knowledgeBaseIDs)
	if len(normalized) == 0 {
		return nil
	}
	items := repositories.KnowledgeBaseRepository.Find(sqls.DB(), sqls.NewCnd().In("id", normalized))
	if len(items) == 0 {
		return nil
	}
	itemMap := make(map[int64]models.KnowledgeBase, len(items))
	for _, item := range items {
		itemMap[item.ID] = item
	}
	results := make([]models.KnowledgeBase, 0, len(normalized))
	for _, id := range normalized {
		if item, ok := itemMap[id]; ok {
			results = append(results, item)
		}
	}
	return results
}

func (s *answer) resolvePrimaryKnowledgeBase(knowledgeBaseIDs []int64) *models.KnowledgeBase {
	items := s.loadKnowledgeBases(knowledgeBaseIDs)
	for _, item := range items {
		return &item
	}
	return nil
}

func (s *answer) resolveAnswerKnowledgeBase(knowledgeBaseIDs []int64, results []RetrieveResult) *models.KnowledgeBase {
	items := s.loadKnowledgeBases(knowledgeBaseIDs)
	if len(items) == 0 {
		return nil
	}
	if len(results) > 0 {
		for _, item := range items {
			if item.ID == results[0].KnowledgeBaseID {
				return &item
			}
		}
	}
	return &items[0]
}

func firstKnowledgeBaseID(ids []int64) int64 {
	normalized := normalizeKnowledgeBaseIDs(ids)
	if len(normalized) == 0 {
		return 0
	}
	return normalized[0]
}

func resolveRerankLimit(requestLimit int, defaultLimit int) int {
	if requestLimit > 0 {
		return requestLimit
	}
	if defaultLimit > 0 {
		return defaultLimit
	}
	return 0
}

func resolveDefaultRerankLimit(items []models.KnowledgeBase) int {
	limit := 0
	for _, item := range items {
		if item.DefaultRerankLimit > limit {
			limit = item.DefaultRerankLimit
		}
	}
	return limit
}

func buildAnswerSystemPrompt(ctx context.Context, answerMode enums.KnowledgeAnswerMode) string {
	key := "knowledge.rag.system.strict"
	if answerMode == enums.KnowledgeAnswerModeAssist {
		key = "knowledge.rag.system.assist"
	}
	return i18nx.Getf(i18nx.LocaleFromContext(ctx), key)
}

func buildFallbackAnswer(ctx context.Context, fallbackMode enums.AIAgentFallbackMode) string {
	key := "knowledge.rag.fallback.noAnswer"
	if fallbackMode == enums.AIAgentFallbackModeSuggestRetry {
		key = "knowledge.rag.fallback.suggestRetry"
	}
	return i18nx.Getf(i18nx.LocaleFromContext(ctx), key)
}

func defaultRetrieveChannel(channel string) string {
	if strings.TrimSpace(channel) == "" {
		return string(enums.KnowledgeRetrieveChannelDebug)
	}
	return channel
}

func defaultRetrieveScene(scene string) string {
	if strings.TrimSpace(scene) == "" {
		return string(enums.KnowledgeRetrieveSceneQA)
	}
	return scene
}

func getAnswerStatusName(ctx context.Context, status enums.KnowledgeAnswerStatus) string {
	key := "knowledge.labels.unknown"
	switch status {
	case enums.KnowledgeAnswerStatusNormal, enums.KnowledgeAnswerStatusNoAnswer, enums.KnowledgeAnswerStatusFallback, enums.KnowledgeAnswerStatusBlocked:
		key = fmt.Sprintf("knowledge.labels.answer.%d", status)
	}
	return i18nx.Getf(i18nx.LocaleFromContext(ctx), key)
}

type generatedKnowledgeAnswer struct {
	ai.ChatCompletionResult
	Status enums.KnowledgeAnswerStatus
}

type knowledgeChatFunc func(context.Context, string, string) (*ai.ChatCompletionResult, error)

// Keep generation independent of persistence so fallback and cancellation paths
// can be tested without a configured model, external service, or database.
func generateKnowledgeAnswer(ctx context.Context, question string, results []RetrieveResult, mode enums.KnowledgeAnswerMode, fallbackMode enums.AIAgentFallbackMode, chat knowledgeChatFunc) (*generatedKnowledgeAnswer, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	result := &generatedKnowledgeAnswer{Status: enums.KnowledgeAnswerStatusNoAnswer}
	result.Content = buildFallbackAnswer(ctx, fallbackMode)
	if len(results) == 0 {
		return result, nil
	}
	contextText := Retrieve.BuildContext(ctx, results, 4000)
	systemPrompt := buildAnswerSystemPrompt(ctx, mode)
	userPrompt := i18nx.Getf(i18nx.LocaleFromContext(ctx), "knowledge.rag.userPrompt", question, contextText)
	llmResult, err := chat(ctx, systemPrompt, userPrompt)
	// An interrupted request is not a model failure and must not be logged as a
	// completed fallback answer after the client has gone away.
	if contextErr := ctx.Err(); contextErr != nil {
		return nil, contextErr
	}
	if err != nil || llmResult == nil {
		result.Status = enums.KnowledgeAnswerStatusFallback
		return result, nil
	}
	result.ChatCompletionResult = *llmResult
	if strings.TrimSpace(llmResult.Content) == "" {
		result.Status = enums.KnowledgeAnswerStatusFallback
		result.Content = buildFallbackAnswer(ctx, fallbackMode)
		return result, nil
	}
	result.Status = enums.KnowledgeAnswerStatusNormal
	return result, nil
}