geo/server/internal/tenant/app/question_expansion_service.go

package app

import (
	"context"
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"errors"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/geo-platform/tenant-api/internal/shared/auditlog"
	"github.com/geo-platform/tenant-api/internal/shared/auth"
	sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache"
	"github.com/geo-platform/tenant-api/internal/shared/llm"
	"github.com/geo-platform/tenant-api/internal/shared/middleware"
	"github.com/geo-platform/tenant-api/internal/shared/response"
)

const (
	questionCombinationMaxItems = 500
	questionAIDistillMaxItems   = 20
	questionDistillCacheTTL     = 5 * time.Minute
	questionDistillTimeout      = 30 * time.Second
)

type QuestionExpansionService struct {
	pool  *pgxpool.Pool
	llm   llm.Client
	brand *BrandService
	cache sharedcache.Cache
}

func NewQuestionExpansionService(pool *pgxpool.Pool, llmClient llm.Client, brand *BrandService) *QuestionExpansionService {
	return &QuestionExpansionService{
		pool:  pool,
		llm:   llmClient,
		brand: brand,
	}
}

func (s *QuestionExpansionService) WithCache(c sharedcache.Cache) *QuestionExpansionService {
	s.cache = c
	return s
}

type QuestionCombinationRequest struct {
	Region   []string `json:"region"`
	Prefix   []string `json:"prefix"`
	Core     []string `json:"core" binding:"required"`
	Industry []string `json:"industry" binding:"required"`
	Suffix   []string `json:"suffix"`
	MaxItems int      `json:"max_items"`
}

type QuestionDistillRequest struct {
	SeedTopic string `json:"seed_topic" binding:"required,min=2"`
}

type QuestionCandidate struct {
	Text          string `json:"text"`
	Layer         string `json:"layer"`
	Intent        string `json:"intent"`
	Source        string `json:"source"`
	MissingSuffix bool   `json:"missing_suffix"`
	TooShort      bool   `json:"too_short"`
	Duplicate     bool   `json:"duplicate"`
	SuggestSkip   bool   `json:"suggest_skip"`
}

type QuestionCandidateResult struct {
	Candidates      []QuestionCandidate `json:"candidates"`
	AIPointsCharged int                 `json:"ai_points_charged,omitempty"`
	CacheHit        bool                `json:"cache_hit,omitempty"`
	Truncated       bool                `json:"truncated,omitempty"`
}

type MaterializeQuestionsRequest struct {
	Source    string                `json:"source"`
	Questions []MaterializeQuestion `json:"questions" binding:"required,min=1"`
}

type MaterializeQuestion struct {
	Text string `json:"text" binding:"required"`
}

type MaterializeQuestionsResult struct {
	CreatedQuestions int          `json:"created_questions"`
	SkippedQuestions []SkipReason `json:"skipped_questions"`
}

type SkipReason struct {
	Text   string `json:"text"`
	Reason string `json:"reason"`
}

type questionBrandContext struct {
	TenantID        int64
	BrandID         int64
	BrandName       string
	PlanCode        string
	CompetitorNames []string
}

type aiDistillPayload struct {
	Candidates []struct {
		Text   string `json:"text"`
		Intent string `json:"intent"`
		Layer  string `json:"layer"`
	} `json:"candidates"`
}

