package catalogverificationtx

import (
	"context"
	"database/sql"
	"errors"
	"regexp"
	"testing"
	"time"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/pingcap-inc/nexuscrm-controlplane/internal/catalogsystemfence"
)

func TestSystemAuthorityFenceIsAcquiredOnceAndReleasedOnClose(t *testing.T) {
	db, mock, token := newOwnedVerificationToken(t)
	defer db.Close()
	expectSystemFenceAcquire(mock, 1)
	if err := token.acquireSystemAuthorityFence(context.Background(), db); err != nil {
		t.Fatal(err)
	}
	if err := token.acquireSystemAuthorityFence(context.Background(), db); err != nil {
		t.Fatal(err)
	}

	mock.ExpectRollback()
	expectSystemFenceRelease(mock, 1)
	expectSystemSessionRestore(mock)
	if err := token.Close(); err != nil {
		t.Fatal(err)
	}
	if token.ActiveFor(db) {
		t.Fatal("closed token remained active")
	}
	assertVerificationExpectations(t, mock)
}

func TestSystemAuthorityFenceUsesRetainedSessionWithOneOpenConnection(
	t *testing.T,
) {
	db, mock, token := newOwnedVerificationToken(t)
	defer db.Close()
	db.SetMaxOpenConns(1)
	expectSystemFenceAcquire(mock, 1)
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	if err := token.acquireSystemAuthorityFence(ctx, db); err != nil {
		t.Fatal(err)
	}
	stats := db.Stats()
	if stats.InUse != 1 || stats.OpenConnections != 1 {
		t.Fatalf(
			"retained session stats = in-use %d, open %d; want exactly one",
			stats.InUse,
			stats.OpenConnections,
		)
	}

	mock.ExpectRollback()
	expectSystemFenceRelease(mock, 1)
	expectSystemSessionRestore(mock)
	if err := token.Close(); err != nil {
		t.Fatal(err)
	}
	assertVerificationExpectations(t, mock)
}

func TestSystemAuthorityFenceIsReleasedAfterCommit(t *testing.T) {
	db, mock, token := newOwnedVerificationToken(t)
	defer db.Close()
	expectSystemFenceAcquire(mock, 1)
	if err := token.acquireSystemAuthorityFence(context.Background(), db); err != nil {
		t.Fatal(err)
	}

	mock.ExpectCommit()
	expectSystemFenceRelease(mock, 1)
	expectSystemSessionRestore(mock)
	if err := token.Commit(); err != nil {
		t.Fatal(err)
	}
	if token.ActiveFor(db) {
		t.Fatal("committed token remained active")
	}
	assertVerificationExpectations(t, mock)
}

func TestSystemAuthorityFenceAcquireFailureDoesNotPoisonToken(t *testing.T) {
	db, mock, token := newOwnedVerificationToken(t)
	defer db.Close()
	expectSystemFenceAcquire(mock, 0)
	if err := token.acquireSystemAuthorityFence(
		context.Background(),
		db,
	); !errors.Is(err, catalogsystemfence.ErrNotAcquired) {
		t.Fatalf("first acquire = %v, want ErrNotAcquired", err)
	}
	if !token.ActiveFor(db) {
		t.Fatal("failed fence acquisition closed the source token")
	}

	expectSystemFenceAcquire(mock, 1)
	if err := token.acquireSystemAuthorityFence(context.Background(), db); err != nil {
		t.Fatal(err)
	}
	mock.ExpectRollback()
	expectSystemFenceRelease(mock, 1)
	expectSystemSessionRestore(mock)
	if err := token.Close(); err != nil {
		t.Fatal(err)
	}
	assertVerificationExpectations(t, mock)
}

func TestSystemAuthorityFenceReleaseFailureDiscardsAfterCommit(
	t *testing.T,
) {
	db, mock, token := newOwnedVerificationToken(t)
	defer db.Close()
	expectSystemFenceAcquire(mock, 1)
	if err := token.acquireSystemAuthorityFence(context.Background(), db); err != nil {
		t.Fatal(err)
	}

	mock.ExpectCommit()
	expectSystemFenceRelease(mock, 0)
	if err := token.Commit(); !errors.Is(
		err,
		catalogsystemfence.ErrNotReleased,
	) {
		t.Fatalf("Commit = %v, want ErrNotReleased", err)
	}
	if token.ActiveFor(db) {
		t.Fatal("failed release left the token active")
	}
	assertVerificationExpectations(t, mock)
}

