package database

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

	"github.com/go-sql-driver/mysql"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/AgentTeam/internal/store/automation"
	"github.com/tidbcloud/AgentTeam/internal/store/repocache"
	"github.com/tidbcloud/AgentTeam/internal/store/resource/activity"
	"gorm.io/gorm"
)

func TestTiDBPublicRepositoryCacheMigration(t *testing.T) {
	db := openTiDBTestDB(t, func(db *gorm.DB) error {
		if err := db.AutoMigrate(&publicRepositoryCacheBeforeGenerations{}); err != nil {
			return err
		}
		legacy := legacyPublicRepositoryCacheFixture(time.Now().UTC())
		if err := db.Create(&legacy).Error; err != nil {
			return err
		}
		return Migrate(db)
	})
	require.True(t, db.Migrator().HasTable(&repocache.Cache{}))
	require.True(t, db.Migrator().HasIndex(&repocache.Cache{}, "idx_public_repository_caches_normalized_remote_url"))
	require.True(t, db.Migrator().HasColumn(&repocache.Cache{}, "Info"))
	require.True(t, db.Migrator().HasColumn(&repocache.Cache{}, "ActiveGeneration"))
	require.True(t, db.Migrator().HasColumn(&repocache.Cache{}, "Revision"))
	var migrated repocache.Cache
	require.NoError(t, db.First(&migrated, "id = ?", "prc_legacy").Error)
	require.Empty(t, migrated.ActiveGeneration)
	require.Equal(t, uint64(0), migrated.Revision)
	require.Equal(t, "legacy-commit", migrated.LastCommit)

	repository := repocache.New(db)
	updateTime := time.Now().UTC().Truncate(time.Second)
	first, err := repository.UpdateSettings(context.Background(), "prc_legacy", false, 3600, "ops", updateTime)
	require.NoError(t, err)
	require.Equal(t, uint64(1), first.Revision)
	second, err := repository.UpdateSettings(context.Background(), "prc_legacy", false, 3600, "ops", updateTime)
	require.NoError(t, err)
	require.Equal(t, uint64(2), second.Revision)

	columns, err := db.Migrator().ColumnTypes(&repocache.Cache{})
	require.NoError(t, err)
	for _, column := range columns {
		if !strings.EqualFold(column.Name(), "info") {
			continue
		}
		nullable, ok := column.Nullable()
		require.True(t, ok)
		require.False(t, nullable)
		_, hasDefault := column.DefaultValue()
		require.False(t, hasDefault)
		return
	}
	t.Fatal("public_repository_caches.info column not found")
}

func TestTiDBLatestCommitProjectionConcurrentFirstWrite(t *testing.T) {
	db := openTiDBTestDB(t, func(db *gorm.DB) error {
		return db.AutoMigrate(&activity.ResourceActivityEvent{}, &activity.ResourceActivityLatestCommit{})
	})
	store := activity.New(db)
	ctx := context.Background()
	base := time.Date(2026, 7, 31, 10, 0, 0, 0, time.UTC)
	events := []activity.ResourceActivityEvent{
		{ID: "act_old", WorkspaceID: "wsp_1", AssetType: activity.AssetTypeKB, AssetID: "src_1", EventType: activity.EventTypeCommit, OccurredAt: base},
		{ID: "act_new", WorkspaceID: "wsp_1", AssetType: activity.AssetTypeKB, AssetID: "src_1", EventType: activity.EventTypeCommit, OccurredAt: base.Add(time.Hour)},
	}
	start := make(chan struct{})
	errs := make(chan error, len(events))
	var wg sync.WaitGroup
	for _, event := range events {
		wg.Add(1)
		go func(event activity.ResourceActivityEvent) {
			defer wg.Done()
			<-start
			_, err := store.UpsertEvent(ctx, event)
			errs <- err
		}(event)
	}
	close(start)
	wg.Wait()
	close(errs)
	for err := range errs {
		require.NoError(t, err)
	}
	var latest activity.ResourceActivityLatestCommit
	require.NoError(t, db.First(&latest, "workspace_id = ? AND asset_type = ? AND asset_id = ?", "wsp_1", activity.AssetTypeKB, "src_1").Error)
	require.Equal(t, "act_new", latest.EventID)
	require.Equal(t, base.Add(time.Hour), latest.OccurredAt)
}

