package app

import (
	"context"
	"errors"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/go-sql-driver/mysql"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/extension/providers/searchengine/runtime"
	"github.com/tidbcloud/AgentTeam/internal/store/app/registry"
	"github.com/tidbcloud/AgentTeam/internal/store/model"
	"github.com/tidbcloud/AgentTeam/internal/store/serviceactor"
)

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

	ctx := context.Background()
	pool := startAppRollbackDB(t, ctx)
	store := New(pool, Options{AppCredentialStorageKey: testAppCredentialStorageKey()})
	now := time.Now().UTC()
	admin := seedAppAdminFixture(t, pool, "search_connection_brave", now)

	installation, err := store.InstallAppForWorkspaceAdmin(
		ctx,
		admin.UserID,
		admin.WorkspaceID,
		InstallAppParams{AppID: registry.BraveSearchAppID},
	)
	require.NoError(t, err)

	selections, err := serviceactor.New(pool).ListInstallationExtensions(
		ctx,
		admin.WorkspaceID,
		installation.ID,
	)
	require.NoError(t, err)
	require.Len(t, selections, 1)
	require.Equal(t, runtime.ProviderBrave, selections[0].ExtensionName)
	require.Equal(t, model.BraveBuiltinExtensionDefinitionDigest, selections[0].DefinitionDigest)

	view, err := store.UpsertSearchEngineConnectionForUser(
		ctx,
		admin.UserID,
		admin.WorkspaceID,
		installation.ID,
		UpsertSearchEngineConnectionParams{
			Config:      map[string]any{},
			Credentials: map[string]any{"apiKey": "initial-secret"},
		},
	)
	require.NoError(t, err)
	require.True(t, view.Active)
	require.True(t, view.APIKeyConfigured)

	var storedCredentials []byte
	require.NoError(t, pool.QueryRow(ctx, `
SELECT credentials
FROM service_connections
WHERE service_installation_id = $1
  AND external_subject_ref = 'default'
  AND status = 'active'
`, installation.ID).Scan(&storedCredentials))
	require.NotContains(t, string(storedCredentials), "initial-secret")
	require.Contains(t, string(storedCredentials), CredentialStorageScheme)

	runtimeConnection, err := store.ResolveSearchEngineRuntimeConnection(
		ctx,
		admin.WorkspaceID,
		installation.ID,
		runtime.ProviderBrave,
	)
	require.NoError(t, err)
	require.Equal(t, "initial-secret", runtimeConnection.APIKey)

	_, err = store.UpsertSearchEngineConnectionForUser(
		ctx,
		admin.UserID,
		admin.WorkspaceID,
		installation.ID,
		UpsertSearchEngineConnectionParams{
			Config:      map[string]any{},
			Credentials: map[string]any{},
		},
	)
	require.NoError(t, err)
	runtimeConnection, err = store.ResolveSearchEngineRuntimeConnection(
		ctx,
		admin.WorkspaceID,
		installation.ID,
		runtime.ProviderBrave,
	)
	require.NoError(t, err)
	require.Equal(t, "initial-secret", runtimeConnection.APIKey)

	peer := seedWorkspaceMemberFixture(
		t,
		pool,
		admin.WorkspaceID,
		"search_connection_brave_peer",
		now,
	)
	_, err = store.GetSearchEngineConnectionForUser(
		ctx,
		peer.UserID,
		admin.WorkspaceID,
		installation.ID,
	)
	require.Error(t, err)
	require.True(t, errors.Is(err, ErrForbidden))

	disconnected, err := store.DeleteSearchEngineConnectionForUser(
		ctx,
		admin.UserID,
		admin.WorkspaceID,
		installation.ID,
	)
	require.NoError(t, err)
	require.False(t, disconnected.Active)
	_, err = store.ResolveSearchEngineRuntimeConnection(
		ctx,
		admin.WorkspaceID,
		installation.ID,
		runtime.ProviderBrave,
	)
	require.Error(t, err)
}

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

	ctx := context.Background()
	pool := startAppRollbackDB(t, ctx)
	store := New(pool, Options{AppCredentialStorageKey: testAppCredentialStorageKey()})
	admin := seedAppAdminFixture(t, pool, "search_connection_searxng", time.Now().UTC())
	installation, err := store.InstallAppForWorkspaceAdmin(
		ctx,
		admin.UserID,
		admin.WorkspaceID,
		InstallAppParams{AppID: registry.SearXNGAppID},
	)
	require.NoError(t, err)

	view, err := store.UpsertSearchEngineConnectionForUser(
		ctx,
		admin.UserID,
		admin.WorkspaceID,
		installation.ID,
		UpsertSearchEngineConnectionParams{
			Config: map[string]any{
				"baseUrl":  "https://search.example.com/base",
				"language": "zh-CN",
			},
			Credentials: map[string]any{},
		},
	)
	require.NoError(t, err)
	require.Equal(t, "https://search.example.com/base", view.Config.BaseURL)
	require.Equal(t, "zh-CN", view.Config.Language)
	require.False(t, view.APIKeyConfigured)

	runtimeConnection, err := store.ResolveSearchEngineRuntimeConnection(
		ctx,
		admin.WorkspaceID,
		installation.ID,
		runtime.ProviderSearXNG,
	)
	require.NoError(t, err)
	require.Equal(t, "https://search.example.com/base", runtimeConnection.BaseURL)
	require.Equal(t, "zh-CN", runtimeConnection.Language)
	require.Empty(t, runtimeConnection.APIKey)
}

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

	ctx := context.Background()
	pool := startAppRollbackDB(t, ctx)
	baseStore := New(pool, Options{
		AppCredentialStorageKey: testAppCredentialStorageKey(),
	})
	admin := seedAppAdminFixture(
		t,
		pool,
		"search_connection_retry_upsert",
		time.Now().UTC(),
	)
	installation, err := baseStore.InstallAppForWorkspaceAdmin(
		ctx,
		admin.UserID,
		admin.WorkspaceID,
		InstallAppParams{AppID: registry.BraveSearchAppID},
	)
	require.NoError(t, err)

	retryingDB := &retrySearchConnectionMutationDB{
		DB:     pool,
		needle: "INSERT INTO service_connections",
	}
	store := New(retryingDB, Options{
		AppCredentialStorageKey: testAppCredentialStorageKey(),
	})
	view, err := store.UpsertSearchEngineConnectionForUser(
		ctx,
		admin.UserID,
		admin.WorkspaceID,
		installation.ID,
		UpsertSearchEngineConnectionParams{
			Config:      map[string]any{},
			Credentials: map[string]any{"apiKey": "retry-secret"},
		},
	)
	require.NoError(t, err)
	require.True(t, retryingDB.wasInjected())
	require.True(t, view.Active)

	var activeCount int
	require.NoError(t, pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM service_connections
WHERE service_installation_id = $1
  AND external_subject_ref = 'default'
  AND status = 'active'
`, installation.ID).Scan(&activeCount))
	require.Equal(t, 1, activeCount)
}

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

	ctx := context.Background()
	pool := startAppRollbackDB(t, ctx)
	baseStore := New(pool, Options{
		AppCredentialStorageKey: testAppCredentialStorageKey(),
	})
	admin := seedAppAdminFixture(
		t,
		pool,
		"search_connection_retry_delete",
		time.Now().UTC(),
	)
	installation, err := baseStore.InstallAppForWorkspaceAdmin(
		ctx,
		admin.UserID,
		admin.WorkspaceID,
		InstallAppParams{AppID: registry.BraveSearchAppID},
	)
	require.NoError(t, err)
	_, err = baseStore.UpsertSearchEngineConnectionForUser(
		ctx,
		admin.UserID,
		admin.WorkspaceID,
		installation.ID,
		UpsertSearchEngineConnectionParams{
			Config:      map[string]any{},
			Credentials: map[string]any{"apiKey": "retry-secret"},
		},
	)
	require.NoError(t, err)

	retryingDB := &retrySearchConnectionMutationDB{
		DB:     pool,
		needle: "UPDATE service_connections\nSET status = 'disabled'",
	}
	store := New(retryingDB, Options{
		AppCredentialStorageKey: testAppCredentialStorageKey(),
	})
	view, err := store.DeleteSearchEngineConnectionForUser(
		ctx,
		admin.UserID,
		admin.WorkspaceID,
		installation.ID,
	)
	require.NoError(t, err)
	require.True(t, retryingDB.wasInjected())
	require.False(t, view.Active)

	var activeCount int
	require.NoError(t, pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM service_connections
WHERE service_installation_id = $1
  AND external_subject_ref = 'default'
  AND status = 'active'
`, installation.ID).Scan(&activeCount))
	require.Zero(t, activeCount)
}

