qc-text-builder/internal/mapping/suggestion_generator.go

package mapping

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

	"qctextbuilder/internal/domain"
	"qctextbuilder/internal/qcclient"
)

type SuggestionGenerator interface {
	Generate(ctx context.Context, req SuggestionRequest) (SuggestionResult, error)
}

type RuleBasedSuggestionGenerator struct{}

func NewRuleBasedSuggestionGenerator() *RuleBasedSuggestionGenerator {
	return &RuleBasedSuggestionGenerator{}
}

func (g *RuleBasedSuggestionGenerator) Generate(_ context.Context, req SuggestionRequest) (SuggestionResult, error) {
	result := SuggestFieldValuesRuleBased(req)
	mappingLogger().Info("rule-based suggestion used",
		"component", "autofill",
		"step", "rule_based_fallback",
		"status", "success",
		"draft_id", strings.TrimSpace(req.DraftID),
		"template_id", req.TemplateID,
		"suggestion_count", len(result.Suggestions),
	)
	return result, nil
}

type LLMSuggestionGenerator struct {
	qc     qcclient.Client
	source string
}

func NewLLMSuggestionGenerator(qc qcclient.Client) *LLMSuggestionGenerator {
	return &LLMSuggestionGenerator{
		qc:     qc,
		source: domain.DraftSuggestionSourceLLM,
	}
}

func NewQCLLMSuggestionGenerator(qc qcclient.Client) *LLMSuggestionGenerator {
	return &LLMSuggestionGenerator{
		qc:     qc,
		source: domain.DraftSuggestionSourceQCLLM,
	}
}

func (g *LLMSuggestionGenerator) Generate(ctx context.Context, req SuggestionRequest) (SuggestionResult, error) {
	started := time.Now()
	source := ""
	if g != nil {
		source = strings.TrimSpace(g.source)
	}
	step := "qc_fallback_request"
	if source == domain.DraftSuggestionSourceLLM {
		step = "llm_request"
	}
	mappingLogger().InfoContext(ctx, "llm suggestion request",
		"component", "autofill",
		"step", step,
		"status", "start",
		"source", source,
		"draft_id", strings.TrimSpace(req.DraftID),
		"template_id", req.TemplateID,
	)
	if g == nil || g.qc == nil {
		mappingLogger().WarnContext(ctx, "llm suggestion request",
			"component", "autofill",
			"step", step,
			"status", "failed",
			"source", source,
			"draft_id", strings.TrimSpace(req.DraftID),
			"template_id", req.TemplateID,
			"error", "llm generator is not configured",
			"duration_ms", time.Since(started).Milliseconds(),
		)
		return SuggestionResult{}, fmt.Errorf("llm generator is not configured")
	}
	if req.TemplateID <= 0 {
		mappingLogger().WarnContext(ctx, "llm suggestion request",
			"component", "autofill",
			"step", step,
			"status", "failed",
			"source", source,
			"draft_id", strings.TrimSpace(req.DraftID),
			"template_id", req.TemplateID,
			"error", "template id is required for llm suggestions",
			"duration_ms", time.Since(started).Milliseconds(),
		)
		return SuggestionResult{}, fmt.Errorf("template id is required for llm suggestions")
	}

	fieldByPath := make(map[string]domain.TemplateField, len(req.Fields))
	for _, field := range req.Fields {
		if !field.IsEnabled || !strings.EqualFold(strings.TrimSpace(field.FieldKind), "text") {
			continue
		}
		fieldByPath[field.Path] = field
	}

	targets := collectSuggestionTargets(req.Fields, req.Existing, req.IncludeFilled)
	customTemplateData := map[string]any{
		"_autofill": map[string]any{
			"masterPrompt": strings.TrimSpace(req.MasterPrompt),
			"promptBlocks": enabledPromptBlocks(req.PromptBlocks),
			"draftContext": llmDraftContextMap(req.DraftContext),
			"semanticSlots": func() map[string]string {
				out := make(map[string]string, len(targets))
				for _, target := range targets {
					out[target.FieldPath] = target.Slot
				}
				return out
			}(),
		},
	}
	resp, _, err := g.qc.GenerateContent(ctx, qcclient.GenerateContentRequest{
		TemplateID:         req.TemplateID,
		GlobalData:         req.GlobalData,
		Empty:              true,
		ToneOfVoice:        contentTone(req.DraftContext),
		TargetAudience:     targetAudience(req),
		CustomTemplateData: customTemplateData,
	})
	if err != nil {
		mappingLogger().WarnContext(ctx, "llm suggestion request",
			"component", "autofill",
			"step", step,
			"status", "failed",
			"source", source,
			"draft_id", strings.TrimSpace(req.DraftID),
			"template_id", req.TemplateID,
			"error", shortErr(err),
			"duration_ms", time.Since(started).Milliseconds(),
		)
		return SuggestionResult{}, err
	}

	out := SuggestionResult{
		Suggestions: make([]Suggestion, 0, len(targets)),
		ByFieldPath: map[string]Suggestion{},
	}
	for _, target := range targets {
		field, ok := fieldByPath[target.FieldPath]
		if !ok {
			continue
		}
		sectionData := resp[field.Section]
		if sectionData == nil {
			continue
		}
		raw, ok := sectionData[field.KeyName]
		if !ok {
			continue
		}
		value := normalizeLLMValue(raw)
		if value == "" {
			continue
		}
		suggestion := Suggestion{
			FieldPath: target.FieldPath,
			Slot:      target.Slot,
			Value:     value,
			Reason:    "llm suggestion from template content generation",
			Source:    firstNonEmpty(strings.TrimSpace(g.source), domain.DraftSuggestionSourceLLM),
		}
		out.Suggestions = append(out.Suggestions, suggestion)
		out.ByFieldPath[target.FieldPath] = suggestion
	}
	mappingLogger().InfoContext(ctx, "llm suggestion request",
		"component", "autofill",
		"step", step,
		"status", "success",
		"source", source,
		"draft_id", strings.TrimSpace(req.DraftID),
		"template_id", req.TemplateID,
		"suggestion_count", len(out.Suggestions),
		"duration_ms", time.Since(started).Milliseconds(),
	)
	return out, nil
}