func TestTiDBAutomationQueueLeaseSemantics(t *testing.T) {
	db := openTiDBAutomationTestDB(t)
	store := automation.New(db)
	ctx := context.Background()

	t.Run("concurrent claim does not duplicate lease", func(t *testing.T) {
		now := time.Date(2026, 6, 13, 10, 0, 0, 0, time.UTC)
		workspaceID := "ws-tidb-concurrent"
		job := createTiDBAutomationQueueJob(t, ctx, store, automationQueueJobFixture{
			WorkspaceID: workspaceID,
			RunID:       "run-tidb-concurrent",
			JobRunID:    "job-tidb-concurrent",
			JobID:       "build",
			MaxAttempts: 3,
			Now:         now,
		})

		const workers = 8
		var wg sync.WaitGroup
		claimedIDs := make(chan int64, workers)
		errs := make(chan error, workers)
		for idx := 0; idx < workers; idx++ {
			wg.Add(1)
			go func(idx int) {
				defer wg.Done()
				claimed, err := store.ClaimJobs(ctx, automation.ClaimJobsRequest{
					WorkspaceID:     workspaceID,
					QueueName:       "default",
					ShardBucket:     0,
					LeaseOwner:      fmt.Sprintf("worker-%d", idx),
					LeaseDuration:   time.Minute,
					Limit:           1,
					Now:             now.Add(time.Duration(idx) * time.Millisecond),
					MaxRunnerLabels: 8,
				})
				if err != nil {
					errs <- err
					return
				}
				for _, claimedJob := range claimed {
					claimedIDs <- claimedJob.ID
				}
			}(idx)
		}
		wg.Wait()
		close(claimedIDs)
		close(errs)

		for err := range errs {
			require.NoError(t, err)
		}
		var ids []int64
		for id := range claimedIDs {
			ids = append(ids, id)
		}
		require.Equal(t, []int64{job.ID}, ids)
	})

	t.Run("renew keeps unexpired lease out of recovery", func(t *testing.T) {
		now := time.Date(2026, 6, 13, 10, 10, 0, 0, time.UTC)
		workspaceID := "ws-tidb-renew"
		job := createTiDBAutomationQueueJob(t, ctx, store, automationQueueJobFixture{
			WorkspaceID: workspaceID,
			RunID:       "run-tidb-renew",
			JobRunID:    "job-tidb-renew",
			JobID:       "build",
			MaxAttempts: 3,
			Now:         now,
		})
		leaseUntil := now.Add(5 * time.Minute)
		claimed, err := store.ClaimJobs(ctx, automation.ClaimJobsRequest{
			WorkspaceID:     workspaceID,
			QueueName:       "default",
			ShardBucket:     0,
			LeaseOwner:      "worker-1",
			LeaseDuration:   5 * time.Minute,
			Limit:           1,
			Now:             now,
			MaxRunnerLabels: 8,
		})
		require.NoError(t, err)
		require.Len(t, claimed, 1)
		require.Equal(t, job.ID, claimed[0].ID)

		renewed, err := store.RenewLeasedJob(ctx, workspaceID, job.ID, "worker-1", now, leaseUntil)
		require.NoError(t, err)
		require.NotNil(t, renewed.LeaseUntil)
		require.Equal(t, leaseUntil, *renewed.LeaseUntil)

		renewedUntil := now.Add(10 * time.Minute)
		renewed, err = store.RenewLeasedJob(ctx, workspaceID, job.ID, "worker-1", now.Add(time.Second), renewedUntil)
		require.NoError(t, err)
		require.NotNil(t, renewed.LeaseUntil)
		require.Equal(t, renewedUntil, *renewed.LeaseUntil)

		recovered, err := store.RecoverExpiredJobs(ctx, now.Add(6*time.Minute), 10)
		require.NoError(t, err)
		require.Zero(t, recovered)
		claimedByOther, err := store.ClaimJobs(ctx, automation.ClaimJobsRequest{
			WorkspaceID:     workspaceID,
			QueueName:       "default",
			ShardBucket:     0,
			LeaseOwner:      "worker-2",
			LeaseDuration:   time.Minute,
			Limit:           1,
			Now:             now.Add(6 * time.Minute),
			MaxRunnerLabels: 8,
		})
		require.NoError(t, err)
		require.Empty(t, claimedByOther)
		stored, err := store.GetJob(ctx, job.ID)
		require.NoError(t, err)
		require.Equal(t, automation.AutomationJobStatusLeased, stored.Status)
		require.Equal(t, "worker-1", stored.LeaseOwner)
		require.NotNil(t, stored.LeaseUntil)
		require.Equal(t, renewedUntil, *stored.LeaseUntil)
	})

	t.Run("expired recovery requeues and fails exhausted jobs", func(t *testing.T) {
		now := time.Date(2026, 6, 13, 10, 20, 0, 0, time.UTC)
		retryJob := createTiDBAutomationQueueJob(t, ctx, store, automationQueueJobFixture{
			WorkspaceID: "ws-tidb-retry",
			RunID:       "run-tidb-retry",
			JobRunID:    "job-tidb-retry",
			JobID:       "build",
			MaxAttempts: 3,
			Now:         now,
		})
		failJob := createTiDBAutomationQueueJob(t, ctx, store, automationQueueJobFixture{
			WorkspaceID: "ws-tidb-fail",
			RunID:       "run-tidb-fail",
			JobRunID:    "job-tidb-fail",
			JobID:       "test",
			MaxAttempts: 1,
			Now:         now,
		})

		requireClaimedJob(t, ctx, store, "ws-tidb-retry", retryJob.ID, "worker-retry", now)
		requireClaimedJob(t, ctx, store, "ws-tidb-fail", failJob.ID, "worker-fail", now)

		recoveryNow := now.Add(2 * time.Second)
		recovered, err := store.RecoverExpiredJobs(ctx, recoveryNow, 10)
		require.NoError(t, err)
		require.EqualValues(t, 2, recovered)

		storedRetry, err := store.GetJob(ctx, retryJob.ID)
		require.NoError(t, err)
		require.Equal(t, automation.AutomationJobStatusReady, storedRetry.Status)
		require.Empty(t, storedRetry.LeaseOwner)
		require.Nil(t, storedRetry.LeaseUntil)
		require.Equal(t, 1, storedRetry.Attempt)
		retryRun, err := store.GetJobRun(ctx, "ws-tidb-retry", "job-tidb-retry")
		require.NoError(t, err)
		require.Equal(t, automation.AutomationJobRunStatusQueued, retryRun.Status)

		claimedRetry, err := store.ClaimJobs(ctx, automation.ClaimJobsRequest{
			WorkspaceID:     "ws-tidb-retry",
			QueueName:       "default",
			ShardBucket:     0,
			LeaseOwner:      "worker-retry-2",
			LeaseDuration:   time.Minute,
			Limit:           1,
			Now:             recoveryNow.Add(31 * time.Second),
			MaxRunnerLabels: 8,
		})
		require.NoError(t, err)
		require.Len(t, claimedRetry, 1)
		require.Equal(t, retryJob.ID, claimedRetry[0].ID)

		storedFail, err := store.GetJob(ctx, failJob.ID)
		require.NoError(t, err)
		require.Equal(t, automation.AutomationJobStatusFailed, storedFail.Status)
		require.Empty(t, storedFail.LeaseOwner)
		require.Nil(t, storedFail.LeaseUntil)
		failJobRun, err := store.GetJobRun(ctx, "ws-tidb-fail", "job-tidb-fail")
		require.NoError(t, err)
		require.Equal(t, automation.AutomationJobRunStatusFailed, failJobRun.Status)
		failRun, err := store.GetRun(ctx, "ws-tidb-fail", "run-tidb-fail")
		require.NoError(t, err)
		require.Equal(t, automation.AutomationRunStatusFailed, failRun.Status)
	})
}

