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

package app

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"

	"github.com/geo-platform/tenant-api/internal/shared/auth"
	"github.com/geo-platform/tenant-api/internal/shared/response"
	"github.com/geo-platform/tenant-api/internal/tenant/repository"
)

type RegenerateArticleResponse struct {
	ArticleID int64 `json:"article_id"`
	TaskID    int64 `json:"task_id"`
}

type articleRegenerationRecord struct {
	ArticleID       int64
	SourceType      string
	TemplateID      *int64
	KolPromptID     *int64
	PromptRuleID    *int64
	GenerateStatus  string
	WizardStateJSON []byte
	SourceTaskID    int64
	TaskType        string
	RequestHash     *string
	InputParamsJSON []byte
	InputParams     map[string]interface{}
}

func (s *ArticleService) Regenerate(ctx context.Context, articleID int64) (*RegenerateArticleResponse, error) {
	actor := auth.MustActor(ctx)
	brandID, err := requireCurrentBrandID(ctx)
	if err != nil {
		return nil, err
	}
	if s == nil || s.llm == nil || s.rabbitMQ == nil {
		return nil, response.ErrServiceUnavailable(50310, "generation_queue_unavailable", "generation queue is not configured")
	}
	if err := s.llm.Validate(); err != nil {
		return nil, response.ErrServiceUnavailable(50311, "llm_unavailable", err.Error())
	}

	quotaRepo := repository.NewQuotaRepository(s.pool)
	balance, err := quotaRepo.GetCurrentBalance(ctx, actor.TenantID, "article_generation")
	if err != nil || balance < 1 {
		return nil, response.ErrForbidden(40301, "quota_insufficient", "generation quota insufficient, please upgrade your plan")
	}

	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return nil, fmt.Errorf("begin regeneration tx: %w", err)
	}
	defer rollbackGenerationTx(tx)

	record, err := loadArticleForRegeneration(ctx, tx, actor.TenantID, brandID, articleID)
	if err != nil {
		return nil, err
	}
	if err := validateArticleRegeneration(record); err != nil {
		return nil, err
	}

	quotaTx := repository.NewQuotaRepository(tx)
	auditTx := repository.NewAuditRepository(tx)

	newBalance := balance - 1
	reason := "generation_reserve"
	referenceType := "generation_task"
	if _, err := quotaTx.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
		TenantID:      actor.TenantID,
		OperatorID:    actor.UserID,
		QuotaType:     "article_generation",
		Delta:         -1,
		BalanceAfter:  newBalance,
		Reason:        &reason,
		ReferenceType: &referenceType,
	}); err != nil {
		return nil, fmt.Errorf("reserve regeneration quota: %w", err)
	}

	reservationID, err := quotaTx.CreateQuotaReservation(ctx, repository.QuotaReservationInput{
		TenantID:       actor.TenantID,
		OperatorID:     actor.UserID,
		QuotaType:      "article_generation",
		ResourceType:   "article",
		ResourceID:     nil,
		ReservedAmount: 1,
		ExpireAt:       time.Now().Add(time.Hour),
	})
	if err != nil {
		return nil, fmt.Errorf("create regeneration reservation: %w", err)
	}
	if err := quotaTx.UpdateQuotaReservationResource(ctx, reservationID, actor.TenantID, articleID); err != nil {
		return nil, fmt.Errorf("update regeneration reservation resource: %w", err)
	}

	title := regenerationInitialTitle(record)
	if err := updateGeneratedArticleState(ctx, tx, articleID, actor.TenantID, title, "generating", record.WizardStateJSON); err != nil {
		return nil, fmt.Errorf("mark article regenerating: %w", err)
	}

	requestHash := record.RequestHash
	if requestHash == nil {
		if hash := strings.TrimSpace(extractString(record.InputParams, "rendered_prompt_hash")); hash != "" {
			requestHash = &hash
		}
	}

	taskID, err := auditTx.CreateGenerationTask(ctx, repository.GenerationTaskInput{
		OperatorID:         actor.UserID,
		TenantID:           actor.TenantID,
		ArticleID:          &articleID,
		QuotaReservationID: &reservationID,
		TaskType:           record.TaskType,
		RequestHash:        requestHash,
		InputParamsJSON:    record.InputParamsJSON,
	})
	if err != nil {
		return nil, fmt.Errorf("create regeneration task: %w", err)
	}

	if record.TaskType == kolGenerationTaskType {
		if _, err := auditTx.CreateTaskRecord(ctx, repository.TaskRecordInput{
			TenantID:     actor.TenantID,
			OperatorID:   actor.UserID,
			TaskType:     kolGenerationTaskRecordType,
			ResourceType: kolGenerationTaskRecordScope,
			ResourceID:   taskID,
			Status:       "queued",
		}); err != nil {
			return nil, fmt.Errorf("create kol regeneration task record: %w", err)
		}
		if err := createKolUsageForRegeneration(ctx, tx, actor.TenantID, taskID, record.InputParams); err != nil {
			return nil, fmt.Errorf("create kol regeneration usage: %w", err)
		}
	}

	if err := tx.Commit(ctx); err != nil {
		return nil, fmt.Errorf("commit regeneration task: %w", err)
	}
	invalidateArticleCaches(ctx, s.cache, actor.TenantID, &articleID)

	task := ClaimedGenerationTask{
		ID:                 taskID,
		TenantID:           actor.TenantID,
		OperatorID:         actor.UserID,
		ArticleID:          articleID,
		QuotaReservationID: reservationID,
		TaskType:           record.TaskType,
		InputParams:        cloneInputParams(record.InputParams),
	}
	if err := publishArticleGenerationTask(ctx, s.rabbitMQ, articleGenerationTaskEnvelope{
		TaskID:    taskID,
		TaskType:  record.TaskType,
		TenantID:  actor.TenantID,
		ArticleID: articleID,
	}); err != nil {
		RecordGenerationTaskEvent(GenerationTaskEventQueueFailure, GenerationTaskLogContext{
			TaskID:        taskID,
			TenantID:      actor.TenantID,
			OperatorID:    actor.UserID,
			ArticleID:     articleID,
			ReservationID: reservationID,
			TaskType:      record.TaskType,
			Stage:         "queue_enqueue",
			StatusBefore:  "queued",
			StatusAfter:   "failed",
			Result:        GenerationTaskResultFailure,
		})
		s.failRegenerationQueue(context.Background(), task, title, err)
		return nil, response.ErrServiceUnavailable(50310, "generation_queue_unavailable", "generation queue is busy, please retry")
	}

	runtimeCfg := s.runtimeConfig()
	if runtimeCfg.StreamEnabled && s.streams != nil {
		s.streams.Start(articleID, taskID, title)
	}

	return &RegenerateArticleResponse{ArticleID: articleID, TaskID: taskID}, nil
}

