package catalogsystemfence

import (
	"context"
	"errors"
	"regexp"
	"testing"

	"github.com/DATA-DOG/go-sqlmock"
)

func TestFenceAcquiresAndReleasesOnCallerOwnedSession(t *testing.T) {
	db, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	conn, err := db.Conn(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()

	expectAcquire(mock, 1)
	fence, err := AcquireOnConn(context.Background(), conn)
	if err != nil {
		t.Fatal(err)
	}
	expectRelease(mock, 1)
	if err := fence.ReleaseOnConn(context.Background(), conn); err != nil {
		t.Fatal(err)
	}
	if err := fence.ReleaseOnConn(context.Background(), conn); err != nil {
		t.Fatalf("second release = %v", err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestFenceAcquiresInsideTransactionAndReleasesAfterRollbackOnSameConn(
	t *testing.T,
) {
	db, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	conn, err := db.Conn(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	mock.ExpectBegin()
	tx, err := conn.BeginTx(context.Background(), nil)
	if err != nil {
		t.Fatal(err)
	}
	expectAcquire(mock, 1)
	expectTransactionSession(mock, 73)
	fence, err := AcquireOnTx(context.Background(), tx)
	if err != nil {
		t.Fatal(err)
	}
	expectOwnershipProbe(mock, 73, 73)
	if err := fence.ValidateOnTx(context.Background(), tx); err != nil {
		t.Fatal(err)
	}
	mock.ExpectRollback()
	if err := tx.Rollback(); err != nil {
		t.Fatal(err)
	}
	expectRelease(mock, 1)
	if err := fence.ReleaseOnConn(context.Background(), conn); err != nil {
		t.Fatal(err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestFenceOwnershipProbeFailsClosedForEveryNonOwnedResult(
	t *testing.T,
) {
	for _, test := range []struct {
		name             string
		ownerSessionID   any
		currentSessionID int64
		want             error
	}{
		{
			name:             "another session owns the lock",
			ownerSessionID:   int64(74),
			currentSessionID: 73,
			want:             ErrOwnershipLost,
		},
		{
			name:             "the lock is free",
			ownerSessionID:   nil,
			currentSessionID: 73,
			want:             ErrOwnershipLost,
		},
		{
			name:             "the transaction session changed",
			ownerSessionID:   int64(73),
			currentSessionID: 74,
			want:             ErrOwnershipLost,
		},
		{
			name:             "the owner result is zero",
			ownerSessionID:   int64(0),
			currentSessionID: 73,
			want:             ErrOwnershipLost,
		},
		{
			name:             "the owner result is malformed",
			ownerSessionID:   "not-a-session-id",
			currentSessionID: 73,
			want:             ErrStateUnknown,
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			db, mock, err := sqlmock.New()
			if err != nil {
				t.Fatal(err)
			}
			defer db.Close()
			conn, err := db.Conn(context.Background())
			if err != nil {
				t.Fatal(err)
			}
			defer conn.Close()
			mock.ExpectBegin()
			tx, err := conn.BeginTx(context.Background(), nil)
			if err != nil {
				t.Fatal(err)
			}
			expectAcquire(mock, 1)
			expectTransactionSession(mock, 73)
			fence, err := AcquireOnTx(context.Background(), tx)
			if err != nil {
				t.Fatal(err)
			}
			expectOwnershipProbe(
				mock,
				test.ownerSessionID,
				test.currentSessionID,
			)
			if err := fence.ValidateOnTx(
				context.Background(),
				tx,
			); !errors.Is(err, test.want) {
				t.Fatalf("ownership probe = %v, want %v", err, test.want)
			}
			if err := fence.ValidateOnTx(
				context.Background(),
				tx,
			); !errors.Is(err, ErrClosed) {
				t.Fatalf("second ownership probe = %v, want ErrClosed", err)
			}
			mock.ExpectRollback()
			if err := tx.Rollback(); err != nil {
				t.Fatal(err)
			}
			if err := mock.ExpectationsWereMet(); err != nil {
				t.Fatal(err)
			}
		})
	}
}

func TestAcquireRejectsAnUnavailableFence(t *testing.T) {
	db, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	conn, err := db.Conn(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	expectAcquire(mock, 0)
	if _, err := AcquireOnConn(
		context.Background(),
		conn,
	); !errors.Is(err, ErrNotAcquired) {
		t.Fatalf("AcquireOnConn = %v, want ErrNotAcquired", err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestAcquireMarksAnAmbiguousResponseUnknown(t *testing.T) {
	db, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	conn, err := db.Conn(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	queryErr := errors.New("response lost")
	mock.ExpectQuery(regexp.QuoteMeta(
		"SELECT GET_LOCK('" + Name + "', 600)",
	)).WillReturnError(queryErr)
	if _, err := AcquireOnConn(
		context.Background(),
		conn,
	); !errors.Is(err, ErrStateUnknown) || !errors.Is(err, queryErr) {
		t.Fatalf("AcquireOnConn = %v, want ambiguous query failure", err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestReleaseRejectsAnUnreleasedFence(t *testing.T) {
	db, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	conn, err := db.Conn(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	expectAcquire(mock, 1)
	fence, err := AcquireOnConn(context.Background(), conn)
	if err != nil {
		t.Fatal(err)
	}
	expectRelease(mock, 0)
	if err := fence.ReleaseOnConn(
		context.Background(),
		conn,
	); !errors.Is(err, ErrNotReleased) ||
		!errors.Is(err, ErrStateUnknown) {
		t.Fatalf("ReleaseOnConn = %v, want ambiguous ErrNotReleased", err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestPessimisticModeGuardRestoresExactPriorMode(t *testing.T) {
	db, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	conn, err := db.Conn(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	expectPessimistic(mock, "optimistic")
	guard, err := ForcePessimistic(context.Background(), conn)
	if err != nil {
		t.Fatal(err)
	}
	mock.ExpectExec(regexp.QuoteMeta(
		"SET SESSION tidb_txn_mode = ?",
	)).WithArgs("optimistic").WillReturnResult(sqlmock.NewResult(0, 0))
	if err := guard.Restore(context.Background(), conn); err != nil {
		t.Fatal(err)
	}
	if err := guard.Restore(context.Background(), conn); err != nil {
		t.Fatalf("second restore = %v", err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestPessimisticModeGuardDiscardsAmbiguousRestore(t *testing.T) {
	db, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	conn, err := db.Conn(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	expectPessimistic(mock, "")
	guard, err := ForcePessimistic(context.Background(), conn)
	if err != nil {
		t.Fatal(err)
	}
	restoreErr := errors.New("restore response lost")
	mock.ExpectExec(regexp.QuoteMeta(
		"SET SESSION tidb_txn_mode = ?",
	)).WithArgs("").WillReturnError(restoreErr)
	if err := guard.Restore(
		context.Background(),
		conn,
	); !errors.Is(err, restoreErr) {
		t.Fatalf("Restore = %v, want restore failure", err)
	}
	_ = conn.Close()
	if open := db.Stats().OpenConnections; open != 0 {
		t.Fatalf("ambiguous restore retained %d pooled connections, want 0", open)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestFenceOperationsRequireOwnedSessionObjects(t *testing.T) {
	if _, err := AcquireOnConn(context.Background(), nil); err == nil {
		t.Fatal("nil connection was accepted")
	}
	if _, err := AcquireOnTx(context.Background(), nil); err == nil {
		t.Fatal("nil transaction was accepted")
	}
	fence := &Fence{active: true}
	if err := fence.ValidateOnTx(
		context.Background(),
		nil,
	); !errors.Is(err, ErrClosed) {
		t.Fatalf("nil validation transaction = %v, want ErrClosed", err)
	}
	if err := fence.ReleaseOnConn(
		context.Background(),
		nil,
	); !errors.Is(err, ErrClosed) {
		t.Fatalf("nil release connection = %v, want ErrClosed", err)
	}
}

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

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

func expectTransactionSession(mock sqlmock.Sqlmock, sessionID int64) {
	mock.ExpectQuery(regexp.QuoteMeta(
		"SELECT CONNECTION_ID()",
	)).WillReturnRows(
		sqlmock.NewRows([]string{"CONNECTION_ID"}).AddRow(sessionID),
	)
}

func expectOwnershipProbe(
	mock sqlmock.Sqlmock,
	ownerSessionID any,
	currentSessionID int64,
) {
	mock.ExpectQuery(regexp.QuoteMeta(
		"SELECT IS_USED_LOCK('" + Name + "'), CONNECTION_ID()",
	)).WillReturnRows(
		sqlmock.NewRows([]string{"IS_USED_LOCK", "CONNECTION_ID"}).
			AddRow(ownerSessionID, currentSessionID),
	)
}

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