func TestSystemAuthorityFenceCloseJoinsReleaseAndRollbackErrors(t *testing.T) {
	db, mock, token := newOwnedVerificationToken(t)
	defer db.Close()
	expectSystemFenceAcquire(mock, 1)
	if err := token.acquireSystemAuthorityFence(context.Background(), db); err != nil {
		t.Fatal(err)
	}

	releaseFailure := errors.New("release unavailable")
	rollbackFailure := errors.New("rollback unavailable")
	mock.ExpectRollback().WillReturnError(rollbackFailure)
	mock.ExpectQuery(regexp.QuoteMeta(
		"SELECT RELEASE_LOCK('" + catalogsystemfence.Name + "')",
	)).WillReturnError(releaseFailure)
	err := token.Close()
	if !errors.Is(err, releaseFailure) || !errors.Is(err, rollbackFailure) {
		t.Fatalf("Close = %v, want both release and rollback failures", err)
	}
	assertVerificationExpectations(t, mock)
}

func TestSystemAuthorityFenceRejectsForeignAndClosedTokensWithoutSQL(t *testing.T) {
	db, mock, token := newOwnedVerificationToken(t)
	defer db.Close()
	foreign, _, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer foreign.Close()

	if err := token.acquireSystemAuthorityFence(
		context.Background(),
		foreign,
	); !errors.Is(err, ErrAuthorityMismatch) {
		t.Fatalf("foreign acquire = %v, want ErrAuthorityMismatch", err)
	}
	mock.ExpectRollback()
	expectSystemSessionRestore(mock)
	if err := token.Close(); err != nil {
		t.Fatal(err)
	}
	if err := token.acquireSystemAuthorityFence(
		context.Background(),
		db,
	); !errors.Is(err, ErrClosed) {
		t.Fatalf("closed acquire = %v, want ErrClosed", err)
	}
	var nilToken *Token
	if err := nilToken.acquireSystemAuthorityFence(
		context.Background(),
		db,
	); !errors.Is(err, ErrClosed) {
		t.Fatalf("nil acquire = %v, want ErrClosed", err)
	}
	assertVerificationExpectations(t, mock)
}

func TestLegacyTokenRejectsSystemFenceWithoutOwnedSession(t *testing.T) {
	db, mock, token := newVerificationToken(t)
	defer db.Close()
	if err := token.acquireSystemAuthorityFence(
		context.Background(),
		db,
	); !errors.Is(err, ErrAuthorityMismatch) {
		t.Fatalf("legacy acquire = %v, want ErrAuthorityMismatch", err)
	}
	mock.ExpectRollback()
	if err := token.Close(); err != nil {
		t.Fatal(err)
	}
	assertVerificationExpectations(t, mock)
}

func TestAmbiguousSystemFenceAcquireInvalidatesAndDiscardsOwnedSession(
	t *testing.T,
) {
	db, mock, token := newOwnedVerificationToken(t)
	defer db.Close()
	acquireFailure := errors.New("GET_LOCK response lost")
	mock.ExpectQuery(regexp.QuoteMeta(
		"SELECT GET_LOCK('" + catalogsystemfence.Name + "', 600)",
	)).WillReturnError(acquireFailure)
	mock.ExpectRollback()

	err := token.acquireSystemAuthorityFence(context.Background(), db)
	if !errors.Is(err, acquireFailure) ||
		!errors.Is(err, catalogsystemfence.ErrStateUnknown) {
		t.Fatalf("ambiguous acquire = %v", err)
	}
	if token.ActiveFor(db) {
		t.Fatal("ambiguous acquire left the owned token active")
	}
	assertVerificationExpectations(t, mock)
}

func TestAmbiguousSystemFenceSessionBindingInvalidatesAndDiscardsOwnedSession(
	t *testing.T,
) {
	db, mock, token := newOwnedVerificationToken(t)
	defer db.Close()
	mock.ExpectQuery(regexp.QuoteMeta(
		"SELECT GET_LOCK('" + catalogsystemfence.Name + "', 600)",
	)).WillReturnRows(sqlmock.NewRows([]string{"GET_LOCK"}).AddRow(1))
	bindFailure := errors.New("CONNECTION_ID response lost")
	mock.ExpectQuery(regexp.QuoteMeta(
		"SELECT CONNECTION_ID()",
	)).WillReturnError(bindFailure)
	mock.ExpectRollback()

	err := token.acquireSystemAuthorityFence(context.Background(), db)
	if !errors.Is(err, bindFailure) ||
		!errors.Is(err, catalogsystemfence.ErrStateUnknown) {
		t.Fatalf("ambiguous session binding = %v", err)
	}
	if token.ActiveFor(db) {
		t.Fatal("ambiguous session binding left the owned token active")
	}
	if open := db.Stats().OpenConnections; open != 0 {
		t.Fatalf("ambiguous session binding retained %d pooled connections", open)
	}
	assertVerificationExpectations(t, mock)
}

