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

package app

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"golang.org/x/sync/singleflight"

	"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/middleware"
	"github.com/geo-platform/tenant-api/internal/shared/response"
)

type PromptRuleService struct {
	pool       *pgxpool.Pool
	auditLogs  *auditlog.AsyncWriter
	cache      sharedcache.Cache
	cacheGroup singleflight.Group
}

func NewPromptRuleService(pool *pgxpool.Pool, auditLogs *auditlog.AsyncWriter) *PromptRuleService {
	return &PromptRuleService{pool: pool, auditLogs: auditLogs}
}

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

// --- Group types ---

type PromptRuleGroupRequest struct {
	Name      string `json:"name" binding:"required"`
	SortOrder *int   `json:"sort_order"`
}

type PromptRuleGroupResponse struct {
	ID        int64  `json:"id"`
	Name      string `json:"name"`
	SortOrder int    `json:"sort_order"`
	CreatedAt string `json:"created_at"`
}

// --- Rule types ---

type PromptRuleRequest struct {
	GroupID           *int64  `json:"group_id"`
	Name              string  `json:"name" binding:"required"`
	PromptContent     string  `json:"prompt_content" binding:"required"`
	Scene             *string `json:"scene"`
	DefaultTone       *string `json:"default_tone"`
	DefaultWordCount  *int    `json:"default_word_count"`
	KnowledgeGroupIDs []int64 `json:"knowledge_group_ids"`
}

type PromptRuleKnowledgeGroup struct {
	ID   int64  `json:"id"`
	Name string `json:"name"`
}

type PromptRuleResponse struct {
	ID                int64                      `json:"id"`
	GroupID           *int64                     `json:"group_id"`
	Name              string                     `json:"name"`
	PromptContent     string                     `json:"prompt_content"`
	Scene             *string                    `json:"scene"`
	DefaultTone       *string                    `json:"default_tone"`
	DefaultWordCount  *int                       `json:"default_word_count"`
	KnowledgeGroupIDs []int64                    `json:"knowledge_group_ids"`
	KnowledgeGroups   []PromptRuleKnowledgeGroup `json:"knowledge_groups"`
	Status            string                     `json:"status"`
	ArticleCount      int                        `json:"article_count"`
	CreatedAt         string                     `json:"created_at"`
	UpdatedAt         string                     `json:"updated_at"`
}

type PromptRuleListParams struct {
	GroupID   *int64  `json:"group_id"`
	Status    *string `json:"status"`
	Keyword   *string `json:"keyword"`
	Page      int     `json:"page"`
	PageSize  int     `json:"page_size"`
	Ungrouped bool    `json:"ungrouped"`
}

type PromptRuleListResponse struct {
	Items []PromptRuleResponse `json:"items"`
	Total int64                `json:"total"`
}

type PromptRuleStatusRequest struct {
	Status string `json:"status" binding:"required,oneof=enabled disabled"`
}

type PromptRuleSimple struct {
	ID   int64  `json:"id"`
	Name string `json:"name"`
}

// --- Group CRUD ---

func (s *PromptRuleService) ListGroups(ctx context.Context) ([]PromptRuleGroupResponse, error) {
	actor := auth.MustActor(ctx)
	return sharedcache.LoadJSON(ctx, s.cache, &s.cacheGroup, promptRuleGroupsCacheKey(actor.TenantID), defaultCacheTTL(), func(loadCtx context.Context) ([]PromptRuleGroupResponse, error) {
		return s.loadPromptRuleGroups(loadCtx, actor.TenantID)
	})
}

