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

package mapping

import (
	"context"
	"encoding/json"
	"fmt"
	"regexp"
	"sort"
	"strings"
	"time"

	"qctextbuilder/internal/domain"
	"qctextbuilder/internal/llmruntime"
)

// AISlotMapper classifies template fields into semantic slots via the active
// LLM provider. It is meant to run once per template (during onboarding) and
// its results are persisted on the manifest fields.
type AISlotMapper struct {
	settings       SettingsReader
	runtimeFactory *llmruntime.Factory
}

func NewAISlotMapper(settings SettingsReader, runtimeFactory *llmruntime.Factory) *AISlotMapper {
	return &AISlotMapper{settings: settings, runtimeFactory: runtimeFactory}
}

type AISlotAssignment struct {
	FieldPath  string  `json:"fieldPath"`
	Slot       string  `json:"slot"`
	Confidence float64 `json:"confidence"`
	Reason     string  `json:"reason"`
}

var aiSlotRepeatedPattern = regexp.MustCompile(`^(service_items|team_items|testimonial_items)\[\d+\]\.(title|description|name)$`)

var aiSlotSingletons = map[string]struct{}{
	"hero.title":        {},
	"intro.title":       {},
	"intro.description": {},
	"about.description": {},
	"cta.text":          {},
}

func isValidSemanticSlot(slot string) bool {
	if slot == "" {
		return false
	}
	if _, ok := aiSlotSingletons[slot]; ok {
		return true
	}
	if !aiSlotRepeatedPattern.MatchString(slot) {
		return false
	}
	// team_items[N].title is not a valid combination — name is the only label role.
	switch {
	case strings.HasPrefix(slot, "team_items[") && strings.HasSuffix(slot, ".title"):
		return false
	case strings.HasPrefix(slot, "service_items[") && strings.HasSuffix(slot, ".name"):
		return false
	}
	return true
}