func TestSystemAuthorityFenceLivenessFailureInvalidatesAndDiscardsSession(
	t *testing.T,
) {
	db, mock, token := newOwnedVerificationToken(t)
	defer db.Close()
	expectSystemFenceAcquire(mock, 1)
	if err := token.acquireSystemAuthorityFence(context.Background(), db); err != nil {
		t.Fatal(err)
	}

	sessionLost := errors.New("retained session dropped")
	mock.ExpectQuery(regexp.QuoteMeta(
		"SELECT IS_USED_LOCK('" + catalogsystemfence.Name + "'), CONNECTION_ID()",
	)).WillReturnError(sessionLost)
	mock.ExpectRollback()
	if token.systemAuthorityFenceActive() {
		t.Fatal("lost system authority session remained active")
	}
	if token.ActiveFor(db) {
		t.Fatal("lost system authority session retained an active token")
	}
	if err := token.Close(); err != nil {
		t.Fatalf("close invalidated token = %v", err)
	}
	if open := db.Stats().OpenConnections; open != 0 {
		t.Fatalf("lost authority session retained %d pooled connections", open)
	}
	assertVerificationExpectations(t, mock)
}

func newOwnedVerificationToken(
	t *testing.T,
) (*sql.DB, sqlmock.Sqlmock, *Token) {
	t.Helper()
	db, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	expectSystemSessionPessimistic(mock)
	conn, err := db.Conn(context.Background())
	if err != nil {
		db.Close()
		t.Fatal(err)
	}
	mode, err := catalogsystemfence.ForcePessimistic(context.Background(), conn)
	if err != nil {
		conn.Close()
		db.Close()
		t.Fatal(err)
	}
	mock.ExpectBegin()
	tx, err := conn.BeginTx(context.WithoutCancel(context.Background()), nil)
	if err != nil {
		conn.Close()
		db.Close()
		t.Fatal(err)
	}
	token, err := NewOwned(db, conn, tx, mode)
	if err != nil {
		tx.Rollback()
		conn.Close()
		db.Close()
		t.Fatal(err)
	}
	return db, mock, token
}

func expectSystemSessionPessimistic(mock sqlmock.Sqlmock) {
	mock.ExpectQuery(regexp.QuoteMeta(
		"SELECT @@SESSION.tidb_txn_mode",
	)).WillReturnRows(
		sqlmock.NewRows([]string{"tidb_txn_mode"}).AddRow("optimistic"),
	)
	mock.ExpectExec(regexp.QuoteMeta(
		"SET SESSION tidb_txn_mode = 'pessimistic'",
	)).WillReturnResult(sqlmock.NewResult(0, 0))
}

func expectSystemSessionRestore(mock sqlmock.Sqlmock) {
	mock.ExpectExec(regexp.QuoteMeta(
		"SET SESSION tidb_txn_mode = ?",
	)).WithArgs("optimistic").WillReturnResult(sqlmock.NewResult(0, 0))
}

func expectSystemFenceAcquire(mock sqlmock.Sqlmock, result int64) {
	mock.ExpectQuery(regexp.QuoteMeta(
		"SELECT GET_LOCK('" + catalogsystemfence.Name + "', 600)",
	)).WillReturnRows(sqlmock.NewRows([]string{"GET_LOCK"}).AddRow(result))
	if result == 1 {
		mock.ExpectQuery(regexp.QuoteMeta(
			"SELECT CONNECTION_ID()",
		)).WillReturnRows(
			sqlmock.NewRows([]string{"CONNECTION_ID"}).AddRow(73),
		)
	}
}

func expectSystemFenceOwnership(mock sqlmock.Sqlmock) {
	mock.ExpectQuery(regexp.QuoteMeta(
		"SELECT IS_USED_LOCK('" + catalogsystemfence.Name + "'), CONNECTION_ID()",
	)).WillReturnRows(
		sqlmock.NewRows([]string{"IS_USED_LOCK", "CONNECTION_ID"}).
			AddRow(73, 73),
	)
}

func expectSystemFenceRelease(mock sqlmock.Sqlmock, result int64) {
	mock.ExpectQuery(regexp.QuoteMeta(
		"SELECT RELEASE_LOCK('" + catalogsystemfence.Name + "')",
	)).WillReturnRows(sqlmock.NewRows([]string{"RELEASE_LOCK"}).AddRow(result))
}