func (s *PromptRuleService) CreateGroup(ctx context.Context, req PromptRuleGroupRequest) (*PromptRuleGroupResponse, error) {
	actor := auth.MustActor(ctx)
	sortOrder := 0
	if req.SortOrder != nil {
		sortOrder = *req.SortOrder
	}
	var id int64
	var ca interface{}
	err := s.pool.QueryRow(ctx, `
		INSERT INTO prompt_rule_groups (tenant_id, name, sort_order)
		VALUES ($1, $2, $3) RETURNING id, created_at
	`, actor.TenantID, req.Name, sortOrder).Scan(&id, &ca)
	if err != nil {
		return nil, response.ErrInternal(50010, "create_failed", "failed to create group")
	}

	afterJSON, _ := json.Marshal(map[string]interface{}{"id": id, "name": req.Name})
	result := "success"
	resourceType := "prompt_rule_group"
	requestID := middleware.RequestIDFromContext(ctx)
	s.auditLogs.Log(auditlog.Entry{
		OperatorID:   actor.UserID,
		TenantID:     &actor.TenantID,
		Module:       "prompt_rule",
		Action:       "create_group",
		ResourceType: &resourceType,
		ResourceID:   &id,
		RequestID:    nilIfEmptyString(requestID),
		AfterJSON:    afterJSON,
		Result:       &result,
	})

	invalidatePromptRuleCaches(ctx, s.cache, actor.TenantID, nil)
	return &PromptRuleGroupResponse{ID: id, Name: req.Name, SortOrder: sortOrder, CreatedAt: fmt.Sprintf("%v", ca)}, nil
}

func (s *PromptRuleService) UpdateGroup(ctx context.Context, id int64, req PromptRuleGroupRequest) error {
	actor := auth.MustActor(ctx)
	sortOrder := 0
	if req.SortOrder != nil {
		sortOrder = *req.SortOrder
	}
	tag, err := s.pool.Exec(ctx, `
		UPDATE prompt_rule_groups SET name = $1, sort_order = $2, updated_at = NOW()
		WHERE id = $3 AND tenant_id = $4 AND deleted_at IS NULL
	`, req.Name, sortOrder, id, actor.TenantID)
	if err != nil || tag.RowsAffected() == 0 {
		return response.ErrNotFound(40430, "group_not_found", "prompt rule group not found")
	}
	invalidatePromptRuleCaches(ctx, s.cache, actor.TenantID, nil)
	return nil
}

func (s *PromptRuleService) DeleteGroup(ctx context.Context, id int64) error {
	actor := auth.MustActor(ctx)

	// Move rules in this group to ungrouped
	_, _ = s.pool.Exec(ctx, `
		UPDATE prompt_rules SET group_id = NULL, updated_at = NOW()
		WHERE group_id = $1 AND tenant_id = $2 AND deleted_at IS NULL
	`, id, actor.TenantID)

	tag, err := s.pool.Exec(ctx, `
		UPDATE prompt_rule_groups SET deleted_at = NOW(), updated_at = NOW()
		WHERE id = $1 AND tenant_id = $2 AND deleted_at IS NULL
	`, id, actor.TenantID)
	if err != nil || tag.RowsAffected() == 0 {
		return response.ErrNotFound(40430, "group_not_found", "prompt rule group not found")
	}

	afterJSON, _ := json.Marshal(map[string]interface{}{"id": id})
	result := "success"
	resourceType := "prompt_rule_group"
	requestID := middleware.RequestIDFromContext(ctx)
	s.auditLogs.Log(auditlog.Entry{
		OperatorID:   actor.UserID,
		TenantID:     &actor.TenantID,
		Module:       "prompt_rule",
		Action:       "delete_group",
		ResourceType: &resourceType,
		ResourceID:   &id,
		RequestID:    nilIfEmptyString(requestID),
		AfterJSON:    afterJSON,
		Result:       &result,
	})
	invalidatePromptRuleCaches(ctx, s.cache, actor.TenantID, nil)
	return nil
}

// --- Rule CRUD ---

func (s *PromptRuleService) List(ctx context.Context, params PromptRuleListParams) (*PromptRuleListResponse, error) {
	actor := auth.MustActor(ctx)
	return sharedcache.LoadJSON(ctx, s.cache, &s.cacheGroup, promptRuleListCacheKey(actor.TenantID, params), defaultCacheTTL(), func(loadCtx context.Context) (*PromptRuleListResponse, error) {
		return s.loadPromptRules(loadCtx, actor.TenantID, params)
	})
}