type CompositeSuggestionGenerator struct {
	Primary  SuggestionGenerator
	Fallback SuggestionGenerator
}

func NewCompositeSuggestionGenerator(primary, fallback SuggestionGenerator) *CompositeSuggestionGenerator {
	return &CompositeSuggestionGenerator{
		Primary:  primary,
		Fallback: fallback,
	}
}

func (g *CompositeSuggestionGenerator) Generate(ctx context.Context, req SuggestionRequest) (SuggestionResult, error) {
	started := time.Now()
	if g == nil {
		return SuggestionResult{}, fmt.Errorf("suggestion generator is not configured")
	}
	if g.Primary == nil {
		mappingLogger().InfoContext(ctx, "autofill fallback",
			"component", "autofill",
			"step", "fallback_attempt",
			"status", "attempted",
			"fallback_generator", generatorLabel(g.Fallback),
			"draft_id", strings.TrimSpace(req.DraftID),
			"template_id", req.TemplateID,
		)
		return generateFallback(ctx, g.Fallback, req)
	}

	primaryResult, err := g.Primary.Generate(ctx, req)
	if err != nil {
		mappingLogger().WarnContext(ctx, "autofill fallback",
			"component", "autofill",
			"step", "primary_failed",
			"status", "failed",
			"primary_generator", generatorLabel(g.Primary),
			"draft_id", strings.TrimSpace(req.DraftID),
			"template_id", req.TemplateID,
			"error", shortErr(err),
		)
		mappingLogger().InfoContext(ctx, "autofill fallback",
			"component", "autofill",
			"step", "qc_fallback",
			"status", "attempted",
			"fallback_generator", generatorLabel(g.Fallback),
			"draft_id", strings.TrimSpace(req.DraftID),
			"template_id", req.TemplateID,
		)
		return generateFallback(ctx, g.Fallback, req)
	}
	primaryResult = normalizeSuggestionResult(primaryResult, req.Fields, req.Existing, req.IncludeFilled)
	if g.Fallback == nil {
		mappingLogger().InfoContext(ctx, "autofill result",
			"component", "autofill",
			"step", "final",
			"status", "success",
			"source_path", "primary_only",
			"suggestion_count", len(primaryResult.Suggestions),
			"draft_id", strings.TrimSpace(req.DraftID),
			"template_id", req.TemplateID,
			"duration_ms", time.Since(started).Milliseconds(),
		)
		return primaryResult, nil
	}

	fallbackResult, fbErr := g.Fallback.Generate(ctx, req)
	if fbErr != nil {
		mappingLogger().WarnContext(ctx, "autofill fallback",
			"component", "autofill",
			"step", "qc_fallback",
			"status", "failed",
			"fallback_generator", generatorLabel(g.Fallback),
			"draft_id", strings.TrimSpace(req.DraftID),
			"template_id", req.TemplateID,
			"error", shortErr(fbErr),
		)
		mappingLogger().InfoContext(ctx, "autofill result",
			"component", "autofill",
			"step", "final",
			"status", "success",
			"source_path", "primary_only_fallback_failed",
			"suggestion_count", len(primaryResult.Suggestions),
			"draft_id", strings.TrimSpace(req.DraftID),
			"template_id", req.TemplateID,
			"duration_ms", time.Since(started).Milliseconds(),
		)
		return primaryResult, nil
	}
	mappingLogger().InfoContext(ctx, "autofill fallback",
		"component", "autofill",
		"step", "qc_fallback",
		"status", "success",
		"fallback_generator", generatorLabel(g.Fallback),
		"draft_id", strings.TrimSpace(req.DraftID),
		"template_id", req.TemplateID,
		"suggestion_count", len(fallbackResult.Suggestions),
	)
	fallbackResult = normalizeSuggestionResult(fallbackResult, req.Fields, req.Existing, req.IncludeFilled)
	merged := primaryResult
	if merged.ByFieldPath == nil {
		merged.ByFieldPath = map[string]Suggestion{}
	}
	for _, suggestion := range fallbackResult.Suggestions {
		if _, exists := merged.ByFieldPath[suggestion.FieldPath]; exists {
			continue
		}
		merged.Suggestions = append(merged.Suggestions, suggestion)
		merged.ByFieldPath[suggestion.FieldPath] = suggestion
	}
	sort.SliceStable(merged.Suggestions, func(i, j int) bool {
		return merged.Suggestions[i].FieldPath < merged.Suggestions[j].FieldPath
	})
	sourcePath := "primary_plus_fallback"
	if len(primaryResult.Suggestions) == 0 && len(fallbackResult.Suggestions) > 0 {
		sourcePath = "fallback_only"
	}
	ruleBasedCount := 0
	for _, suggestion := range merged.Suggestions {
		if strings.EqualFold(strings.TrimSpace(suggestion.Source), domain.DraftSuggestionSourceFallbackRuleBased) {
			ruleBasedCount++
		}
	}
	if ruleBasedCount > 0 {
		mappingLogger().InfoContext(ctx, "autofill fallback",
			"component", "autofill",
			"step", "rule_based_fallback",
			"status", "used",
			"draft_id", strings.TrimSpace(req.DraftID),
			"template_id", req.TemplateID,
			"suggestion_count", ruleBasedCount,
		)
	}
	mappingLogger().InfoContext(ctx, "autofill result",
		"component", "autofill",
		"step", "final",
		"status", "success",
		"source_path", sourcePath,
		"suggestion_count", len(merged.Suggestions),
		"draft_id", strings.TrimSpace(req.DraftID),
		"template_id", req.TemplateID,
		"duration_ms", time.Since(started).Milliseconds(),
	)
	return merged, nil
}