func (s *QuestionExpansionService) GenerateByCombination(ctx context.Context, brandID int64, req QuestionCombinationRequest) (*QuestionCandidateResult, error) {
	actor := auth.MustActor(ctx)
	brandCtx, err := s.loadQuestionBrandContext(ctx, actor.TenantID, brandID)
	if err != nil {
		return nil, err
	}

	core := normalizeQuestionParts(req.Core)
	industry := normalizeQuestionParts(req.Industry)
	if len(core) == 0 {
		return nil, response.ErrBadRequest(40091, "invalid_params", "core_required")
	}
	if len(industry) == 0 {
		return nil, response.ErrBadRequest(40091, "invalid_params", "industry_required")
	}

	region := optionalQuestionParts(req.Region)
	prefix := optionalQuestionParts(req.Prefix)
	suffix := optionalQuestionParts(req.Suffix)
	limit := req.MaxItems
	if limit <= 0 || limit > questionCombinationMaxItems {
		limit = questionCombinationMaxItems
	}

	existing, err := s.loadExistingQuestionKeys(ctx, actor.TenantID, brandID)
	if err != nil {
		return nil, err
	}

	candidates := make([]QuestionCandidate, 0, minInt(limit, 64))
	seen := map[string]struct{}{}
	truncated := false

	for _, rg := range region {
		for _, pf := range prefix {
			for _, co := range core {
				for _, in := range industry {
					for _, sf := range suffix {
						text := strings.TrimSpace(strings.Join(nonEmptyParts(rg, pf, co, in, sf), ""))
						if text == "" {
							continue
						}
						if sf != "" && !strings.HasSuffix(text, "?") && !strings.HasSuffix(text, "？") {
							if !strings.Contains(sf, "?") && !strings.Contains(sf, "？") {
								text += "？"
							}
						}
						key := normalizeQuestionKey(text)
						if _, ok := seen[key]; ok {
							continue
						}
						seen[key] = struct{}{}
						if len(candidates) >= limit {
							truncated = true
							continue
						}
						candidates = append(candidates, buildCandidate(text, QuestionSourceCombination, sf == "", existing, brandCtx))
					}
				}
			}
		}
	}

	return &QuestionCandidateResult{Candidates: candidates, Truncated: truncated}, nil
}

func (s *QuestionExpansionService) GenerateByAIDistill(ctx context.Context, brandID int64, req QuestionDistillRequest) (*QuestionCandidateResult, error) {
	actor := auth.MustActor(ctx)
	seedTopic := normalizeQuestionText(req.SeedTopic)
	if utf8.RuneCountInString(seedTopic) < 2 {
		return nil, response.ErrBadRequest(40091, "invalid_params", "seed_topic is required")
	}

	brandCtx, err := s.loadQuestionBrandContext(ctx, actor.TenantID, brandID)
	if err != nil {
		return nil, err
	}

	existing, err := s.loadExistingQuestionKeys(ctx, actor.TenantID, brandID)
	if err != nil {
		return nil, err
	}

	cacheKey := questionDistillCacheKey(actor.TenantID, brandID, questionDistillPromptVersion, brandCtx.BrandName, brandCtx.CompetitorNames, seedTopic)
	if s.cache != nil {
		if raw, cacheErr := s.cache.Get(ctx, cacheKey); cacheErr == nil && len(raw) > 0 {
			var cached []QuestionCandidate
			if json.Unmarshal(raw, &cached) == nil {
				return &QuestionCandidateResult{Candidates: cached, CacheHit: true}, nil
			}
		}
	}

	if s.llm == nil {
		return nil, response.ErrServiceUnavailable(50311, "llm_unavailable", "llm client is not configured")
	}
	if err := s.llm.Validate(); err != nil {
		return nil, response.ErrServiceUnavailable(50311, "llm_unavailable", err.Error())
	}

	resourceType := "question_distill"
	resourceUID := cacheKey
	reservation, err := ReserveAIPoints(ctx, s.pool, s.cache, AIPointReserveInput{
		TenantID:     actor.TenantID,
		OperatorID:   actor.UserID,
		UsageType:    AIUsageTypeQuestionDistill,
		ResourceType: &resourceType,
		ResourceUID:  &resourceUID,
		MeteredText:  seedTopic,
		FixedPoints:  1,
		Metadata: map[string]any{
			"brand_id":       brandID,
			"seed_topic":     seedTopic,
			"prompt_version": questionDistillPromptVersion,
		},
	})
	if err != nil {
		return nil, err
	}

	prompt := buildQuestionDistillPrompt(brandCtx.BrandName, brandCtx.CompetitorNames, seedTopic)
	result, err := s.llm.Generate(ctx, llm.GenerateRequest{
		Prompt:          prompt,
		Timeout:         questionDistillTimeout,
		MaxOutputTokens: 1600,
		ResponseFormat: &llm.ResponseFormat{
			Type:        llm.ResponseFormatTypeJSONSchema,
			Name:        "question_distill_candidates",
			Description: "Question expansion candidates for a brand.",
			SchemaJSON:  questionDistillSchema,
			Strict:      true,
		},
	}, nil)
	if err != nil {
		_ = RefundAIPoints(context.Background(), s.pool, s.cache, actor.TenantID, actor.UserID, *reservation, err.Error())
		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
			return nil, response.ErrServiceUnavailable(50312, "llm_timeout", "AI generation timed out")
		}
		return nil, response.ErrServiceUnavailable(50311, "llm_unavailable", err.Error())
	}

	var payload aiDistillPayload
	if err := json.Unmarshal([]byte(result.Content), &payload); err != nil {
		_ = RefundAIPoints(context.Background(), s.pool, s.cache, actor.TenantID, actor.UserID, *reservation, err.Error())
		return nil, response.ErrBadRequest(50210, "llm_invalid_output", "AI output could not be parsed")
	}

	candidates := make([]QuestionCandidate, 0, minInt(len(payload.Candidates), questionAIDistillMaxItems))
	seen := map[string]struct{}{}
	for _, item := range payload.Candidates {
		if len(candidates) >= questionAIDistillMaxItems {
			break
		}
		text := normalizeQuestionText(item.Text)
		key := normalizeQuestionKey(text)
		if text == "" {
			continue
		}
		if _, ok := seen[key]; ok {
			continue
		}
		seen[key] = struct{}{}
		candidates = append(candidates, buildCandidate(text, QuestionSourceAIDistill, false, existing, brandCtx))
	}

	if err := CompleteAIPoints(context.Background(), s.pool, actor.TenantID, *reservation, result.Model); err != nil {
		return nil, response.ErrInternal(50127, "ai_points_commit_failed", "failed to confirm ai points")
	}

	if s.cache != nil {
		if raw, err := json.Marshal(candidates); err == nil {
			_ = s.cache.Set(context.Background(), cacheKey, raw, questionDistillCacheTTL)
		}
	}

	return &QuestionCandidateResult{
		Candidates:      candidates,
		AIPointsCharged: reservation.Points,
	}, nil
}

