package service_test

import (
	"context"
	"fmt"
	"slices"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"gorm.io/gorm"

	"github.com/pingcap/agent-git-service/internal/db"
	"github.com/pingcap/agent-git-service/internal/embedding"
	"github.com/pingcap/agent-git-service/internal/gitstore"
	"github.com/pingcap/agent-git-service/internal/service"
	"github.com/pingcap/agent-git-service/internal/testharness/testdb"
	"github.com/pingcap/agent-git-service/internal/wikicatalog"
)

func openTenantTestDB(t *testing.T, prefix string) *gorm.DB {
	t.Helper()
	gdb, cleanup := testdb.OpenRaw(t, prefix)
	t.Cleanup(cleanup)
	return gdb
}

// TestDBForCtxUsesTenantDB verifies that DBForCtx returns the tenant DB
// when one is injected into the context.
func TestDBForCtxUsesTenantDB(t *testing.T) {
	defaultDB := openTenantTestDB(t, "tenant_default")
	tenantDB := openTenantTestDB(t, "tenant_context")

	svc := &service.Service{DB: defaultDB}

	// Inject tenant DB into context
	ctx := service.ContextWithDB(context.Background(), tenantDB)

	// Verify DBForCtx returns the tenant DB
	got := svc.DBForCtx(ctx)

	// The returned DB should be derived from tenantDB, not defaultDB.
	gotSQL, _ := got.DB()
	tenantSQL, _ := tenantDB.DB()
	defaultSQL, _ := defaultDB.DB()

	if gotSQL != tenantSQL {
		t.Error("DBForCtx did not return the context-injected DB")
	}
	if gotSQL == defaultSQL {
		t.Error("DBForCtx returned the default DB instead of the context DB")
	}
}

func TestWikiBackgroundGitIngestStateIsTenantScoped_Issue1448(t *testing.T) {
	defaultDB := openMigratedServiceTestDB(t)
	tenantDB := openMigratedServiceTestDB(t)

	store, err := gitstore.New(t.TempDir())
	if err != nil {
		t.Fatalf("gitstore.New: %v", err)
	}
	blobStore := wikicatalog.NewBlobStore(t.TempDir())
	wikiCat := wikicatalog.New(defaultDB, blobStore)
	svc, err := service.New(service.Config{
		DB:          defaultDB,
		Git:         store,
		WikiCatalog: wikiCat,
		WikiBlob:    blobStore,
		BaseURL:     "http://localhost:8080",
		Embedder:    embedding.NopEmbedder{},
	})
	if err != nil {
		t.Fatalf("service.New: %v", err)
	}
	wikiCat.DBFor = svc.DBForCtx
	wikiCat.OnChangeSetCommitted = svc.WikiCatalogPostCommit

	defaultCtx := context.Background()
	tenantCtx := service.ContextWithDB(context.Background(), tenantDB)
	for _, tc := range []struct {
		ctx  context.Context
		gdb  *gorm.DB
		user string
	}{
		{ctx: defaultCtx, gdb: defaultDB, user: "shared-owner"},
		{ctx: tenantCtx, gdb: tenantDB, user: "shared-owner"},
	} {
		owner := db.User{Login: tc.user, Name: tc.user, Type: db.TypeUser}
		if err := tc.gdb.Create(&owner).Error; err != nil {
			t.Fatalf("create %s owner: %v", tc.user, err)
		}
		if _, err := svc.CreateRepo(tc.ctx, service.CreateRepoInput{OwnerLogin: owner.Login, Name: "shared-name", AutoInit: true}); err != nil {
			t.Fatalf("CreateRepo %s: %v", tc.user, err)
		}
		repoFullName := owner.Login + "/shared-name"
		if err := tc.gdb.Model(&db.Repository{}).Where("full_name = ?", repoFullName).Update("has_wiki", true).Error; err != nil {
			t.Fatalf("set has_wiki for %s: %v", repoFullName, err)
		}
	}

	defaultRepo := "shared-owner/shared-name"
	tenantRepo := "shared-owner/shared-name"
	if !svc.ClaimWikiBackgroundGitIngestForTest(defaultCtx, defaultRepo) {
		t.Fatal("expected to claim default tenant git ingest slot")
	}
	defer svc.ReleaseWikiBackgroundGitIngestForTest(defaultCtx, defaultRepo)

	if !svc.IsWikiBackgroundGitIngestRunning(defaultCtx, defaultRepo) {
		t.Fatal("expected default tenant git ingest state to be visible in default context")
	}
	if svc.IsWikiBackgroundGitIngestRunning(tenantCtx, tenantRepo) {
		t.Fatal("tenant-scoped git ingest state leaked from default DB into tenant DB")
	}
}