func generateFallback(ctx context.Context, fallback SuggestionGenerator, req SuggestionRequest) (SuggestionResult, error) {
	if fallback == nil {
		return SuggestionResult{}, fmt.Errorf("fallback suggestion generator is not configured")
	}
	started := time.Now()
	label := generatorLabel(fallback)
	mappingLogger().InfoContext(ctx, "autofill fallback",
		"component", "autofill",
		"step", "fallback_attempt",
		"status", "attempted",
		"fallback_generator", label,
		"draft_id", strings.TrimSpace(req.DraftID),
		"template_id", req.TemplateID,
	)
	result, err := fallback.Generate(ctx, req)
	if err != nil {
		mappingLogger().WarnContext(ctx, "autofill fallback",
			"component", "autofill",
			"step", "fallback_attempt",
			"status", "failed",
			"fallback_generator", label,
			"draft_id", strings.TrimSpace(req.DraftID),
			"template_id", req.TemplateID,
			"error", shortErr(err),
			"duration_ms", time.Since(started).Milliseconds(),
		)
		return SuggestionResult{}, err
	}
	mappingLogger().InfoContext(ctx, "autofill fallback",
		"component", "autofill",
		"step", "fallback_attempt",
		"status", "success",
		"fallback_generator", label,
		"draft_id", strings.TrimSpace(req.DraftID),
		"template_id", req.TemplateID,
		"suggestion_count", len(result.Suggestions),
		"duration_ms", time.Since(started).Milliseconds(),
	)
	return normalizeSuggestionResult(result, req.Fields, req.Existing, req.IncludeFilled), nil
}

