package serviceactor

import (
	"context"
	"database/sql"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

func TestReadyDeliveryClaimQueriesAvoidCoalesceReadyOr(t *testing.T) {
	t.Parallel()

	noRetryQuery := readyDeliveryClaimNoRetrySQL()
	dueRetryQuery := readyDeliveryClaimDueRetrySQL()

	for _, query := range []string{noRetryQuery, dueRetryQuery} {
		require.NotContains(t, query, "COALESCE(delivery.next_retry_at")
		require.NotContains(t, query, "delivery.next_retry_at IS NULL OR")
		require.NotContains(t, query, "FOR UPDATE")
		require.Contains(t, query, "SELECT delivery.id, delivery.created_at, delivery.next_retry_at")
		require.Contains(t, query, "NOT EXISTS")
		require.Contains(t, query, "delivery.delivery_policy = $4")
		require.Contains(t, query, "delivery.service_connection_ref = $5")
		require.NotContains(t, query, "attempt.lease_expires_at")
	}
	require.Contains(t, noRetryQuery, "delivery.next_retry_at IS NULL")
	require.Contains(t, noRetryQuery, "ORDER BY delivery.created_at, delivery.id")
	require.Contains(t, dueRetryQuery, "delivery.next_retry_at IS NOT NULL")
	require.Contains(t, dueRetryQuery, "delivery.next_retry_at <= $6")
	require.Contains(t, dueRetryQuery, "ORDER BY delivery.next_retry_at, delivery.id")

	noRetryLockQuery := readyDeliveryClaimLockSQL("delivery.next_retry_at IS NULL", "$7")
	dueRetryLockQuery := readyDeliveryClaimLockSQL(`delivery.next_retry_at IS NOT NULL
  AND delivery.next_retry_at <= $6`, "$7, $8")
	for _, query := range []string{noRetryLockQuery, dueRetryLockQuery} {
		require.NotContains(t, query, "COALESCE(delivery.next_retry_at")
		require.NotContains(t, query, "delivery.next_retry_at IS NULL OR")
		require.Contains(t, query, "delivery.id IN")
		require.Contains(t, query, "FOR UPDATE")
		require.Contains(t, query, "NOT EXISTS")
	}
}

func TestMergeReadyDeliveryClaimCandidatesPreservesReadyOrderAndLimit(t *testing.T) {
	t.Parallel()

	now := time.Date(2026, 6, 5, 11, 0, 0, 0, time.UTC)
	retryAt := now.Add(2 * time.Second)
	tieReadyAt := now.Add(4 * time.Second)
	noRetryCandidates := []readyDeliveryClaimCandidate{
		{ID: "del_no_retry_1", CreatedAt: now.Add(time.Second)},
		{ID: "del_z_no_retry_tie", CreatedAt: tieReadyAt},
	}
	dueRetryCandidates := []readyDeliveryClaimCandidate{
		{ID: "del_retry_1", CreatedAt: now.Add(-time.Minute), NextRetryAt: &retryAt},
		{ID: "del_a_retry_tie", CreatedAt: now.Add(-time.Minute), NextRetryAt: &tieReadyAt},
	}

	merged := mergeReadyDeliveryClaimCandidates(noRetryCandidates, dueRetryCandidates, 3)

	gotIDs := make([]string, 0, len(merged))
	for _, candidate := range merged {
		gotIDs = append(gotIDs, candidate.ID)
	}
	require.Equal(t, []string{"del_no_retry_1", "del_retry_1", "del_a_retry_tie"}, gotIDs)
}

func TestListReadyDeliveryClaimsLocksOnlyMergedCandidates(t *testing.T) {
	t.Parallel()

	now := time.Date(2026, 6, 5, 11, 0, 0, 0, time.UTC)
	retryAt1 := now.Add(2 * time.Second)
	retryAt2 := now.Add(3 * time.Second)
	retryAt3 := now.Add(6 * time.Second)
	db := &readyClaimCapture{
		rows: []coredb.Rows{
			&readyClaimCandidateRows{candidates: []readyDeliveryClaimCandidate{
				{ID: " del_no_retry_1 ", CreatedAt: now.Add(time.Second)},
				{ID: "del_no_retry_2", CreatedAt: now.Add(4 * time.Second)},
				{ID: "del_no_retry_3", CreatedAt: now.Add(5 * time.Second)},
			}},
			&readyClaimCandidateRows{candidates: []readyDeliveryClaimCandidate{
				{ID: "del_retry_1", CreatedAt: now.Add(-time.Minute), NextRetryAt: &retryAt1},
				{ID: "del_retry_2", CreatedAt: now.Add(-time.Minute), NextRetryAt: &retryAt2},
				{ID: "del_retry_3", CreatedAt: now.Add(-time.Minute), NextRetryAt: &retryAt3},
			}},
			&readyClaimCandidateRows{},
			&readyClaimCandidateRows{},
		},
	}

	deliveries, err := listReadyDeliveryClaimsForUpdate(context.Background(), db, deliveryReadyClaimQuery{
		WorkspaceID: " wsp_1 ",
		Now:         now,
		Limit:       3,
	})

	require.NoError(t, err)
	require.Empty(t, deliveries)
	require.Len(t, db.queries, 4)
	require.NotContains(t, strings.Join(strings.Fields(db.queries[0]), " "), "FOR UPDATE")
	require.NotContains(t, strings.Join(strings.Fields(db.queries[1]), " "), "FOR UPDATE")
	noRetryLockQuery := strings.Join(strings.Fields(db.queries[2]), " ")
	require.Contains(t, noRetryLockQuery, "delivery.id IN ($7)")
	require.Contains(t, noRetryLockQuery, "FOR UPDATE")
	require.Equal(t, []any{"wsp_1", "", "", "", "", now, "del_no_retry_1"}, db.args[2])
	dueRetryLockQuery := strings.Join(strings.Fields(db.queries[3]), " ")
	require.Contains(t, dueRetryLockQuery, "delivery.id IN ($7, $8)")
	require.Contains(t, dueRetryLockQuery, "FOR UPDATE")
	require.Equal(t, []any{"wsp_1", "", "", "", "", now, "del_retry_1", "del_retry_2"}, db.args[3])
}

type readyClaimCapture struct {
	queries []string
	args    [][]any
	rows    []coredb.Rows
}

func (db *readyClaimCapture) Query(_ context.Context, query string, args ...any) (coredb.Rows, error) {
	index := len(db.queries)
	db.queries = append(db.queries, query)
	db.args = append(db.args, append([]any(nil), args...))
	if index >= len(db.rows) {
		return &readyClaimCandidateRows{}, nil
	}
	return db.rows[index], nil
}

func (db *readyClaimCapture) QueryRow(context.Context, string, ...any) coredb.Row {
	return readyClaimRow{err: sql.ErrNoRows}
}

type readyClaimCandidateRows struct {
	candidates []readyDeliveryClaimCandidate
	index      int
}

func (rows *readyClaimCandidateRows) Close() {}

func (rows *readyClaimCandidateRows) Err() error { return nil }

func (rows *readyClaimCandidateRows) Next() bool {
	return rows.index < len(rows.candidates)
}

func (rows *readyClaimCandidateRows) Scan(dest ...any) error {
	if rows.index >= len(rows.candidates) || len(dest) != 3 {
		return sql.ErrNoRows
	}
	candidate := rows.candidates[rows.index]
	rows.index++
	id, ok := dest[0].(*string)
	if !ok {
		return sql.ErrNoRows
	}
	createdAt, ok := dest[1].(*time.Time)
	if !ok {
		return sql.ErrNoRows
	}
	nextRetryAt, ok := dest[2].(*sql.NullTime)
	if !ok {
		return sql.ErrNoRows
	}
	*id = candidate.ID
	*createdAt = candidate.CreatedAt
	if candidate.NextRetryAt != nil {
		*nextRetryAt = sql.NullTime{Time: *candidate.NextRetryAt, Valid: true}
	} else {
		*nextRetryAt = sql.NullTime{}
	}
	return nil
}

type readyClaimRow struct {
	err error
}

func (row readyClaimRow) Scan(...any) error {
	if row.err != nil {
		return row.err
	}
	return nil
}