func (s *QuestionExpansionService) ClassifyMetadata(ctx context.Context, brandID int64, texts []string) ([]ClassifiedQuestion, error) {
	actor := auth.MustActor(ctx)
	brandCtx, err := s.loadQuestionBrandContext(ctx, actor.TenantID, brandID)
	if err != nil {
		return nil, err
	}
	result := make([]ClassifiedQuestion, 0, len(texts))
	for _, text := range texts {
		trimmed := normalizeQuestionText(text)
		if trimmed == "" {
			continue
		}
		result = append(result, classifyQuestionText(trimmed, brandCtx.BrandName, brandCtx.CompetitorNames))
	}
	return result, nil
}

func (s *QuestionExpansionService) MaterializeQuestions(ctx context.Context, brandID int64, req MaterializeQuestionsRequest) (*MaterializeQuestionsResult, error) {
	actor := auth.MustActor(ctx)
	source := strings.TrimSpace(req.Source)
	if source == "" {
		source = QuestionSourceManual
	}
	if !isValidExternalQuestionSource(source) {
		return nil, response.ErrBadRequest(40092, "invalid_enum", "source is invalid")
	}
	if len(req.Questions) == 0 {
		return nil, response.ErrBadRequest(40093, "no_valid_questions", "questions is required")
	}

	brandCtx, err := s.loadQuestionBrandContext(ctx, actor.TenantID, brandID)
	if err != nil {
		return nil, err
	}
	existing, err := s.loadExistingQuestionKeys(ctx, actor.TenantID, brandID)
	if err != nil {
		return nil, err
	}
	currentCount, err := s.countTenantQuestions(ctx, actor.TenantID)
	if err != nil {
		return nil, err
	}

	maxQuestions := s.brand.currentLimits().QuestionLimitForPlan(brandCtx.PlanCode)
	remaining := maxQuestions - currentCount
	if remaining < 0 {
		remaining = 0
	}

	type acceptedQuestion struct {
		Text   string `json:"text"`
		Layer  string `json:"layer"`
		Intent string `json:"intent"`
		Source string `json:"source"`
	}

	accepted := make([]acceptedQuestion, 0, len(req.Questions))
	skipped := make([]SkipReason, 0)
	inBatch := map[string]struct{}{}

	for _, item := range req.Questions {
		text := normalizeQuestionText(item.Text)
		key := normalizeQuestionKey(text)
		if text == "" || utf8.RuneCountInString(text) < 4 {
			skipped = append(skipped, SkipReason{Text: text, Reason: "too_short"})
			continue
		}
		if !questionLooksValid(text) {
			skipped = append(skipped, SkipReason{Text: text, Reason: "invalid_text"})
			continue
		}
		if _, ok := existing[key]; ok {
			skipped = append(skipped, SkipReason{Text: text, Reason: "duplicate"})
			continue
		}
		if _, ok := inBatch[key]; ok {
			skipped = append(skipped, SkipReason{Text: text, Reason: "duplicate_in_batch"})
			continue
		}
		if remaining <= 0 {
			skipped = append(skipped, SkipReason{Text: text, Reason: "quota_exceeded"})
			continue
		}
		classified := classifyQuestionText(text, brandCtx.BrandName, brandCtx.CompetitorNames)
		accepted = append(accepted, acceptedQuestion{
			Text:   text,
			Layer:  classified.Layer,
			Intent: classified.Intent,
			Source: source,
		})
		inBatch[key] = struct{}{}
		remaining--
	}

	if len(accepted) == 0 {
		return nil, response.ErrBadRequest(40093, "no_valid_questions", skippedReasonDetail(skipped))
	}

	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return nil, response.ErrInternal(50010, "materialize_failed", "failed to begin transaction")
	}
	defer func() {
		_ = tx.Rollback(ctx)
	}()

	if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtext($1), hashtext($2))`, tenantQuestionQuotaLockKey(actor.TenantID), "questions"); err != nil {
		return nil, response.ErrInternal(50010, "materialize_failed", "failed to lock question quota")
	}

	bucketID, err := s.ensureDefaultQuestionBucket(ctx, tx, actor.TenantID, brandID, brandCtx.BrandName)
	if err != nil {
		return nil, err
	}

	currentCount, err = countTenantQuestionsTx(ctx, tx, actor.TenantID)
	if err != nil {
		return nil, response.ErrInternal(50010, "materialize_failed", "failed to recount questions")
	}
	remaining = maxQuestions - currentCount
	if remaining < 0 {
		remaining = 0
	}
	if remaining < len(accepted) {
		overflow := accepted[remaining:]
		accepted = accepted[:remaining]
		for _, q := range overflow {
			skipped = append(skipped, SkipReason{Text: q.Text, Reason: "quota_exceeded_concurrent"})
		}
	}
	if len(accepted) == 0 {
		return nil, response.ErrBadRequest(40093, "no_valid_questions", skippedReasonDetail(skipped))
	}

	payload, err := json.Marshal(accepted)
	if err != nil {
		return nil, response.ErrInternal(50010, "materialize_failed", "failed to encode questions")
	}
	rows, err := tx.Query(ctx, `
		INSERT INTO brand_questions (
			tenant_id, brand_id, keyword_id, question_text,
			layer, intent, source, status
		)
		SELECT $1, $2, $3, q.text, q.layer, q.intent, q.source, 'active'
		FROM jsonb_to_recordset($4::jsonb)
			AS q(text TEXT, layer TEXT, intent TEXT, source TEXT)
		ON CONFLICT (tenant_id, brand_id, lower(btrim(question_text)))
			WHERE deleted_at IS NULL
		DO NOTHING
		RETURNING question_text
	`, actor.TenantID, brandID, bucketID, payload)
	if err != nil {
		return nil, response.ErrInternal(50010, "materialize_failed", "failed to insert questions")
	}
	inserted := map[string]struct{}{}
	for rows.Next() {
		var text string
		if err := rows.Scan(&text); err != nil {
			rows.Close()
			return nil, response.ErrInternal(50010, "materialize_failed", "failed to scan inserted questions")
		}
		inserted[normalizeQuestionKey(text)] = struct{}{}
	}
	if err := rows.Err(); err != nil {
		rows.Close()
		return nil, response.ErrInternal(50010, "materialize_failed", "failed to iterate inserted questions")
	}
	rows.Close()

	for _, q := range accepted {
		if _, ok := inserted[normalizeQuestionKey(q.Text)]; !ok {
			skipped = append(skipped, SkipReason{Text: q.Text, Reason: "duplicate_concurrent"})
		}
	}

	if err := tx.Commit(ctx); err != nil {
		return nil, response.ErrInternal(50010, "materialize_failed", "failed to commit questions")
	}

	invalidateBrandCaches(ctx, s.cache, actor.TenantID, brandID)
	s.logQuestionExpansionAudit(ctx, actor.UserID, actor.TenantID, brandID, source, len(inserted), skipped)
	return &MaterializeQuestionsResult{
		CreatedQuestions: len(inserted),
		SkippedQuestions: skipped,
	}, nil
}

func (s *QuestionExpansionService) loadQuestionBrandContext(ctx context.Context, tenantID, brandID int64) (*questionBrandContext, error) {
	var brandName string
	err := s.pool.QueryRow(ctx, `
		SELECT name
		FROM brands
		WHERE id = $1 AND tenant_id = $2 AND deleted_at IS NULL
	`, brandID, tenantID).Scan(&brandName)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, response.ErrNotFound(40420, "brand_not_found", "brand not found")
		}
		return nil, response.ErrInternal(50010, "query_failed", "failed to load brand")
	}

	rows, err := s.pool.Query(ctx, `
		SELECT name
		FROM competitors
		WHERE brand_id = $1 AND tenant_id = $2 AND deleted_at IS NULL
		ORDER BY id ASC
	`, brandID, tenantID)
	if err != nil {
		return nil, response.ErrInternal(50010, "query_failed", "failed to load competitors")
	}
	defer rows.Close()
	competitors := make([]string, 0)
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			return nil, response.ErrInternal(50010, "scan_failed", err.Error())
		}
		competitors = append(competitors, name)
	}

	plan, err := s.brand.loadBrandLibraryPlan(ctx, tenantID)
	if err != nil {
		return nil, err
	}

	return &questionBrandContext{
		TenantID:        tenantID,
		BrandID:         brandID,
		BrandName:       brandName,
		PlanCode:        plan.PlanCode,
		CompetitorNames: competitors,
	}, nil
}

func (s *QuestionExpansionService) loadExistingQuestionKeys(ctx context.Context, tenantID, brandID int64) (map[string]struct{}, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT lower(btrim(question_text))
		FROM brand_questions
		WHERE tenant_id = $1 AND brand_id = $2 AND deleted_at IS NULL
	`, tenantID, brandID)
	if err != nil {
		return nil, response.ErrInternal(50010, "query_failed", "failed to load existing questions")
	}
	defer rows.Close()

	result := map[string]struct{}{}
	for rows.Next() {
		var key string
		if err := rows.Scan(&key); err != nil {
			return nil, response.ErrInternal(50010, "scan_failed", err.Error())
		}
		result[key] = struct{}{}
	}
	return result, nil
}

