package calendar_test

import (
	"context"
	"errors"
	"fmt"
	"regexp"
	"testing"
	"time"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/store/workflow/calendar"
)

const scheduleSourcePageQueryPattern = `(?s)JSON_REMOVE\(JSON_EXTRACT\(revision\.config_json, '\$\.source'\), '\$\.schedule\.inputs'\).*FROM workflow_binding_sources source.*source\.next_occurrence_at > \?.*source\.next_occurrence_at = \?.*source\.id > \?.*ORDER BY source\.next_occurrence_at ASC, source\.id ASC.*LIMIT \?`

type transactionOptionsRecorder struct {
	coredb.DB
	options coredb.TxOptions
}

func (db *transactionOptionsRecorder) BeginTx(ctx context.Context, options coredb.TxOptions) (coredb.Tx, error) {
	db.options = options
	return db.DB.BeginTx(ctx, options)
}

func TestVisitSnapshotUsesOneTiDBCompatibleRepeatableReadSnapshot(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	t.Cleanup(func() { _ = sqlDB.Close() })
	database := &transactionOptionsRecorder{DB: coredb.WrapSQLDB(sqlDB)}
	store := calendar.New(database)
	from := time.Date(2026, 8, 17, 16, 0, 0, 0, time.UTC)
	to := from.Add(72 * time.Hour)

	mock.ExpectBegin()
	mock.ExpectQuery(`(?s)FROM workflow_invocations i.*JOIN workflow_run_instances r.*WHERE i\.workspace_id = \?.*i\.source_occurred_at >= \?.*i\.source_occurred_at < \?.*ORDER BY i\.source_occurred_at ASC, i\.id ASC.*LIMIT \?`).
		WithArgs("wsp_1", from, to, 5001).
		WillReturnRows(sqlmock.NewRows([]string{
			"invocation_id", "run_id", "plan_id", "plan_version_id", "workflow_name",
			"plan_key", "plan_display_name", "binding_id", "binding_key", "binding_display_name",
			"binding_revision_id", "binding_source_id", "source_kind", "occurrence_key", "occurs_at",
			"status", "queued_at", "started_at", "completed_at", "canceled_at",
		}))
	cursorStart := time.Date(1000, 1, 1, 0, 0, 0, 0, time.UTC)
	mock.ExpectQuery(scheduleSourcePageQueryPattern).
		WithArgs("wsp_1", to, cursorStart, cursorStart, "", 256).
		WillReturnRows(sqlmock.NewRows([]string{
			"source_id", "binding_id", "revision_id", "activation", "next_occurrence_at",
			"plan_id", "plan_version_id", "source_json", "binding_key", "binding_name", "plan_key", "plan_name",
		}))
	mock.ExpectCommit()

	var runs []calendar.Run
	var sources []calendar.ScheduleSource
	err = store.VisitSnapshot(context.Background(), calendar.ListSnapshotParams{
		WorkspaceID: " wsp_1 ", From: from, To: to, RunLimit: 5001,
	}, calendar.SnapshotVisitor{
		Runs: func(values []calendar.Run) error {
			runs = append(runs, values...)
			return nil
		},
		Source: func(value calendar.ScheduleSource) error {
			sources = append(sources, value)
			return nil
		},
	})
	require.NoError(t, err)
	require.Empty(t, runs)
	require.Empty(t, sources)
	require.Equal(t, coredb.TxOptions{IsoLevel: coredb.RepeatableRead}, database.options)
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestVisitSnapshotRunQueryDoesNotUseRecentRunOrderingOrImplicitLimit(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	t.Cleanup(func() { _ = sqlDB.Close() })
	store := calendar.New(coredb.WrapSQLDB(sqlDB))
	from := time.Date(2026, 8, 18, 0, 0, 0, 0, time.UTC)
	to := from.Add(time.Hour)
	wantErr := context.Canceled

	mock.ExpectBegin()
	mock.ExpectQuery(regexp.QuoteMeta("ORDER BY i.source_occurred_at ASC, i.id ASC\nLIMIT ?")).
		WithArgs("wsp_1", from, to, 137).
		WillReturnError(wantErr)
	mock.ExpectRollback()

	err = store.VisitSnapshot(context.Background(), calendar.ListSnapshotParams{
		WorkspaceID: "wsp_1", From: from, To: to, RunLimit: 137,
	}, calendar.SnapshotVisitor{
		Runs:   func([]calendar.Run) error { return nil },
		Source: func(calendar.ScheduleSource) error { return nil },
	})
	require.ErrorIs(t, err, wantErr)
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestVisitSnapshotStopsScanningSourcesWhenVisitorReturnsError(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	t.Cleanup(func() { _ = sqlDB.Close() })
	store := calendar.New(coredb.WrapSQLDB(sqlDB))
	from := time.Date(2026, 8, 18, 0, 0, 0, 0, time.UTC)
	to := from.Add(time.Hour)
	stopErr := errors.New("item limit reached")

	mock.ExpectBegin()
	mock.ExpectQuery(`(?s)FROM workflow_invocations i.*LIMIT \?`).
		WithArgs("wsp_1", from, to, 5001).
		WillReturnRows(sqlmock.NewRows([]string{
			"invocation_id", "run_id", "plan_id", "plan_version_id", "workflow_name",
			"plan_key", "plan_display_name", "binding_id", "binding_key", "binding_display_name",
			"binding_revision_id", "binding_source_id", "source_kind", "occurrence_key", "occurs_at",
			"status", "queued_at", "started_at", "completed_at", "canceled_at",
		}))
	cursorStart := time.Date(1000, 1, 1, 0, 0, 0, 0, time.UTC)
	mock.ExpectQuery(scheduleSourcePageQueryPattern).
		WithArgs("wsp_1", to, cursorStart, cursorStart, "", 256).
		WillReturnRows(sqlmock.NewRows([]string{
			"source_id", "binding_id", "revision_id", "activation", "next_occurrence_at",
			"plan_id", "plan_version_id", "source_json", "binding_key", "binding_name", "plan_key", "plan_name",
		}).
			AddRow("wbs_1", "wbd_1", "wbr_1", 1, from, "pln_1", "plv_1",
				[]byte(`{"type":"schedule","schedule":{"mode":"absolute","scheduledAt":"2026-08-18T00:00:00Z"}}`),
				"daily", "Daily", "daily", "Daily").
			AddRow("wbs_2", "wbd_2", "wbr_2", 1, from, "pln_2", "plv_2",
				[]byte(`{"type":"schedule","schedule":{"mode":"absolute","scheduledAt":"2026-08-18T00:00:00Z"}}`),
				"weekly", "Weekly", "weekly", "Weekly"))
	mock.ExpectRollback()

	visited := 0
	err = store.VisitSnapshot(context.Background(), calendar.ListSnapshotParams{
		WorkspaceID: "wsp_1", From: from, To: to, RunLimit: 5001,
	}, calendar.SnapshotVisitor{
		Runs: func([]calendar.Run) error { return nil },
		Source: func(calendar.ScheduleSource) error {
			visited++
			return stopErr
		},
	})
	require.ErrorIs(t, err, stopErr)
	require.Equal(t, 1, visited)
	require.NoError(t, mock.ExpectationsWereMet())
}

func TestVisitSnapshotPagesSourcesWithinTheSameTransaction(t *testing.T) {
	sqlDB, mock, err := sqlmock.New()
	require.NoError(t, err)
	t.Cleanup(func() { _ = sqlDB.Close() })
	store := calendar.New(coredb.WrapSQLDB(sqlDB))
	from := time.Date(2026, 8, 18, 0, 0, 0, 0, time.UTC)
	to := from.Add(time.Hour)
	sourceColumns := []string{
		"source_id", "binding_id", "revision_id", "activation", "next_occurrence_at",
		"plan_id", "plan_version_id", "source_json", "binding_key", "binding_name", "plan_key", "plan_name",
	}
	sourceJSON := []byte(`{"type":"schedule","schedule":{"mode":"absolute","scheduledAt":"2026-08-18T00:00:00Z"}}`)

	mock.ExpectBegin()
	mock.ExpectQuery(`(?s)FROM workflow_invocations i.*LIMIT \?`).
		WithArgs("wsp_1", from, to, 5001).
		WillReturnRows(sqlmock.NewRows([]string{
			"invocation_id", "run_id", "plan_id", "plan_version_id", "workflow_name",
			"plan_key", "plan_display_name", "binding_id", "binding_key", "binding_display_name",
			"binding_revision_id", "binding_source_id", "source_kind", "occurrence_key", "occurs_at",
			"status", "queued_at", "started_at", "completed_at", "canceled_at",
		}))
	firstPage := sqlmock.NewRows(sourceColumns)
	for index := 0; index < 256; index++ {
		firstPage.AddRow(
			fmt.Sprintf("wbs_%03d", index), "wbd_1", "wbr_1", 1, from.Add(time.Duration(index)*time.Microsecond),
			"pln_1", "plv_1", sourceJSON, "daily", "Daily", "daily", "Daily",
		)
	}
	cursorStart := time.Date(1000, 1, 1, 0, 0, 0, 0, time.UTC)
	mock.ExpectQuery(scheduleSourcePageQueryPattern).
		WithArgs("wsp_1", to, cursorStart, cursorStart, "", 256).
		WillReturnRows(firstPage)
	lastAt := from.Add(255 * time.Microsecond)
	mock.ExpectQuery(scheduleSourcePageQueryPattern).
		WithArgs("wsp_1", to, lastAt, lastAt, "wbs_255", 256).
		WillReturnRows(sqlmock.NewRows(sourceColumns).
			AddRow("wbs_256", "wbd_1", "wbr_1", 1, from.Add(256*time.Microsecond),
				"pln_1", "plv_1", sourceJSON, "daily", "Daily", "daily", "Daily"))
	mock.ExpectCommit()

	visited := 0
	err = store.VisitSnapshot(context.Background(), calendar.ListSnapshotParams{
		WorkspaceID: "wsp_1", From: from, To: to, RunLimit: 5001,
	}, calendar.SnapshotVisitor{
		Runs: func([]calendar.Run) error { return nil },
		Source: func(calendar.ScheduleSource) error {
			visited++
			return nil
		},
	})
	require.NoError(t, err)
	require.Equal(t, 257, visited)
	require.NoError(t, mock.ExpectationsWereMet())
}