func (s *PromptRuleService) Detail(ctx context.Context, id int64) (*PromptRuleResponse, error) {
	actor := auth.MustActor(ctx)
	record, found, err := sharedcache.LoadJSONWithEmpty(ctx, s.cache, &s.cacheGroup, promptRuleDetailCacheKey(actor.TenantID, id), defaultCacheTTL(), defaultCacheEmptyTTL(), func(loadCtx context.Context) (*PromptRuleResponse, bool, error) {
		return s.loadPromptRuleDetail(loadCtx, actor.TenantID, id)
	})
	if err != nil {
		return nil, err
	}
	if !found || record == nil {
		return nil, response.ErrNotFound(40431, "rule_not_found", "prompt rule not found")
	}
	return record, nil
}

func (s *PromptRuleService) Create(ctx context.Context, req PromptRuleRequest) (*PromptRuleResponse, error) {
	actor := auth.MustActor(ctx)
	knowledgeGroups, err := s.validatePromptRuleKnowledgeGroups(ctx, actor.TenantID, req.KnowledgeGroupIDs)
	if err != nil {
		return nil, err
	}

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

	var id int64
	var ca interface{}
	err = tx.QueryRow(ctx, `
		INSERT INTO prompt_rules (tenant_id, group_id, name, prompt_content, scene, default_tone, default_word_count)
		VALUES ($1, $2, $3, $4, $5, $6, $7) RETURNING id, created_at
	`, actor.TenantID, req.GroupID, req.Name, req.PromptContent, req.Scene, req.DefaultTone, req.DefaultWordCount).Scan(&id, &ca)
	if err != nil {
		return nil, response.ErrInternal(50010, "create_failed", "failed to create prompt rule")
	}

	if err := s.replacePromptRuleKnowledgeGroups(ctx, tx, id, promptRuleKnowledgeGroupIDs(knowledgeGroups)); err != nil {
		return nil, err
	}
	if err := tx.Commit(ctx); err != nil {
		return nil, response.ErrInternal(50010, "create_failed", "failed to commit prompt rule")
	}

	afterJSON, _ := json.Marshal(map[string]interface{}{"id": id, "name": req.Name})
	result := "success"
	resourceType := "prompt_rule"
	requestID := middleware.RequestIDFromContext(ctx)
	s.auditLogs.Log(auditlog.Entry{
		OperatorID:   actor.UserID,
		TenantID:     &actor.TenantID,
		Module:       "prompt_rule",
		Action:       "create",
		ResourceType: &resourceType,
		ResourceID:   &id,
		RequestID:    nilIfEmptyString(requestID),
		AfterJSON:    afterJSON,
		Result:       &result,
	})

	invalidatePromptRuleCaches(ctx, s.cache, actor.TenantID, &id)
	return &PromptRuleResponse{
		ID:                id,
		GroupID:           req.GroupID,
		Name:              req.Name,
		PromptContent:     req.PromptContent,
		Scene:             req.Scene,
		DefaultTone:       req.DefaultTone,
		DefaultWordCount:  req.DefaultWordCount,
		KnowledgeGroupIDs: promptRuleKnowledgeGroupIDs(knowledgeGroups),
		KnowledgeGroups:   knowledgeGroups,
		Status:            "enabled",
		ArticleCount:      0,
		CreatedAt:         fmt.Sprintf("%v", ca),
		UpdatedAt:         fmt.Sprintf("%v", ca),
	}, nil
}