type automationQueueJobFixture struct {
	WorkspaceID string
	RunID       string
	JobRunID    string
	JobID       string
	MaxAttempts int
	Now         time.Time
}

func createTiDBAutomationQueueJob(t *testing.T, ctx context.Context, store *automation.Store, fixture automationQueueJobFixture) automation.AutomationJob {
	t.Helper()
	_, err := store.CreateRun(ctx, automation.AutomationRun{
		ID:          fixture.RunID,
		WorkspaceID: fixture.WorkspaceID,
		Name:        "TiDB queue lease test",
		Status:      automation.AutomationRunStatusRunning,
		CreatedAt:   fixture.Now,
		UpdatedAt:   fixture.Now,
	})
	require.NoError(t, err)
	require.NoError(t, store.CreateJobRuns(ctx, []automation.AutomationJobRun{{
		ID:          fixture.JobRunID,
		WorkspaceID: fixture.WorkspaceID,
		RunID:       fixture.RunID,
		JobID:       fixture.JobID,
		JobType:     "github_actions_job",
		Status:      automation.AutomationJobRunStatusQueued,
		ScheduledAt: fixture.Now,
		CreatedAt:   fixture.Now,
		UpdatedAt:   fixture.Now,
	}}))
	job, err := store.EnqueueJob(ctx, automation.AutomationJob{
		WorkspaceID:    fixture.WorkspaceID,
		RunID:          fixture.RunID,
		JobRunID:       fixture.JobRunID,
		QueueName:      "default",
		Status:         automation.AutomationJobStatusReady,
		ScheduledAt:    fixture.Now,
		MaxAttempts:    fixture.MaxAttempts,
		IdempotencyKey: fixture.RunID + ":" + fixture.JobID,
		ShardBucket:    0,
		CreatedAt:      fixture.Now,
		UpdatedAt:      fixture.Now,
	})
	require.NoError(t, err)
	return job
}