// TestDBForCtxFallsBackToDefault verifies that DBForCtx falls back to
// the default DB when no tenant DB is in the context.
func TestDBForCtxFallsBackToDefault(t *testing.T) {
	defaultDB := openTenantTestDB(t, "tenant_fallback")

	svc := &service.Service{DB: defaultDB}

	// Create context WITHOUT tenant DB
	ctx := context.Background()

	// Verify DBForCtx falls back to default DB
	got := svc.DBForCtx(ctx)
	gotSQL, _ := got.DB()
	defaultSQL, _ := defaultDB.DB()

	if gotSQL != defaultSQL {
		t.Error("DBForCtx did not fall back to s.DB when no context DB was set")
	}
}

// TestBackgroundGoroutinePropagatesTenantDB verifies that background
// goroutines properly propagate tenant DB context.
func TestBackgroundGoroutinePropagatesTenantDB(t *testing.T) {
	defaultDB := openTenantTestDB(t, "tenant_bg_default")
	tenantDB := openTenantTestDB(t, "tenant_bg_tenant")

	svc := &service.Service{
		DB:  defaultDB,
		Ctx: context.Background(),
	}

	// Track which DB was used in the background goroutine
	var bgDB *gorm.DB
	done := make(chan struct{})

	// Simulate the pattern used in embedding_hook.go and workflow_dispatch.go
	requestCtx := service.ContextWithDB(context.Background(), tenantDB)
	bgCtx := svc.ServerCtx()
	if db, ok := service.DBFromContext(requestCtx); ok {
		bgCtx = service.ContextWithDB(bgCtx, db)
	}

	svc.Wg.Add(1)
	go func() {
		defer svc.Wg.Done()
		bgDB = svc.DBForCtx(bgCtx)
		close(done)
	}()

	<-done

	if bgDB == nil {
		t.Fatal("Background goroutine did not use any DB")
	}

	// Verify the background goroutine used the tenant DB, not the default
	bgSQL, _ := bgDB.DB()
	tenantSQL, _ := tenantDB.DB()
	defaultSQL, _ := defaultDB.DB()

	if bgSQL != tenantSQL {
		t.Error("Background goroutine did not use the tenant DB")
	}
	if bgSQL == defaultSQL {
		t.Error("Background goroutine used the default DB instead of tenant DB")
	}
}

// TestContextHelpers verifies that context helper functions work correctly.
func TestContextHelpers(t *testing.T) {
	testDB := openTenantTestDB(t, "tenant_helpers")

	// Test ContextWithDB and DBFromContext
	ctx := service.ContextWithDB(context.Background(), testDB)

	retrievedDB, ok := service.DBFromContext(ctx)
	if !ok {
		t.Fatal("DBFromContext returned ok=false for context with DB")
	}

	if retrievedDB != testDB {
		t.Fatal("DBFromContext returned a different DB than what was injected")
	}

	// Test with context without DB
	emptyCtx := context.Background()
	_, ok = service.DBFromContext(emptyCtx)
	if ok {
		t.Fatal("DBFromContext returned ok=true for context without DB")
	}
}