func (s *PromptRuleService) Update(ctx context.Context, id int64, req PromptRuleRequest) error {
	actor := auth.MustActor(ctx)
	knowledgeGroups, err := s.validatePromptRuleKnowledgeGroups(ctx, actor.TenantID, req.KnowledgeGroupIDs)
	if err != nil {
		return err
	}

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

	tag, err := tx.Exec(ctx, `
		UPDATE prompt_rules SET name = $1, group_id = $2, prompt_content = $3,
		       scene = $4, default_tone = $5, default_word_count = $6, updated_at = NOW()
		WHERE id = $7 AND tenant_id = $8 AND deleted_at IS NULL
	`, req.Name, req.GroupID, req.PromptContent, req.Scene, req.DefaultTone, req.DefaultWordCount, id, actor.TenantID)
	if err != nil || tag.RowsAffected() == 0 {
		return response.ErrNotFound(40431, "rule_not_found", "prompt rule not found")
	}
	if err := s.replacePromptRuleKnowledgeGroups(ctx, tx, id, promptRuleKnowledgeGroupIDs(knowledgeGroups)); err != nil {
		return err
	}
	if err := tx.Commit(ctx); err != nil {
		return response.ErrInternal(50010, "update_failed", "failed to commit prompt rule")
	}
	invalidatePromptRuleCaches(ctx, s.cache, actor.TenantID, &id)
	return nil
}

func (s *PromptRuleService) Delete(ctx context.Context, id int64) error {
	actor := auth.MustActor(ctx)
	tag, err := s.pool.Exec(ctx, `
		UPDATE prompt_rules SET deleted_at = NOW(), updated_at = NOW()
		WHERE id = $1 AND tenant_id = $2 AND deleted_at IS NULL
	`, id, actor.TenantID)
	if err != nil || tag.RowsAffected() == 0 {
		return response.ErrNotFound(40431, "rule_not_found", "prompt rule not found")
	}

	afterJSON, _ := json.Marshal(map[string]interface{}{"id": id})
	result := "success"
	resourceType := "prompt_rule"
	requestID := middleware.RequestIDFromContext(ctx)
	s.auditLogs.Log(auditlog.Entry{
		OperatorID:   actor.UserID,
		TenantID:     &actor.TenantID,
		Module:       "prompt_rule",
		Action:       "delete",
		ResourceType: &resourceType,
		ResourceID:   &id,
		RequestID:    nilIfEmptyString(requestID),
		AfterJSON:    afterJSON,
		Result:       &result,
	})
	invalidatePromptRuleCaches(ctx, s.cache, actor.TenantID, &id)
	return nil
}

func (s *PromptRuleService) ToggleStatus(ctx context.Context, id int64, req PromptRuleStatusRequest) error {
	actor := auth.MustActor(ctx)
	tag, err := s.pool.Exec(ctx, `
		UPDATE prompt_rules SET status = $1, updated_at = NOW()
		WHERE id = $2 AND tenant_id = $3 AND deleted_at IS NULL
	`, req.Status, id, actor.TenantID)
	if err != nil || tag.RowsAffected() == 0 {
		return response.ErrNotFound(40431, "rule_not_found", "prompt rule not found")
	}
	invalidatePromptRuleCaches(ctx, s.cache, actor.TenantID, &id)
	return nil
}

func (s *PromptRuleService) ListSimple(ctx context.Context) ([]PromptRuleSimple, error) {
	actor := auth.MustActor(ctx)
	return sharedcache.LoadJSON(ctx, s.cache, &s.cacheGroup, promptRuleSimpleCacheKey(actor.TenantID), defaultCacheTTL(), func(loadCtx context.Context) ([]PromptRuleSimple, error) {
		return s.loadPromptRuleSimple(loadCtx, actor.TenantID)
	})
}