// MapFields returns one assignment per input field whose Slot is a valid
// semantic slot. Fields the model decided to skip are omitted.
func (m *AISlotMapper) MapFields(ctx context.Context, templateID int64, fields []domain.TemplateField) ([]AISlotAssignment, error) {
	started := time.Now()
	if m == nil || m.settings == nil || m.runtimeFactory == nil {
		return nil, fmt.Errorf("ai slot mapper is not configured")
	}

	textFields := make([]domain.TemplateField, 0, len(fields))
	for _, f := range fields {
		if !f.IsEnabled {
			continue
		}
		if !strings.EqualFold(strings.TrimSpace(f.FieldKind), "text") {
			continue
		}
		textFields = append(textFields, f)
	}
	if len(textFields) == 0 {
		return []AISlotAssignment{}, nil
	}

	settings, err := m.settings.GetSettings(ctx)
	if err != nil || settings == nil {
		return nil, fmt.Errorf("llm settings are not available")
	}
	provider := domain.NormalizeLLMProvider(settings.LLMActiveProvider)
	model := domain.NormalizeLLMModel(provider, settings.LLMActiveModel)
	if strings.TrimSpace(model) == "" {
		return nil, fmt.Errorf("no active llm model configured")
	}
	apiKey := domain.LLMAPIKeyForProvider(provider, *settings)
	if provider != domain.LLMProviderOllama && strings.TrimSpace(apiKey) == "" {
		return nil, fmt.Errorf("api key for provider %s is not configured", provider)
	}
	baseURL := strings.TrimSpace(settings.LLMBaseURL)

	client, err := m.runtimeFactory.ClientFor(provider)
	if err != nil {
		return nil, err
	}

	systemPrompt, userPrompt := buildAISlotPrompts(textFields)
	mappingLogger().InfoContext(ctx, "ai slot mapping",
		"component", "automapper",
		"step", "request",
		"status", "start",
		"provider", provider,
		"model", model,
		"template_id", templateID,
		"field_count", len(textFields),
	)
	temperature := 0.0
	maxTokens := domain.NormalizeLLMMaxTokens(settings.LLMMaxTokens)
	if maxTokens < 2048 {
		maxTokens = 2048
	}
	raw, err := client.Generate(ctx, llmruntime.Request{
		Provider:     provider,
		Model:        model,
		BaseURL:      baseURL,
		APIKey:       strings.TrimSpace(apiKey),
		Temperature:  &temperature,
		MaxTokens:    &maxTokens,
		SystemPrompt: systemPrompt,
		UserPrompt:   userPrompt,
	})
	if err != nil {
		mappingLogger().WarnContext(ctx, "ai slot mapping",
			"component", "automapper",
			"step", "request",
			"status", "failed",
			"provider", provider,
			"model", model,
			"template_id", templateID,
			"error", shortErr(err),
			"duration_ms", time.Since(started).Milliseconds(),
		)
		return nil, fmt.Errorf("ai slot mapper request failed (provider=%s model=%s): %w", provider, model, err)
	}

	parsed, err := parseAISlotResponse(raw)
	if err != nil {
		mappingLogger().WarnContext(ctx, "ai slot mapping",
			"component", "automapper",
			"step", "parse",
			"status", "failed",
			"provider", provider,
			"model", model,
			"template_id", templateID,
			"error", shortErr(err),
			"response_snippet", providerLogSnippet(raw, 1500),
			"duration_ms", time.Since(started).Milliseconds(),
		)
		return nil, fmt.Errorf("ai slot mapper returned invalid json: %w", err)
	}

	allowed := make(map[string]struct{}, len(textFields))
	for _, f := range textFields {
		allowed[strings.TrimSpace(f.Path)] = struct{}{}
	}
	out := make([]AISlotAssignment, 0, len(parsed))
	seen := make(map[string]struct{}, len(parsed))
	for _, item := range parsed {
		path := strings.TrimSpace(item.FieldPath)
		if path == "" {
			continue
		}
		if _, ok := allowed[path]; !ok {
			continue
		}
		if _, dup := seen[path]; dup {
			continue
		}
		slot := strings.TrimSpace(item.Slot)
		if !isValidSemanticSlot(slot) {
			continue
		}
		seen[path] = struct{}{}
		out = append(out, AISlotAssignment{
			FieldPath:  path,
			Slot:       slot,
			Confidence: item.Confidence,
			Reason:     strings.TrimSpace(item.Reason),
		})
	}
	sort.SliceStable(out, func(i, j int) bool { return out[i].FieldPath < out[j].FieldPath })

	mappingLogger().InfoContext(ctx, "ai slot mapping",
		"component", "automapper",
		"step", "result",
		"status", "success",
		"provider", provider,
		"model", model,
		"template_id", templateID,
		"input_count", len(textFields),
		"mapped_count", len(out),
		"duration_ms", time.Since(started).Milliseconds(),
	)
	return out, nil
}