func TestRepositoryMutationOwnerPinsScopedDatabase(t *testing.T) {
	defaultDB := openMigratedServiceTestDB(t)
	scopedDB := openMigratedServiceTestDB(t)
	store, err := gitstore.New(t.TempDir())
	if err != nil {
		t.Fatalf("gitstore.New: %v", err)
	}
	svc, err := service.New(service.Config{DB: defaultDB, Git: store, BaseURL: "http://localhost:8080"})
	if err != nil {
		t.Fatalf("service.New: %v", err)
	}

	owner := db.User{Login: "scoped-owner", Name: "scoped-owner", Type: db.TypeUser}
	if err := defaultDB.Create(&owner).Error; err != nil {
		t.Fatalf("create default owner: %v", err)
	}
	repository, err := svc.CreateRepo(context.Background(), service.CreateRepoInput{
		OwnerLogin: owner.Login,
		Name:       "scoped-repository",
		AutoInit:   true,
	})
	if err != nil {
		t.Fatalf("CreateRepo default: %v", err)
	}

	scopedOwner := db.User{ID: owner.ID, Login: owner.Login, Name: owner.Name, Type: owner.Type}
	if err := scopedDB.Create(&scopedOwner).Error; err != nil {
		t.Fatalf("create scoped owner: %v", err)
	}
	scopedRepository := db.Repository{
		ID: repository.ID, Name: repository.Name, FullName: repository.FullName,
		OwnerID: scopedOwner.ID, DefaultBranch: repository.DefaultBranch,
	}
	if err := scopedDB.Create(&scopedRepository).Error; err != nil {
		t.Fatalf("create scoped repository: %v", err)
	}
	scopedFreeze := db.RepositoryMigrationFreeze{RepositoryID: repository.ID, State: "sealed", Epoch: 7}
	if err := scopedDB.Create(&scopedFreeze).Error; err != nil {
		t.Fatalf("create scoped freeze: %v", err)
	}

	scopedCtx := service.ContextWithDB(context.Background(), scopedDB)
	if err := svc.CreateRepositoryBranch(scopedCtx, repository.ID, repository.FullName, "sealed-state-ignored", "main"); err != nil {
		t.Fatalf("scoped legacy sealed state blocked mutation: %v", err)
	}
	if _, err := svc.Git.HeadSHA(scopedCtx, repository.FullName, "sealed-state-ignored"); err != nil {
		t.Fatalf("resolve mutation with ignored legacy state: %v", err)
	}

	if err := defaultDB.Model(&db.RepositoryMigrationFreeze{}).Where("repository_id = ?", repository.ID).
		Updates(map[string]any{"state": "sealed", "epoch": 2}).Error; err != nil {
		t.Fatalf("seal default repository: %v", err)
	}
	if err := svc.CreateRepositoryBranch(scopedCtx, repository.ID, repository.FullName, "scoped-authority", "main"); err != nil {
		t.Fatalf("scoped repository mutation: %v", err)
	}
	if _, err := svc.Git.HeadSHA(scopedCtx, repository.FullName, "scoped-authority"); err != nil {
		t.Fatalf("resolve scoped branch: %v", err)
	}

	emptyScopedDB := openMigratedServiceTestDB(t)
	emptyScopedCtx := service.ContextWithDB(context.Background(), emptyScopedDB)
	if err := svc.CreateRepositoryBranch(emptyScopedCtx, repository.ID, repository.FullName, "must-not-use-default", "main"); err == nil {
		t.Fatal("missing scoped identity fell back to the default database")
	}
	if _, err := svc.Git.HeadSHA(emptyScopedCtx, repository.FullName, "must-not-use-default"); err == nil {
		t.Fatal("default-only authority published through an empty scoped database")
	}

	var capture atomic.Bool
	var queryMu sync.Mutex
	var queries []string
	const callbackName = "test:scoped_repository_mutation_query_shape"
	if err := scopedDB.Callback().Query().After("gorm:query").Register(callbackName, func(tx *gorm.DB) {
		if !capture.Load() {
			return
		}
		queryMu.Lock()
		queries = append(queries, tx.Statement.SQL.String())
		queryMu.Unlock()
	}); err != nil {
		t.Fatalf("register query callback: %v", err)
	}
	t.Cleanup(func() { _ = scopedDB.Callback().Query().Remove(callbackName) })

	captureQueries := func(branch string) []string {
		queryMu.Lock()
		queries = nil
		queryMu.Unlock()
		capture.Store(true)
		err := svc.CreateRepositoryBranch(scopedCtx, repository.ID, repository.FullName, branch, "main")
		capture.Store(false)
		if err != nil {
			t.Fatalf("CreateRepositoryBranch(%s): %v", branch, err)
		}
		queryMu.Lock()
		defer queryMu.Unlock()
		return append([]string(nil), queries...)
	}

	withoutUnrelated := captureQueries("query-shape-zero")
	unrelated := make([]db.Repository, 128)
	for index := range unrelated {
		unrelated[index] = db.Repository{
			Name: fmt.Sprintf("unrelated-%03d", index), FullName: fmt.Sprintf("%s/unrelated-%03d", owner.Login, index),
			OwnerID: scopedOwner.ID, DefaultBranch: "main",
		}
	}
	if err := scopedDB.CreateInBatches(unrelated, 64).Error; err != nil {
		t.Fatalf("create unrelated repositories: %v", err)
	}
	withUnrelated := captureQueries("query-shape-many")
	if !slices.Equal(withoutUnrelated, withUnrelated) {
		t.Fatalf("repository mutation query shape changed with unrelated repositories:\nzero=%q\nmany=%q", withoutUnrelated, withUnrelated)
	}
}