func (s *PromptRuleService) loadPromptRuleGroups(ctx context.Context, tenantID int64) ([]PromptRuleGroupResponse, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT id, name, sort_order, created_at
		FROM prompt_rule_groups WHERE tenant_id = $1 AND deleted_at IS NULL
		ORDER BY sort_order, created_at
	`, tenantID)
	if err != nil {
		return nil, response.ErrInternal(50010, "query_failed", "failed to list groups")
	}
	defer rows.Close()

	items := make([]PromptRuleGroupResponse, 0)
	for rows.Next() {
		var item PromptRuleGroupResponse
		var createdAt interface{}
		if err := rows.Scan(&item.ID, &item.Name, &item.SortOrder, &createdAt); err != nil {
			return nil, response.ErrInternal(50010, "scan_failed", err.Error())
		}
		item.CreatedAt = fmt.Sprintf("%v", createdAt)
		items = append(items, item)
	}
	return items, nil
}

func (s *PromptRuleService) loadPromptRules(ctx context.Context, tenantID int64, params PromptRuleListParams) (*PromptRuleListResponse, error) {
	if params.Page < 1 {
		params.Page = 1
	}
	if params.PageSize < 1 || params.PageSize > 100 {
		params.PageSize = 20
	}
	offset := (params.Page - 1) * params.PageSize

	var total int64
	var rows interface{ Close() }

	if params.Ungrouped {
		if err := s.pool.QueryRow(ctx, `
			SELECT COUNT(*) FROM prompt_rules
			WHERE tenant_id = $1 AND deleted_at IS NULL AND group_id IS NULL
		`, tenantID).Scan(&total); err != nil {
			return nil, response.ErrInternal(50010, "count_failed", "failed to count rules")
		}

		r, err := s.pool.Query(ctx, `
			SELECT pr.id, pr.group_id, pr.name, pr.prompt_content,
			       pr.scene, pr.default_tone, pr.default_word_count, pr.status,
			       pr.created_at, pr.updated_at,
			       COUNT(a.id)::INT AS article_count
			FROM prompt_rules pr
			LEFT JOIN articles a ON a.prompt_rule_id = pr.id AND a.deleted_at IS NULL
			WHERE pr.tenant_id = $1 AND pr.deleted_at IS NULL AND pr.group_id IS NULL
			GROUP BY pr.id
			ORDER BY pr.created_at DESC
			LIMIT $2 OFFSET $3
		`, tenantID, params.PageSize, offset)
		if err != nil {
			return nil, response.ErrInternal(50010, "query_failed", "failed to list rules")
		}
		rows = r
	} else {
		if err := s.pool.QueryRow(ctx, `
			SELECT COUNT(*) FROM prompt_rules
			WHERE tenant_id = $1 AND deleted_at IS NULL
			  AND ($2::bigint IS NULL OR group_id = $2)
			  AND ($3::text IS NULL OR status = $3)
			  AND ($4::text IS NULL OR name ILIKE '%' || $4 || '%')
		`, tenantID, params.GroupID, params.Status, params.Keyword).Scan(&total); err != nil {
			return nil, response.ErrInternal(50010, "count_failed", "failed to count rules")
		}

		r, err := s.pool.Query(ctx, `
			SELECT pr.id, pr.group_id, pr.name, pr.prompt_content,
			       pr.scene, pr.default_tone, pr.default_word_count, pr.status,
			       pr.created_at, pr.updated_at,
			       COUNT(a.id)::INT AS article_count
			FROM prompt_rules pr
			LEFT JOIN articles a ON a.prompt_rule_id = pr.id AND a.deleted_at IS NULL
			WHERE pr.tenant_id = $1 AND pr.deleted_at IS NULL
			  AND ($2::bigint IS NULL OR pr.group_id = $2)
			  AND ($3::text IS NULL OR pr.status = $3)
			  AND ($4::text IS NULL OR pr.name ILIKE '%' || $4 || '%')
			GROUP BY pr.id
			ORDER BY pr.created_at DESC
			LIMIT $5 OFFSET $6
		`, tenantID, params.GroupID, params.Status, params.Keyword, params.PageSize, offset)
		if err != nil {
			return nil, response.ErrInternal(50010, "query_failed", "failed to list rules")
		}
		rows = r
	}

	pgxRows := rows.(interface {
		Close()
		Next() bool
		Scan(dest ...interface{}) error
		Err() error
	})
	defer pgxRows.Close()

	items := make([]PromptRuleResponse, 0)
	for pgxRows.Next() {
		var item PromptRuleResponse
		var createdAt interface{}
		var updatedAt interface{}
		if err := pgxRows.Scan(&item.ID, &item.GroupID, &item.Name, &item.PromptContent,
			&item.Scene, &item.DefaultTone, &item.DefaultWordCount, &item.Status,
			&createdAt, &updatedAt, &item.ArticleCount); err != nil {
			return nil, response.ErrInternal(50010, "scan_failed", err.Error())
		}
		item.CreatedAt = fmt.Sprintf("%v", createdAt)
		item.UpdatedAt = fmt.Sprintf("%v", updatedAt)
		item.KnowledgeGroupIDs = []int64{}
		item.KnowledgeGroups = []PromptRuleKnowledgeGroup{}
		items = append(items, item)
	}
	groupMap, err := s.loadPromptRuleKnowledgeGroupMap(ctx, tenantID, collectPromptRuleIDs(items))
	if err != nil {
		return nil, err
	}
	for idx := range items {
		items[idx].KnowledgeGroups = groupMap[items[idx].ID]
		items[idx].KnowledgeGroupIDs = promptRuleKnowledgeGroupIDs(items[idx].KnowledgeGroups)
	}

	return &PromptRuleListResponse{Items: items, Total: total}, nil
}

func (s *PromptRuleService) loadPromptRuleDetail(ctx context.Context, tenantID, ruleID int64) (*PromptRuleResponse, bool, error) {
	var item PromptRuleResponse
	var createdAt interface{}
	var updatedAt interface{}
	var articleCount int
	err := s.pool.QueryRow(ctx, `
		SELECT pr.id, pr.group_id, pr.name, pr.prompt_content,
		       pr.scene, pr.default_tone, pr.default_word_count, pr.status,
		       pr.created_at, pr.updated_at,
		       (SELECT COUNT(*) FROM articles WHERE prompt_rule_id = pr.id AND deleted_at IS NULL)::INT
		FROM prompt_rules pr
		WHERE pr.id = $1 AND pr.tenant_id = $2 AND pr.deleted_at IS NULL
	`, ruleID, tenantID).Scan(&item.ID, &item.GroupID, &item.Name, &item.PromptContent,
		&item.Scene, &item.DefaultTone, &item.DefaultWordCount, &item.Status,
		&createdAt, &updatedAt, &articleCount)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, false, nil
		}
		return nil, false, response.ErrInternal(50010, "query_failed", "failed to fetch prompt rule")
	}
	item.CreatedAt = fmt.Sprintf("%v", createdAt)
	item.UpdatedAt = fmt.Sprintf("%v", updatedAt)
	item.ArticleCount = articleCount

	groupMap, err := s.loadPromptRuleKnowledgeGroupMap(ctx, tenantID, []int64{ruleID})
	if err != nil {
		return nil, false, err
	}
	item.KnowledgeGroups = groupMap[ruleID]
	item.KnowledgeGroupIDs = promptRuleKnowledgeGroupIDs(item.KnowledgeGroups)
	return &item, true, nil
}

func (s *PromptRuleService) loadPromptRuleSimple(ctx context.Context, tenantID int64) ([]PromptRuleSimple, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT id, name FROM prompt_rules
		WHERE tenant_id = $1 AND deleted_at IS NULL AND status = 'enabled'
		ORDER BY name
	`, tenantID)
	if err != nil {
		return nil, response.ErrInternal(50010, "query_failed", "failed to list rules")
	}
	defer rows.Close()

	items := make([]PromptRuleSimple, 0)
	for rows.Next() {
		var item PromptRuleSimple
		if err := rows.Scan(&item.ID, &item.Name); err != nil {
			return nil, response.ErrInternal(50010, "scan_failed", err.Error())
		}
		items = append(items, item)
	}
	return items, nil
}