func (s *QuestionExpansionService) countTenantQuestions(ctx context.Context, tenantID int64) (int, error) {
	return countTenantQuestionsTx(ctx, s.pool, tenantID)
}

func countTenantQuestionsTx(ctx context.Context, db interface {
	QueryRow(context.Context, string, ...any) pgx.Row
}, tenantID int64) (int, error) {
	var count int
	err := db.QueryRow(ctx, `
		SELECT COUNT(*)::INT
		FROM brand_questions
		WHERE tenant_id = $1 AND deleted_at IS NULL
	`, tenantID).Scan(&count)
	return count, err
}

func (s *QuestionExpansionService) ensureDefaultQuestionBucket(ctx context.Context, tx pgx.Tx, tenantID, brandID int64, brandName string) (int64, error) {
	return ensureDefaultQuestionBucketTx(ctx, tx, tenantID, brandID, brandName)
}

func (s *QuestionExpansionService) logQuestionExpansionAudit(ctx context.Context, operatorID, tenantID, brandID int64, source string, created int, skipped []SkipReason) {
	if s == nil || s.brand == nil || s.brand.auditLogs == nil {
		return
	}
	afterJSON, _ := json.Marshal(map[string]interface{}{
		"brand_id":          brandID,
		"source":            source,
		"created_questions": created,
		"skipped_count":     len(skipped),
	})
	result := "success"
	resourceType := "brand"
	requestID := middleware.RequestIDFromContext(ctx)
	s.brand.auditLogs.Log(auditlog.Entry{
		OperatorID:   operatorID,
		TenantID:     &tenantID,
		Module:       "brand",
		Action:       "question_expansion.materialize",
		ResourceType: &resourceType,
		ResourceID:   &brandID,
		RequestID:    nilIfEmptyString(requestID),
		AfterJSON:    afterJSON,
		Result:       &result,
	})
}

