package catalogwriter

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

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/go-sql-driver/mysql"
	"github.com/pingcap/tidb/pkg/parser"
	_ "github.com/pingcap/tidb/pkg/parser/test_driver"
)

const lockQuery = `SELECT OrgId, RuntimeAudience, CatalogRevision, LastSecurityEpoch, ActiveOperationId, WriterGeneration
FROM _System_CatalogOrgState WHERE OrgId = ? FOR UPDATE`

func TestSchemaOwnsOnlySharedOrgWriterState(t *testing.T) {
	joined := strings.Join(SchemaStatementsForTest(), "\n")
	for _, required := range []string{
		"_System_CatalogOrgState",
		"CatalogRevision",
		"LastSecurityEpoch",
		"ActiveOperationId",
		"WriterGeneration",
		"ADD COLUMN IF NOT EXISTS",
	} {
		if !strings.Contains(joined, required) {
			t.Fatalf("writer schema is missing %q", required)
		}
	}
	for _, forbidden := range []string{"AppInstallationId", "AppId", "Layout", "Profile", "Role"} {
		if strings.Contains(joined, forbidden) {
			t.Fatalf("writer schema leaked %q", forbidden)
		}
	}
}

func TestWriterGenerationSchemaParsesWithTiDB(t *testing.T) {
	p := parser.New()
	for _, statement := range SchemaStatementsForTest() {
		if _, _, err := p.Parse(statement, "", ""); err != nil {
			t.Fatalf("writer schema does not parse with TiDB: %v\n%s", err, statement)
		}
	}
}