func (s *PromptRuleService) validatePromptRuleKnowledgeGroups(
	ctx context.Context,
	tenantID int64,
	groupIDs []int64,
) ([]PromptRuleKnowledgeGroup, error) {
	groupIDs = normalizeInt64IDs(groupIDs)
	if len(groupIDs) == 0 {
		return []PromptRuleKnowledgeGroup{}, nil
	}

	rows, err := s.pool.Query(ctx, `
		SELECT id, name
		FROM knowledge_groups
		WHERE tenant_id = $1 AND id = ANY($2::bigint[]) AND deleted_at IS NULL
		ORDER BY sort_order, created_at, id
	`, tenantID, groupIDs)
	if err != nil {
		return nil, response.ErrInternal(50010, "knowledge_group_query_failed", "failed to query knowledge groups")
	}
	defer rows.Close()

	groups := make([]PromptRuleKnowledgeGroup, 0, len(groupIDs))
	found := make(map[int64]struct{}, len(groupIDs))
	for rows.Next() {
		var item PromptRuleKnowledgeGroup
		if err := rows.Scan(&item.ID, &item.Name); err != nil {
			return nil, response.ErrInternal(50010, "knowledge_group_query_failed", err.Error())
		}
		groups = append(groups, item)
		found[item.ID] = struct{}{}
	}

	if len(found) != len(groupIDs) {
		return nil, response.ErrNotFound(40451, "knowledge_group_not_found", "knowledge group not found")
	}

	return groups, nil
}