func buildCandidate(text, source string, missingSuffix bool, existing map[string]struct{}, brandCtx *questionBrandContext) QuestionCandidate {
	classified := classifyQuestionText(text, brandCtx.BrandName, brandCtx.CompetitorNames)
	key := normalizeQuestionKey(text)
	_, duplicate := existing[key]
	tooShort := utf8.RuneCountInString(classified.Text) < 4
	invalid := !questionLooksValid(classified.Text)
	return QuestionCandidate{
		Text:          classified.Text,
		Layer:         classified.Layer,
		Intent:        classified.Intent,
		Source:        source,
		MissingSuffix: missingSuffix,
		TooShort:      tooShort,
		Duplicate:     duplicate,
		SuggestSkip:   tooShort || invalid || duplicate,
	}
}

func normalizeQuestionParts(values []string) []string {
	result := make([]string, 0, len(values))
	seen := map[string]struct{}{}
	for _, value := range values {
		trimmed := strings.TrimSpace(value)
		if trimmed == "" {
			continue
		}
		key := strings.ToLower(trimmed)
		if _, ok := seen[key]; ok {
			continue
		}
		seen[key] = struct{}{}
		result = append(result, trimmed)
	}
	return result
}

func optionalQuestionParts(values []string) []string {
	parts := normalizeQuestionParts(values)
	if len(parts) == 0 {
		return []string{""}
	}
	return parts
}