func TestKickBackgroundWikiGitIngest_UsesCallerTenantContext_Issue1448(t *testing.T) {
	defaultDB := openMigratedServiceTestDB(t)
	tenantDB := openMigratedServiceTestDB(t)

	store, err := gitstore.New(t.TempDir())
	if err != nil {
		t.Fatalf("gitstore.New: %v", err)
	}
	blobStore := wikicatalog.NewBlobStore(t.TempDir())
	wikiCat := wikicatalog.New(defaultDB, blobStore)
	svc, err := service.New(service.Config{
		DB:          defaultDB,
		Git:         store,
		WikiCatalog: wikiCat,
		WikiBlob:    blobStore,
		BaseURL:     "http://localhost:8080",
		Embedder:    embedding.NopEmbedder{},
	})
	if err != nil {
		t.Fatalf("service.New: %v", err)
	}
	wikiCat.DBFor = svc.DBForCtx
	wikiCat.OnChangeSetCommitted = svc.WikiCatalogPostCommit

	tenantCtx := service.ContextWithDB(context.Background(), tenantDB)
	owner := db.User{Login: "tenant-owner", Name: "tenant-owner", Type: db.TypeUser}
	if err := tenantDB.Create(&owner).Error; err != nil {
		t.Fatalf("create tenant owner: %v", err)
	}
	if _, err := svc.CreateRepo(tenantCtx, service.CreateRepoInput{OwnerLogin: owner.Login, Name: "tenant-wiki", AutoInit: true}); err != nil {
		t.Fatalf("CreateRepo: %v", err)
	}
	repoFullName := owner.Login + "/tenant-wiki"
	if err := tenantDB.Model(&db.Repository{}).Where("full_name = ?", repoFullName).Update("has_wiki", true).Error; err != nil {
		t.Fatalf("set has_wiki: %v", err)
	}
	if _, err := svc.PutWikiPage(tenantCtx, repoFullName, "home", "v1", "create home", ""); err != nil {
		t.Fatalf("PutWikiPage home: %v", err)
	}
	if _, err := writeRepositoryFileFixture(t, svc, tenantCtx, repoFullName+".wiki", "master", "about.md", "add about", []byte("about body")); err != nil {
		t.Fatalf("git write about: %v", err)
	}

	started := make(chan struct{}, 1)
	svc.SetWikiBackgroundGitIngestStartedHookForTest(func(fullName string) {
		if fullName == repoFullName {
			started <- struct{}{}
		}
	})
	defer svc.SetWikiBackgroundGitIngestStartedHookForTest(nil)

	svc.KickBackgroundWikiGitIngest(tenantCtx, repoFullName)

	select {
	case <-started:
	case <-time.After(3 * time.Second):
		t.Fatal("timed out waiting for background git ingest to start")
	}
	svc.Wg.Wait()

	pages, err := svc.ListWikiPages(tenantCtx, repoFullName, service.ListWikiPagesOptions{Recursive: true})
	if err != nil {
		t.Fatalf("ListWikiPages after background git ingest: %v", err)
	}
	if len(pages) != 2 {
		t.Fatalf("pages = %+v, want 2 pages after background git ingest", pages)
	}
}
