package message

import (
	"context"
	"fmt"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
)

func resolveWorkflowRunSources(ctx context.Context, queryer coredb.Queryer, workspaceID string, messageIDs []string) (map[string]ContributionSource, error) {
	messageIDs = normalizeContributionSourceMessageIDs(messageIDs)
	if queryer == nil || strings.TrimSpace(workspaceID) == "" || len(messageIDs) == 0 {
		return map[string]ContributionSource{}, nil
	}
	cte, cteArgs, _, err := coredb.BuildStringValuesCTE("requested_messages", "message_id", messageIDs, 2)
	if err != nil {
		return nil, err
	}
	args := append([]any{strings.TrimSpace(workspaceID)}, cteArgs...)
	rows, err := queryer.Query(ctx, fmt.Sprintf(`
WITH %s
SELECT requested_messages.message_id, COALESCE(effect.id, JSON_UNQUOTE(JSON_EXTRACT(m.metadata, '$.workflow_effect_id'))),
       run.id,
       run.plan_id,
       run.workflow_name_snapshot,
       CASE WHEN effect.effect_kind = 'request_approval' THEN 'approval_feedback' ELSE invocation.source_kind END
FROM requested_messages JOIN scope_messages m ON m.workspace_id = $1 AND m.id = requested_messages.message_id
LEFT JOIN workflow_run_effects effect ON effect.workspace_id = m.workspace_id AND (
       (effect.effect_kind = 'send_message' AND effect.message_id = m.id)
    OR (effect.effect_kind = 'request_approval' AND JSON_UNQUOTE(JSON_EXTRACT(effect.result_json, '$.approvalFeedbackMessageId')) = m.id)
)
JOIN workflow_run_instances run ON run.workspace_id = m.workspace_id AND run.id = COALESCE(effect.run_id, JSON_UNQUOTE(JSON_EXTRACT(m.metadata, '$.workflow_run_id')))
JOIN workflow_invocations invocation ON invocation.workspace_id = run.workspace_id AND invocation.id = run.invocation_id
WHERE effect.id IS NOT NULL
   OR (COALESCE(JSON_UNQUOTE(JSON_EXTRACT(m.metadata, '$.workflow_run_id')), '') <> ''
   AND COALESCE(JSON_UNQUOTE(JSON_EXTRACT(m.metadata, '$.workflow_effect_id')), '') <> '')
ORDER BY effect.id DESC
`, cte), args...)
	if err != nil {
		return nil, fmt.Errorf("load workflow Run contribution sources: %w", err)
	}
	defer rows.Close()
	byMessageID := make(map[string]ContributionSource)
	for rows.Next() {
		var messageID, effectID, runID, planID, workflowName, sourceKind string
		if err := rows.Scan(&messageID, &effectID, &runID, &planID, &workflowName, &sourceKind); err != nil {
			return nil, fmt.Errorf("scan workflow Run contribution source: %w", err)
		}
		messageID = strings.TrimSpace(messageID)
		if messageID == "" {
			continue
		}
		if _, exists := byMessageID[messageID]; exists {
			continue
		}
		byMessageID[messageID] = ContributionSourceForWorkflowRun(WorkflowRunContributionSourceParams{
			RunID: runID, EffectID: effectID, PlanID: planID,
			WorkflowName: workflowName, SourceKind: sourceKind,
		})
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate workflow Run contribution sources: %w", err)
	}
	return byMessageID, nil
}