func (s *PromptRuleService) loadPromptRuleKnowledgeGroupMap(
	ctx context.Context,
	tenantID int64,
	ruleIDs []int64,
) (map[int64][]PromptRuleKnowledgeGroup, error) {
	result := make(map[int64][]PromptRuleKnowledgeGroup, len(ruleIDs))
	ruleIDs = normalizeInt64IDs(ruleIDs)
	if len(ruleIDs) == 0 {
		return result, nil
	}

	rows, err := s.pool.Query(ctx, `
		SELECT prkg.prompt_rule_id, kg.id, kg.name
		FROM prompt_rule_knowledge_groups prkg
		JOIN knowledge_groups kg ON kg.id = prkg.knowledge_group_id
		WHERE prkg.prompt_rule_id = ANY($1::bigint[])
		  AND kg.tenant_id = $2
		  AND kg.deleted_at IS NULL
		ORDER BY kg.sort_order, kg.created_at, kg.id
	`, ruleIDs, tenantID)
	if err != nil {
		return nil, response.ErrInternal(50010, "knowledge_group_query_failed", "failed to query prompt rule knowledge groups")
	}
	defer rows.Close()

	for rows.Next() {
		var (
			ruleID int64
			group  PromptRuleKnowledgeGroup
		)
		if err := rows.Scan(&ruleID, &group.ID, &group.Name); err != nil {
			return nil, response.ErrInternal(50010, "knowledge_group_query_failed", err.Error())
		}
		result[ruleID] = append(result[ruleID], group)
	}

	return result, nil
}

func (s *PromptRuleService) replacePromptRuleKnowledgeGroups(
	ctx context.Context,
	tx pgx.Tx,
	promptRuleID int64,
	groupIDs []int64,
) error {
	if _, err := tx.Exec(ctx, `
		DELETE FROM prompt_rule_knowledge_groups
		WHERE prompt_rule_id = $1
	`, promptRuleID); err != nil {
		return response.ErrInternal(50010, "update_failed", "failed to update prompt rule knowledge groups")
	}

	for _, groupID := range normalizeInt64IDs(groupIDs) {
		if _, err := tx.Exec(ctx, `
			INSERT INTO prompt_rule_knowledge_groups (prompt_rule_id, knowledge_group_id)
			VALUES ($1, $2)
		`, promptRuleID, groupID); err != nil {
			return response.ErrInternal(50010, "update_failed", "failed to update prompt rule knowledge groups")
		}
	}

	return nil
}

func collectPromptRuleIDs(items []PromptRuleResponse) []int64 {
	ids := make([]int64, 0, len(items))
	for _, item := range items {
		ids = append(ids, item.ID)
	}
	return normalizeInt64IDs(ids)
}

func promptRuleKnowledgeGroupIDs(groups []PromptRuleKnowledgeGroup) []int64 {
	ids := make([]int64, 0, len(groups))
	for _, group := range groups {
		ids = append(ids, group.ID)
	}
	return normalizeInt64IDs(ids)
}