func nonEmptyParts(values ...string) []string {
	result := make([]string, 0, len(values))
	for _, value := range values {
		if strings.TrimSpace(value) != "" {
			result = append(result, strings.TrimSpace(value))
		}
	}
	return result
}

func questionDistillCacheKey(tenantID, brandID int64, promptVersion, brandName string, competitors []string, seedTopic string) string {
	raw, _ := json.Marshal(struct {
		TenantID      int64    `json:"tenant_id"`
		BrandID       int64    `json:"brand_id"`
		PromptVersion string   `json:"prompt_version"`
		SchemaVersion string   `json:"schema_version"`
		BrandName     string   `json:"brand_name"`
		Competitors   []string `json:"competitors"`
		SeedTopic     string   `json:"seed_topic"`
	}{
		TenantID:      tenantID,
		BrandID:       brandID,
		PromptVersion: promptVersion,
		SchemaVersion: "question_distill_schema_v1",
		BrandName:     strings.TrimSpace(brandName),
		Competitors:   competitors,
		SeedTopic:     strings.TrimSpace(seedTopic),
	})
	sum := sha1.Sum(raw)
	return "question_distill:" + hex.EncodeToString(sum[:])
}

func skippedReasonDetail(skipped []SkipReason) string {
	if len(skipped) == 0 {
		return "no valid questions"
	}
	raw, err := json.Marshal(skipped)
	if err != nil {
		return "no valid questions"
	}
	return string(raw)
}

func isUniqueQuestionConstraintError(err error) bool {
	var pgErr *pgconn.PgError
	return errors.As(err, &pgErr) && pgErr.Code == "23505" && pgErr.ConstraintName == "uk_brand_question_text_active"
}