package main

import (
	"context"
	"database/sql"
	"fmt"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

var (
	messageSearchCommandPackageFixtures = map[testkit.TiDBTestRuntime]*messageSearchCommandPackageFixtureState{
		testkit.TiDBTestRuntimeFullCluster: {},
		testkit.TiDBTestRuntimeUnistore:    {},
	}
)

type messageSearchCommandPackageFixtureState struct {
	once    sync.Once
	dsn     string
	db      *sql.DB
	cleanup func(context.Context) error
	err     error
}

func TestMain(m *testing.M) {
	os.Exit(runMessageSearchCommandPackageTests(m))
}

func TestStartMessageSearchCommandTestDBReusesPackageFixtureAndResetsState(t *testing.T) {
	ctx, firstDSN, firstDB := startMessageSearchCommandTestDB(t, testkit.TiDBTestRuntimeUnistore)
	seedScopeMessageSearchCommandMessages(t, ctx, firstDB)

	var initialMessages int
	require.NoError(t, firstDB.QueryRowContext(ctx, `
SELECT COUNT(*)
FROM scope_messages
WHERE workspace_id = 'wsp_scope_cmd_backfill'
`).Scan(&initialMessages))
	require.Equal(t, 2, initialMessages)

	ctx, secondDSN, secondDB := startMessageSearchCommandTestDB(t, testkit.TiDBTestRuntimeUnistore)
	require.Equal(t, firstDSN, secondDSN)
	require.Same(t, firstDB, secondDB)

	var resetMessages int
	require.NoError(t, secondDB.QueryRowContext(ctx, `
SELECT COUNT(*)
FROM scope_messages
WHERE workspace_id = 'wsp_scope_cmd_backfill'
`).Scan(&resetMessages))
	require.Zero(t, resetMessages)

	var stateCount int
	require.NoError(t, secondDB.QueryRowContext(ctx, `
SELECT COUNT(*)
FROM scope_message_search_state
WHERE workspace_id = 'wsp_scope_cmd_backfill'
`).Scan(&stateCount))
	require.Zero(t, stateCount)
}

func runMessageSearchCommandPackageTests(m *testing.M) int {
	code := m.Run()

	for _, runtime := range []testkit.TiDBTestRuntime{
		testkit.TiDBTestRuntimeFullCluster,
		testkit.TiDBTestRuntimeUnistore,
	} {
		fixture := messageSearchCommandPackageFixtures[runtime]
		if fixture == nil || fixture.cleanup == nil {
			continue
		}

		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
		err := fixture.cleanup(ctx)
		cancel()
		if err != nil {
			fmt.Fprintf(os.Stderr, "cleanup cmd/agentteam-tidb-admin %s fixture: %v\n", runtime, err)
			if code == 0 {
				code = 1
			}
		}
	}
	return code
}

func startMessageSearchCommandTestDB(t *testing.T, runtime testkit.TiDBTestRuntime) (context.Context, string, *sql.DB) {
	t.Helper()

	ctx := context.Background()
	dsn, db, _ := sharedMessageSearchCommandTestDB(t, runtime)
	resetMessageSearchCommandTestDB(t, ctx, db)
	return ctx, dsn, db
}

func sharedMessageSearchCommandTestDB(t *testing.T, runtime testkit.TiDBTestRuntime) (string, *sql.DB, *messageSearchCommandPackageFixtureState) {
	t.Helper()

	fixture := messageSearchCommandPackageFixtures[runtime]
	require.NotNil(t, fixture, "unsupported cmd/agentteam-tidb-admin test runtime: %s", runtime)

	fixture.once.Do(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
		defer cancel()

		dsn, cleanupDatabase, err := startMessageSearchCommandInitializedFixture(ctx, runtime)
		if err != nil {
			fixture.err = err
			return
		}

		db, err := testkit.OpenTiDBContext(ctx, dsn)
		if err != nil {
			cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 2*time.Minute)
			defer cleanupCancel()
			_ = cleanupDatabase(cleanupCtx)
			fixture.err = err
			return
		}

		fixture.dsn = dsn
		fixture.db = db
		fixture.cleanup = func(cleanupCtx context.Context) error {
			if cleanupCtx == nil {
				cleanupCtx = context.Background()
			}

			var cleanupErr error
			if err := db.Close(); err != nil {
				cleanupErr = err
			}
			if err := cleanupDatabase(cleanupCtx); cleanupErr == nil && err != nil {
				cleanupErr = err
			}
			return cleanupErr
		}
	})

	if reason, ok := testkit.SkipReason(fixture.err); ok {
		t.Skip(reason)
	}
	require.NoError(t, fixture.err)
	require.NotEmpty(t, fixture.dsn)
	require.NotNil(t, fixture.db)
	return fixture.dsn, fixture.db, fixture
}

func startMessageSearchCommandInitializedFixture(ctx context.Context, runtime testkit.TiDBTestRuntime) (string, func(context.Context) error, error) {
	if ctx == nil {
		ctx = context.Background()
	}

	dsn, cleanup, err := testkit.StartInitializedTiDBWithRuntimeAndCleanup(ctx, runtime)
	if err != nil {
		return "", nil, err
	}
	if dsn == "" || cleanup == nil {
		return "", nil, fmt.Errorf("start cmd/agentteam-tidb-admin %s fixture: empty dsn or cleanup", runtime)
	}
	return dsn, cleanup, nil
}

func resetMessageSearchCommandTestDB(t *testing.T, ctx context.Context, db *sql.DB) {
	t.Helper()

	statements := []string{
		`DELETE FROM scope_message_search_docs
WHERE workspace_id = 'wsp_scope_cmd_backfill'`,
		`DELETE FROM scope_message_search_state
WHERE workspace_id = 'wsp_scope_cmd_backfill'`,
		`DELETE FROM scope_messages
WHERE workspace_id = 'wsp_scope_cmd_backfill'`,
		`DELETE FROM scope_memberships
WHERE scope_id = 'scp_scope_cmd_backfill' OR member_id = 'mem_scope_cmd_backfill'`,
		`DELETE FROM scopes
WHERE id = 'scp_scope_cmd_backfill'`,
		`DELETE FROM members
WHERE id = 'mem_scope_cmd_backfill'`,
		`DELETE FROM workspaces
WHERE id = 'wsp_scope_cmd_backfill'`,
		`DELETE FROM organizations
WHERE id = 'org_scope_cmd_backfill'`,
		`DELETE FROM users
WHERE id = 'usr_scope_cmd_backfill'`,
	}

	for _, statement := range statements {
		_, err := db.ExecContext(ctx, statement)
		require.NoError(t, err, statement)
	}
}