type retrySearchConnectionMutationDB struct {
	coredb.DB
	mu       sync.Mutex
	needle   string
	injected bool
}

func (db *retrySearchConnectionMutationDB) BeginTx(
	ctx context.Context,
	options coredb.TxOptions,
) (coredb.Tx, error) {
	tx, err := db.DB.BeginTx(ctx, options)
	if err != nil {
		return nil, err
	}
	return &retrySearchConnectionMutationTx{Tx: tx, parent: db}, nil
}

func (db *retrySearchConnectionMutationDB) wasInjected() bool {
	db.mu.Lock()
	defer db.mu.Unlock()
	return db.injected
}

type retrySearchConnectionMutationTx struct {
	coredb.Tx
	parent *retrySearchConnectionMutationDB
}

func (tx *retrySearchConnectionMutationTx) Exec(
	ctx context.Context,
	query string,
	args ...any,
) (coredb.Result, error) {
	if strings.Contains(query, tx.parent.needle) {
		tx.parent.mu.Lock()
		if !tx.parent.injected {
			tx.parent.injected = true
			tx.parent.mu.Unlock()
			return nil, &mysql.MySQLError{
				Number:  1213,
				Message: "Deadlock found when trying to get lock; try restarting transaction",
			}
		}
		tx.parent.mu.Unlock()
	}
	return tx.Tx.Exec(ctx, query, args...)
}