func requireClaimedJob(t *testing.T, ctx context.Context, store *automation.Store, workspaceID string, jobID int64, leaseOwner string, now time.Time) {
	t.Helper()
	claimed, err := store.ClaimJobs(ctx, automation.ClaimJobsRequest{
		WorkspaceID:     workspaceID,
		QueueName:       "default",
		ShardBucket:     0,
		LeaseOwner:      leaseOwner,
		LeaseDuration:   time.Second,
		Limit:           1,
		Now:             now,
		MaxRunnerLabels: 8,
	})
	require.NoError(t, err)
	require.Len(t, claimed, 1)
	require.Equal(t, jobID, claimed[0].ID)
}

func openTiDBAutomationTestDB(t *testing.T) *gorm.DB {
	t.Helper()
	return openTiDBTestDB(t, func(db *gorm.DB) error {
		if err := applyAutomationMigrations(db); err != nil {
			return err
		}
		return VerifySchema(db)
	})
}

func openTiDBTestDB(t *testing.T, migrate func(*gorm.DB) error) *gorm.DB {
	t.Helper()
	dsn := tidbTestDSN(t)

	isolatedDSN, cleanup := createIsolatedTiDBDatabase(t, dsn)
	t.Cleanup(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		require.NoError(t, cleanup(ctx))
	})

	db, err := open(Config{DSN: isolatedDSN}, false)
	require.NoError(t, err)
	require.NoError(t, migrate(db))
	t.Cleanup(func() {
		require.NoError(t, Close(db))
	})
	return db
}

func tidbTestDSN(t *testing.T) string {
	t.Helper()
	dsn := strings.TrimSpace(os.Getenv("MESHX_TIDB_TEST_DSN"))
	if dsn == "" {
		dsn = strings.TrimSpace(os.Getenv("AGENTTEAM_TIDB_TEST_DSN"))
	}
	if dsn == "" {
		t.Skip("set MESHX_TIDB_TEST_DSN or AGENTTEAM_TIDB_TEST_DSN to run MeshX TiDB queue lease integration tests")
	}
	return dsn
}

func createIsolatedTiDBDatabase(t *testing.T, dsn string) (string, func(context.Context) error) {
	t.Helper()
	cfg, err := mysql.ParseDSN(dsn)
	require.NoError(t, err)
	configureTiDBTestDSN(cfg)

	adminCfg := *cfg
	adminCfg.DBName = ""
	adminDB, err := sql.Open("mysql", adminCfg.FormatDSN())
	require.NoError(t, err)
	adminDB.SetMaxOpenConns(1)
	adminDB.SetMaxIdleConns(1)
	t.Cleanup(func() {
		require.NoError(t, adminDB.Close())
	})

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	require.NoError(t, adminDB.PingContext(ctx))
	requireTiDBTestDatabase(t, ctx, adminDB)

	databaseName := fmt.Sprintf("meshx_tidb_queue_lease_%d", time.Now().UnixNano())
	_, err = adminDB.ExecContext(ctx, "CREATE DATABASE "+quoteTiDBIdentifier(databaseName))
	require.NoError(t, err)

	cfg.DBName = databaseName
	cleanup := func(ctx context.Context) error {
		if ctx == nil {
			ctx = context.Background()
		}
		_, err := adminDB.ExecContext(ctx, "DROP DATABASE IF EXISTS "+quoteTiDBIdentifier(databaseName))
		return err
	}
	return cfg.FormatDSN(), cleanup
}

func requireTiDBTestDatabase(t *testing.T, ctx context.Context, db *sql.DB) {
	t.Helper()
	var version string
	require.NoError(t, db.QueryRowContext(ctx, "SELECT VERSION()").Scan(&version))
	if !strings.Contains(strings.ToLower(version), "tidb") {
		t.Skipf("MeshX TiDB queue lease integration tests require TiDB, got database version %q", version)
	}
}

func configureTiDBTestDSN(cfg *mysql.Config) {
	cfg.ParseTime = true
	if cfg.Loc == nil {
		cfg.Loc = time.UTC
	}
	if cfg.Params == nil {
		cfg.Params = map[string]string{}
	}
	if _, ok := cfg.Params["time_zone"]; !ok {
		cfg.Params["time_zone"] = "'+00:00'"
	}
}

func quoteTiDBIdentifier(identifier string) string {
	return "`" + strings.ReplaceAll(identifier, "`", "``") + "`"
}