func TestLockOrCreateReturnsExistingLockedState(t *testing.T) {
	db, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	mock.ExpectBegin()
	tx, err := db.BeginTx(context.Background(), nil)
	if err != nil {
		t.Fatal(err)
	}
	mock.ExpectQuery(regexp.QuoteMeta(lockQuery)).WithArgs("org-1").WillReturnRows(
		sqlmock.NewRows([]string{
			"OrgId", "RuntimeAudience", "CatalogRevision", "LastSecurityEpoch",
			"ActiveOperationId", "WriterGeneration",
		}).AddRow("org-1", "runtime", []byte("7"), []byte("11"), "command-1", []byte("5")))

	state, err := LockOrCreate(context.Background(), tx, "org-1", "runtime")
	if err != nil {
		t.Fatal(err)
	}
	want := State{
		OrgID: "org-1", RuntimeAudience: "runtime", CatalogRevision: 7,
		SecurityEpoch: 11, ActiveCommandID: "command-1", WriterGeneration: 5,
	}
	if state != want {
		t.Fatalf("state = %#v, want %#v", state, want)
	}
	mock.ExpectRollback()
	_ = tx.Rollback()
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestLockOrCreateInitializesZeroState(t *testing.T) {
	db, mock, tx := beginMockTx(t)
	defer db.Close()
	mock.ExpectQuery(regexp.QuoteMeta(lockQuery)).WithArgs("org-new").WillReturnError(sql.ErrNoRows)
	mock.ExpectExec(regexp.QuoteMeta(`INSERT INTO _System_CatalogOrgState
(OrgId, RuntimeAudience, CatalogRevision, LastSecurityEpoch, ActiveOperationId, WriterGeneration, CreatedDate, LastModifiedDate)
VALUES (?, ?, 0, 0, NULL, 0, CURRENT_TIMESTAMP(6), CURRENT_TIMESTAMP(6))`)).
		WithArgs("org-new", "runtime").WillReturnResult(sqlmock.NewResult(1, 1))

	state, err := LockOrCreate(context.Background(), tx, "org-new", "runtime")
	if err != nil {
		t.Fatal(err)
	}
	if state != (State{OrgID: "org-new", RuntimeAudience: "runtime"}) {
		t.Fatalf("new state = %#v", state)
	}
	mock.ExpectRollback()
	_ = tx.Rollback()
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestLockOrCreateRejectsAudienceDriftAndDuplicateInitialization(t *testing.T) {
	t.Run("collation alias", func(t *testing.T) {
		db, mock, tx := beginMockTx(t)
		defer db.Close()
		mock.ExpectQuery(regexp.QuoteMeta(lockQuery)).WithArgs("org-1").WillReturnRows(
			sqlmock.NewRows([]string{
				"OrgId", "RuntimeAudience", "CatalogRevision", "LastSecurityEpoch",
				"ActiveOperationId", "WriterGeneration",
			}).AddRow("ORG-1", "runtime", []byte("0"), []byte("0"), nil, []byte("0")))
		if _, err := LockOrCreate(context.Background(), tx, "org-1", "runtime"); !errors.Is(err, ErrConflict) {
			t.Fatalf("byte-different org alias = %v", err)
		}
		mock.ExpectRollback()
		_ = tx.Rollback()
		if err := mock.ExpectationsWereMet(); err != nil {
			t.Fatal(err)
		}
	})
	t.Run("audience drift", func(t *testing.T) {
		db, mock, tx := beginMockTx(t)
		defer db.Close()
		mock.ExpectQuery(regexp.QuoteMeta(lockQuery)).WithArgs("org-1").WillReturnRows(
			sqlmock.NewRows([]string{
				"OrgId", "RuntimeAudience", "CatalogRevision", "LastSecurityEpoch",
				"ActiveOperationId", "WriterGeneration",
			}).AddRow("org-1", "runtime-a", []byte("0"), []byte("0"), nil, []byte("0")))
		if _, err := LockOrCreate(context.Background(), tx, "org-1", "runtime-b"); !errors.Is(err, ErrConflict) {
			t.Fatalf("audience drift = %v", err)
		}
		mock.ExpectRollback()
		_ = tx.Rollback()
		if err := mock.ExpectationsWereMet(); err != nil {
			t.Fatal(err)
		}
	})
	t.Run("duplicate initialization", func(t *testing.T) {
		db, mock, tx := beginMockTx(t)
		defer db.Close()
		mock.ExpectQuery(regexp.QuoteMeta(lockQuery)).WithArgs("org-1").WillReturnError(sql.ErrNoRows)
		mock.ExpectExec("INSERT INTO _System_CatalogOrgState").WithArgs("org-1", "runtime").
			WillReturnError(&mysql.MySQLError{Number: 1062, Message: "duplicate"})
		if _, err := LockOrCreate(context.Background(), tx, "org-1", "runtime"); !errors.Is(err, ErrConflict) {
			t.Fatalf("duplicate initialization = %v", err)
		}
		mock.ExpectRollback()
		_ = tx.Rollback()
		if err := mock.ExpectationsWereMet(); err != nil {
			t.Fatal(err)
		}
	})
}

func TestReserveAndCommitUseOneExactCommandFence(t *testing.T) {
	db, mock, tx := beginMockTx(t)
	defer db.Close()
	state := State{
		OrgID: "org-1", RuntimeAudience: "runtime", CatalogRevision: 7,
		SecurityEpoch: 11, WriterGeneration: 5,
	}
	mock.ExpectExec(regexp.QuoteMeta(`UPDATE _System_CatalogOrgState
SET ActiveOperationId = ?, WriterGeneration = ?, LastModifiedDate = CURRENT_TIMESTAMP(6)
WHERE OrgId = ? AND ActiveOperationId IS NULL AND CatalogRevision = ? AND LastSecurityEpoch = ? AND WriterGeneration = ?`)).
		WithArgs("command-1", "6", "org-1", "7", "11", "5").
		WillReturnResult(sqlmock.NewResult(0, 1))
	generation, err := ReserveGeneration(context.Background(), tx, state, "command-1")
	if err != nil {
		t.Fatal(err)
	}
	if generation != 6 {
		t.Fatalf("generation = %d, want 6", generation)
	}
	mock.ExpectExec(regexp.QuoteMeta(`UPDATE _System_CatalogOrgState
SET CatalogRevision = ?, LastSecurityEpoch = ?, ActiveOperationId = NULL, LastModifiedDate = CURRENT_TIMESTAMP(6)
WHERE OrgId = ? AND ActiveOperationId = ?`)).
		WithArgs("8", "12", "org-1", "command-1").WillReturnResult(sqlmock.NewResult(0, 1))
	if err := Commit(context.Background(), tx, "org-1", "command-1", 8, 12); err != nil {
		t.Fatal(err)
	}
	mock.ExpectRollback()
	_ = tx.Rollback()
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestReserveGenerationRejectsExhaustionBeforeMutation(t *testing.T) {
	db, mock, tx := beginMockTx(t)
	defer db.Close()
	state := State{
		OrgID: "org-1", RuntimeAudience: "runtime",
		WriterGeneration: ^uint64(0),
	}
	if _, err := ReserveGeneration(
		context.Background(), tx, state, "command-1",
	); !errors.Is(err, ErrLineageExhausted) {
		t.Fatalf("exhausted reservation generation = %v", err)
	}
	mock.ExpectRollback()
	_ = tx.Rollback()
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestReleaseClearsExactCommandFenceWithoutChangingLineage(t *testing.T) {
	db, mock, tx := beginMockTx(t)
	defer db.Close()
	mock.ExpectExec(regexp.QuoteMeta(`UPDATE _System_CatalogOrgState
SET ActiveOperationId = NULL, LastModifiedDate = CURRENT_TIMESTAMP(6)
WHERE OrgId = ? AND ActiveOperationId = ?`)).
		WithArgs("org-1", "command-1").WillReturnResult(sqlmock.NewResult(0, 1))

	if err := Release(context.Background(), tx, "org-1", "command-1"); err != nil {
		t.Fatal(err)
	}
	mock.ExpectRollback()
	_ = tx.Rollback()
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestReleaseGenerationClearsOnlyExactReservationIncarnation(t *testing.T) {
	db, mock, tx := beginMockTx(t)
	defer db.Close()
	mock.ExpectExec(regexp.QuoteMeta(`UPDATE _System_CatalogOrgState
SET ActiveOperationId = NULL, LastModifiedDate = CURRENT_TIMESTAMP(6)
WHERE OrgId = ? AND ActiveOperationId = ? AND WriterGeneration = ?`)).
		WithArgs("org-1", "command-1", "7").
		WillReturnResult(sqlmock.NewResult(0, 1))

	if err := ReleaseGeneration(
		context.Background(), tx, "org-1", "command-1", 7,
	); err != nil {
		t.Fatal(err)
	}
	mock.ExpectRollback()
	_ = tx.Rollback()
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatal(err)
	}
}

func TestReleaseRequiresExactFenceAndFailsClosedWhenItIsLost(t *testing.T) {
	t.Run("identity required", func(t *testing.T) {
		for _, input := range []struct {
			orgID     string
			commandID string
		}{
			{orgID: "", commandID: "command-1"},
			{orgID: "org-1", commandID: ""},
			{orgID: "  ", commandID: "command-1"},
			{orgID: "org-1", commandID: "  "},
		} {
			db, mock, tx := beginMockTx(t)
			if err := Release(context.Background(), tx, input.orgID, input.commandID); err == nil {
				t.Fatal("release accepted an incomplete fence identity")
			}
			mock.ExpectRollback()
			_ = tx.Rollback()
			if err := mock.ExpectationsWereMet(); err != nil {
				t.Fatal(err)
			}
			_ = db.Close()
		}
	})

	t.Run("lost fence", func(t *testing.T) {
		db, mock, tx := beginMockTx(t)
		defer db.Close()
		mock.ExpectExec(regexp.QuoteMeta(`UPDATE _System_CatalogOrgState
SET ActiveOperationId = NULL, LastModifiedDate = CURRENT_TIMESTAMP(6)
WHERE OrgId = ? AND ActiveOperationId = ?`)).
			WithArgs("org-1", "command-1").WillReturnResult(sqlmock.NewResult(0, 0))
		if err := Release(context.Background(), tx, "org-1", "command-1"); !errors.Is(err, ErrConflict) {
			t.Fatalf("lost fence = %v", err)
		}
		mock.ExpectRollback()
		_ = tx.Rollback()
		if err := mock.ExpectationsWereMet(); err != nil {
			t.Fatal(err)
		}
	})
}

func TestLineageAllocationAxesAreIndependent(t *testing.T) {
	state := State{CatalogRevision: 7, SecurityEpoch: ^uint64(0)}
	nextCatalog, err := NextCatalogRevision(state)
	if err != nil || nextCatalog != 8 {
		t.Fatalf("next Catalog revision = %d, %v", nextCatalog, err)
	}
	if _, err := NextSecurityEpoch(state); !errors.Is(err, ErrLineageExhausted) {
		t.Fatalf("exhausted security epoch = %v", err)
	}

	state = State{CatalogRevision: ^uint64(0), SecurityEpoch: 11}
	if _, err := NextCatalogRevision(state); !errors.Is(err, ErrLineageExhausted) {
		t.Fatalf("exhausted Catalog revision = %v", err)
	}
	nextSecurity, err := NextSecurityEpoch(state)
	if err != nil || nextSecurity != 12 {
		t.Fatalf("next security epoch = %d, %v", nextSecurity, err)
	}
}

func TestReserveAndCommitFailClosedOnLostFence(t *testing.T) {
	for name, run := range map[string]func(context.Context, *sql.Tx) error{
		"reserve": func(ctx context.Context, tx *sql.Tx) error {
			return Reserve(ctx, tx, State{OrgID: "org-1", CatalogRevision: 7, SecurityEpoch: 11}, "command-1")
		},
		"commit": func(ctx context.Context, tx *sql.Tx) error {
			return Commit(ctx, tx, "org-1", "command-1", 8, 12)
		},
	} {
		t.Run(name, func(t *testing.T) {
			db, mock, tx := beginMockTx(t)
			defer db.Close()
			mock.ExpectExec("UPDATE _System_CatalogOrgState").WillReturnResult(sqlmock.NewResult(0, 0))
			if err := run(context.Background(), tx); !errors.Is(err, ErrConflict) {
				t.Fatalf("lost fence = %v", err)
			}
			mock.ExpectRollback()
			_ = tx.Rollback()
			if err := mock.ExpectationsWereMet(); err != nil {
				t.Fatal(err)
			}
		})
	}
}

func beginMockTx(t *testing.T) (*sql.DB, sqlmock.Sqlmock, *sql.Tx) {
	t.Helper()
	db, mock, err := sqlmock.New()
	if err != nil {
		t.Fatal(err)
	}
	mock.ExpectBegin()
	tx, err := db.BeginTx(context.Background(), nil)
	if err != nil {
		t.Fatal(err)
	}
	return db, mock, tx
}