func loadArticleForRegeneration(ctx context.Context, tx pgx.Tx, tenantID, brandID, articleID int64) (*articleRegenerationRecord, error) {
	record := &articleRegenerationRecord{}
	var templateID, kolPromptID, promptRuleID sql.NullInt64
	err := tx.QueryRow(ctx, `
		SELECT id, source_type, template_id, kol_prompt_id, prompt_rule_id, generate_status, wizard_state_json
		FROM articles
		WHERE id = $1 AND tenant_id = $2 AND brand_id = $3 AND deleted_at IS NULL
		FOR UPDATE
	`, articleID, tenantID, brandID).Scan(
		&record.ArticleID,
		&record.SourceType,
		&templateID,
		&kolPromptID,
		&promptRuleID,
		&record.GenerateStatus,
		&record.WizardStateJSON,
	)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, response.ErrNotFound(40411, "article_not_found", "article not found")
		}
		return nil, response.ErrInternal(50010, "query_failed", "failed to fetch article")
	}
	record.TemplateID = nullableSQLInt64(templateID)
	record.KolPromptID = nullableSQLInt64(kolPromptID)
	record.PromptRuleID = nullableSQLInt64(promptRuleID)

	var requestHash sql.NullString
	err = tx.QueryRow(ctx, `
		SELECT id, task_type, request_hash, input_params_json
		FROM generation_tasks
		WHERE article_id = $1
		  AND tenant_id = $2
		  AND status = 'failed'
		ORDER BY (COALESCE(input_params_json, '{}'::jsonb) = '{}'::jsonb), created_at DESC, id DESC
		LIMIT 1
	`, articleID, tenantID).Scan(
		&record.SourceTaskID,
		&record.TaskType,
		&requestHash,
		&record.InputParamsJSON,
	)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, response.ErrConflict(40914, "generation_task_missing", "failed generation task is missing")
		}
		return nil, response.ErrInternal(50010, "query_failed", "failed to fetch generation task")
	}
	if requestHash.Valid {
		record.RequestHash = &requestHash.String
	}
	if len(record.InputParamsJSON) == 0 {
		record.InputParamsJSON = []byte("{}")
	}
	params := map[string]interface{}{}
	if err := json.Unmarshal(record.InputParamsJSON, &params); err != nil {
		return nil, response.ErrConflict(40915, "generation_params_invalid", "failed generation params are invalid")
	}
	record.InputParams = params
	return record, nil
}