func buildAISlotPrompts(fields []domain.TemplateField) (string, string) {
	type fieldPayload struct {
		Path         string `json:"path"`
		KeyName      string `json:"keyName"`
		Section      string `json:"section"`
		Block        string `json:"block,omitempty"`
		DisplayOrder int    `json:"displayOrder"`
		SampleValue  string `json:"sampleValue,omitempty"`
	}
	payload := make([]fieldPayload, 0, len(fields))
	for _, f := range fields {
		payload = append(payload, fieldPayload{
			Path:         strings.TrimSpace(f.Path),
			KeyName:      strings.TrimSpace(f.KeyName),
			Section:      strings.TrimSpace(f.Section),
			Block:        semanticExtractBlockID(f),
			DisplayOrder: f.DisplayOrder,
			SampleValue:  truncateSample(f.SampleValue, 200),
		})
	}
	body := map[string]any{
		"fields": payload,
	}
	bodyJSON, _ := json.MarshalIndent(body, "", "  ")

	system := strings.TrimSpace(`
You are a website-template field mapper. You classify each text field of a QC
website template into ONE semantic slot, or skip it.

Allowed slots:
- hero.title                          (one per template, topmost large heading)
- intro.title, intro.description      (one each, secondary heading/intro block)
- about.description                   (one, longer "about us" paragraph)
- cta.text                            (one, call-to-action button / highlight)
- service_items[N].title              (N = 0,1,2... one per service card)
- service_items[N].description
- team_items[N].name
- team_items[N].description
- testimonial_items[N].title
- testimonial_items[N].name
- testimonial_items[N].description

Rules:
1. Use the section (e.g. "services", "testimonials") as a strong hint.
2. Use the block prefix (e.g. "c8987", "r1865") to group title+description from
   the same UI block together.
3. Use displayOrder to decide which block is hero vs intro vs about (lower order = earlier on page).
4. Use sampleValue length: short (<60 chars) = title/name; long (>120 chars) = description.
5. For repeated slots, pair each block to ONE item index; use the same N for
   title and description coming from the same block. Index from 0, increasing
   by block displayOrder.
6. If you cannot map a field to one of the listed slots, output "" (empty) — never invent slots.
7. Each fieldPath appears exactly once in the output.

Return ONLY JSON, no markdown, no commentary:
{"mappings":[{"fieldPath":"...","slot":"hero.title","confidence":0.9,"reason":"..."}]}
`)

	user := "Classify the following fields. Respond with JSON only.\n\n" + string(bodyJSON)
	return system, user
}

func parseAISlotResponse(raw string) ([]AISlotAssignment, error) {
	content := strings.TrimSpace(raw)
	if content == "" {
		return nil, fmt.Errorf("empty response")
	}
	candidates := []string{content}
	if fenced := extractFencedJSON(content); fenced != "" {
		candidates = append([]string{fenced}, candidates...)
	}
	if obj := extractJSONObject(content); obj != "" {
		candidates = append(candidates, obj)
	}

	var firstErr error
	for _, candidate := range candidates {
		items, err := parseAISlotPayload(candidate)
		if err == nil {
			return items, nil
		}
		if firstErr == nil {
			firstErr = err
		}
	}
	if firstErr != nil {
		return nil, firstErr
	}
	return nil, fmt.Errorf("response is not a valid mapping json")
}

func parseAISlotPayload(raw string) ([]AISlotAssignment, error) {
	var root any
	if err := json.Unmarshal([]byte(raw), &root); err != nil {
		return nil, err
	}
	var itemsRaw []any
	switch value := root.(type) {
	case map[string]any:
		list, ok := value["mappings"].([]any)
		if !ok {
			// also tolerate {"suggestions":...} or any single array property
			for _, key := range []string{"suggestions", "results", "items"} {
				if arr, ok := value[key].([]any); ok {
					list = arr
					break
				}
			}
		}
		if list == nil {
			return nil, fmt.Errorf("object must contain \"mappings\" array")
		}
		itemsRaw = list
	case []any:
		itemsRaw = value
	default:
		return nil, fmt.Errorf("payload must be object or array")
	}
	out := make([]AISlotAssignment, 0, len(itemsRaw))
	for idx, item := range itemsRaw {
		m, ok := item.(map[string]any)
		if !ok {
			return nil, fmt.Errorf("entry #%d is not an object", idx+1)
		}
		path := strings.TrimSpace(anyToString(m["fieldPath"]))
		if path == "" {
			path = strings.TrimSpace(anyToString(m["path"]))
		}
		if path == "" {
			continue
		}
		confidence := 0.0
		if v, ok := m["confidence"].(float64); ok {
			confidence = v
		}
		out = append(out, AISlotAssignment{
			FieldPath:  path,
			Slot:       strings.TrimSpace(anyToString(m["slot"])),
			Confidence: confidence,
			Reason:     strings.TrimSpace(anyToString(m["reason"])),
		})
	}
	return out, nil
}

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