func normalizeSuggestionResult(result SuggestionResult, fields []domain.TemplateField, existing map[string]string, includeFilled bool) SuggestionResult {
	allowed := make(map[string]SemanticSlotTarget)
	for _, target := range collectSuggestionTargets(fields, existing, includeFilled) {
		if _, exists := allowed[target.FieldPath]; exists {
			continue
		}
		allowed[target.FieldPath] = target
	}
	out := SuggestionResult{
		Suggestions: make([]Suggestion, 0, len(result.Suggestions)),
		ByFieldPath: map[string]Suggestion{},
	}
	for _, suggestion := range result.Suggestions {
		fieldPath := strings.TrimSpace(suggestion.FieldPath)
		if fieldPath == "" {
			continue
		}
		target, ok := allowed[fieldPath]
		if !ok {
			continue
		}
		value := strings.TrimSpace(suggestion.Value)
		if value == "" {
			continue
		}
		normalized := suggestion
		normalized.FieldPath = fieldPath
		if strings.TrimSpace(normalized.Slot) == "" {
			normalized.Slot = target.Slot
		}
		normalized.Value = value
		if strings.TrimSpace(normalized.Source) == "" {
			normalized.Source = domain.DraftSuggestionSourceFallbackRuleBased
		}
		if _, exists := out.ByFieldPath[fieldPath]; exists {
			continue
		}
		out.Suggestions = append(out.Suggestions, normalized)
		out.ByFieldPath[fieldPath] = normalized
	}
	sort.SliceStable(out.Suggestions, func(i, j int) bool {
		return out.Suggestions[i].FieldPath < out.Suggestions[j].FieldPath
	})
	return out
}

func collectSuggestionTargets(fields []domain.TemplateField, existing map[string]string, includeFilled bool) []SemanticSlotTarget {
	normalizedExisting := existing
	if normalizedExisting == nil {
		normalizedExisting = map[string]string{}
	}
	mappingResult := MapTemplateFieldsToSemanticSlots(fields)
	targets := append([]SemanticSlotTarget(nil), mappingResult.Targets...)
	sort.SliceStable(targets, func(i, j int) bool {
		if targets[i].FieldPath == targets[j].FieldPath {
			return targets[i].Slot < targets[j].Slot
		}
		return targets[i].FieldPath < targets[j].FieldPath
	})

	out := make([]SemanticSlotTarget, 0, len(targets))
	seen := map[string]struct{}{}
	for _, target := range targets {
		if _, exists := seen[target.FieldPath]; exists {
			continue
		}
		if !includeFilled && strings.TrimSpace(normalizedExisting[target.FieldPath]) != "" {
			continue
		}
		out = append(out, target)
		seen[target.FieldPath] = struct{}{}
	}
	return out
}

func enabledPromptBlocks(blocks []domain.PromptBlockConfig) []map[string]string {
	out := make([]map[string]string, 0, len(blocks))
	for _, block := range blocks {
		if !block.Enabled {
			continue
		}
		entry := map[string]string{"id": strings.TrimSpace(block.ID)}
		if label := strings.TrimSpace(block.Label); label != "" {
			entry["label"] = label
		}
		if instruction := strings.TrimSpace(block.Instruction); instruction != "" {
			entry["instruction"] = instruction
		}
		out = append(out, entry)
	}
	return out
}

func llmDraftContextMap(ctx *domain.DraftContext) map[string]any {
	if ctx == nil {
		return map[string]any{}
	}
	return map[string]any{
		"businessType":   strings.TrimSpace(ctx.LLM.BusinessType),
		"websiteUrl":     strings.TrimSpace(ctx.LLM.WebsiteURL),
		"websiteSummary": strings.TrimSpace(ctx.LLM.WebsiteSummary),
		"styleProfile": map[string]string{
			"localeStyle":        strings.TrimSpace(ctx.LLM.StyleProfile.LocaleStyle),
			"marketStyle":        strings.TrimSpace(ctx.LLM.StyleProfile.MarketStyle),
			"addressMode":        strings.TrimSpace(ctx.LLM.StyleProfile.AddressMode),
			"contentTone":        strings.TrimSpace(ctx.LLM.StyleProfile.ContentTone),
			"promptInstructions": strings.TrimSpace(ctx.LLM.StyleProfile.PromptInstructions),
		},
	}
}

func contentTone(ctx *domain.DraftContext) string {
	if ctx == nil {
		return ""
	}
	return strings.TrimSpace(ctx.LLM.StyleProfile.ContentTone)
}

func targetAudience(req SuggestionRequest) string {
	ctx := suggestionContextFrom(req.GlobalData, req.DraftContext)
	parts := make([]string, 0, 4)
	if ctx.BusinessType != "" {
		parts = append(parts, "businessType="+ctx.BusinessType)
	}
	if ctx.LocaleStyle != "" {
		parts = append(parts, "locale="+ctx.LocaleStyle)
	}
	if ctx.MarketStyle != "" {
		parts = append(parts, "market="+ctx.MarketStyle)
	}
	if ctx.AddressMode != "" {
		parts = append(parts, "addressMode="+ctx.AddressMode)
	}
	return strings.Join(parts, ", ")
}

func normalizeLLMValue(raw any) string {
	switch value := raw.(type) {
	case string:
		return strings.TrimSpace(value)
	default:
		return ""
	}
}