func TestSearchEngineConnectionWorkspaceIdentityMismatchFailsClosed(
	t *testing.T,
) {
	t.Parallel()
	testSearchEngineConnectionIdentityMismatchFailsClosed(t, "workspace")
}

func TestSearchEngineConnectionDefinitionIdentityMismatchFailsClosed(
	t *testing.T,
) {
	t.Parallel()
	testSearchEngineConnectionIdentityMismatchFailsClosed(t, "definition")
}

func testSearchEngineConnectionIdentityMismatchFailsClosed(
	t *testing.T,
	mismatch string,
) {
	t.Helper()
	ctx := context.Background()
	pool := startAppRollbackDB(t, ctx)
	store := New(pool, Options{
		AppCredentialStorageKey: testAppCredentialStorageKey(),
	})
	now := time.Now().UTC()
	admin := seedAppAdminFixture(
		t,
		pool,
		"search_identity_"+mismatch,
		now,
	)
	other := seedAppAdminFixture(
		t,
		pool,
		"search_identity_other_"+mismatch,
		now,
	)
	installation, err := store.InstallAppForWorkspaceAdmin(
		ctx,
		admin.UserID,
		admin.WorkspaceID,
		InstallAppParams{AppID: registry.BraveSearchAppID},
	)
	require.NoError(t, err)
	_, err = store.UpsertSearchEngineConnectionForUser(
		ctx,
		admin.UserID,
		admin.WorkspaceID,
		installation.ID,
		UpsertSearchEngineConnectionParams{
			Config:      map[string]any{},
			Credentials: map[string]any{"apiKey": "tenant-secret"},
		},
	)
	require.NoError(t, err)
	if mismatch == "workspace" {
		mustExec(t, pool, `
UPDATE service_connections SET workspace_id = $1
WHERE service_installation_id = $2 AND status = 'active'
`, other.WorkspaceID, installation.ID)
	} else {
		mustExec(t, pool, `
UPDATE service_connections SET service_definition_id = $1
WHERE service_installation_id = $2 AND status = 'active'
`, registry.SearXNGAppID, installation.ID)
	}
	assertSearchConnectionIdentityFailsClosed(
		t,
		ctx,
		store,
		admin,
		installation.ID,
	)
}

func assertSearchConnectionIdentityFailsClosed(
	t *testing.T,
	ctx context.Context,
	store *Store,
	admin appAdminFixture,
	installationID string,
) {
	t.Helper()
	_, err := store.GetSearchEngineConnectionForUser(
		ctx,
		admin.UserID,
		admin.WorkspaceID,
		installationID,
	)
	require.ErrorIs(t, err, ErrConflict)

	_, err = store.ResolveSearchEngineRuntimeConnection(
		ctx,
		admin.WorkspaceID,
		installationID,
		runtime.ProviderBrave,
	)
	require.ErrorIs(t, err, runtime.ErrRuntimeConnectionUnavailable)

	_, err = store.UpsertSearchEngineConnectionForUser(
		ctx,
		admin.UserID,
		admin.WorkspaceID,
		installationID,
		UpsertSearchEngineConnectionParams{
			Config:      map[string]any{},
			Credentials: map[string]any{},
		},
	)
	require.ErrorIs(t, err, ErrConflict)

	_, err = store.DeleteSearchEngineConnectionForUser(
		ctx,
		admin.UserID,
		admin.WorkspaceID,
		installationID,
	)
	require.ErrorIs(t, err, ErrConflict)
}