func validateArticleRegeneration(record *articleRegenerationRecord) error {
	if record == nil {
		return response.ErrNotFound(40411, "article_not_found", "article not found")
	}
	if record.GenerateStatus != "failed" {
		return response.ErrConflict(40913, "article_not_regeneratable", "article cannot be regenerated in its current state")
	}

	switch record.SourceType {
	case "template":
		if record.TaskType == "template" && record.TemplateID != nil {
			return nil
		}
	case "custom_generation":
		if record.TaskType == "custom_generation" {
			return nil
		}
	case articleImitationSourceType:
		if record.TaskType == articleImitationTaskType {
			return nil
		}
	case kolGenerationTaskType:
		if record.TaskType == kolGenerationTaskType {
			return nil
		}
	}

	return response.ErrConflict(40913, "article_not_regeneratable", "article cannot be regenerated in its current state")
}

func regenerationInitialTitle(record *articleRegenerationRecord) string {
	if record == nil {
		return "Generated Article"
	}

	switch record.TaskType {
	case "template":
		return resolveArticleTitle(record.InputParams, "")
	case "custom_generation":
		if title := strings.TrimSpace(extractString(record.InputParams, "task_name")); title != "" {
			return title
		}
		return promptRuleGeneratingTitlePlaceholder
	case articleImitationTaskType:
		return buildImitationInitialTitle(extractString(record.InputParams, "source_title"))
	case kolGenerationTaskType:
		if title := strings.TrimSpace(extractString(record.InputParams, "prompt_name")); title != "" {
			return title
		}
	}

	if title := strings.TrimSpace(extractString(record.InputParams, "title")); title != "" {
		return title
	}
	return "Generated Article"
}

func (s *ArticleService) failRegenerationQueue(ctx context.Context, task ClaimedGenerationTask, title string, taskErr error) {
	if task.TaskType == kolGenerationTaskType {
		_ = HandleKolGenerationTaskFailure(ctx, s.pool, s.cache, task, "queue_enqueue", taskErr)
	} else {
		_ = HandleClaimedGenerationTaskFailure(ctx, s.pool, nil, nil, task, taskErr)
	}

	runtimeCfg := s.runtimeConfig()
	if runtimeCfg.StreamEnabled && s.streams != nil {
		s.streams.Fail(task.ArticleID, task.ID, title, formatGenerationFailure("queue_enqueue", taskErr))
	}
}

func createKolUsageForRegeneration(ctx context.Context, tx pgx.Tx, tenantID, taskID int64, params map[string]interface{}) error {
	packageID := int64(extractInt(params, "package_id"))
	promptID := int64(extractInt(params, "prompt_id"))
	if packageID <= 0 || promptID <= 0 {
		return nil
	}

	promptRevisionNo := extractInt(params, "prompt_revision_no")
	if promptRevisionNo <= 0 {
		promptRevisionNo = 1
	}
	subscriptionID := optionalInt64FromParams(params, "subscription_id")
	subscriptionPromptID := optionalInt64FromParams(params, "subscription_prompt_id")
	variablesSnapshot, err := marshalOptionalJSONField(params, "variables")
	if err != nil {
		return err
	}
	schemaSnapshot, err := marshalOptionalJSONField(params, "schema_snapshot")
	if err != nil {
		return err
	}
	renderedPromptHash := strings.TrimSpace(extractString(params, "rendered_prompt_hash"))
	var renderedPromptHashPtr *string
	if renderedPromptHash != "" {
		renderedPromptHashPtr = &renderedPromptHash
	}

	_, err = repository.NewKolUsageRepository(tx).Create(ctx, repository.CreateKolUsageLogInput{
		TenantID:             tenantID,
		SubscriptionID:       subscriptionID,
		SubscriptionPromptID: subscriptionPromptID,
		PackageID:            packageID,
		PromptID:             promptID,
		PromptRevisionNo:     promptRevisionNo,
		GenerationTaskID:     &taskID,
		VariablesSnapshot:    variablesSnapshot,
		SchemaSnapshot:       schemaSnapshot,
		RenderedPromptHash:   renderedPromptHashPtr,
	})
	return err
}

func marshalOptionalJSONField(params map[string]interface{}, key string) ([]byte, error) {
	if len(params) == 0 {
		return nil, nil
	}
	value, ok := params[key]
	if !ok || value == nil {
		return nil, nil
	}
	return json.Marshal(value)
}

func optionalInt64FromParams(params map[string]interface{}, key string) *int64 {
	value := int64(extractInt(params, key))
	if value <= 0 {
		return nil
	}
	return &value
}

func nullableSQLInt64(value sql.NullInt64) *int64 {
	if !value.Valid {
		return nil
	}
	result := value.Int64
	return &result
}