package serving

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

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

const (
	accountGateTestID      = "account-gate-1"
	accountGateTestAccount = "connector-account-1"
	accountGateTestEpochID = "serving-epoch-1"
)

var (
	accountGateTestScope       = strings.Repeat("a", 64)
	accountGateTestEpochDigest = "sha256:" + strings.Repeat("b", 64)
)

func TestAccountServingGateStoredStateIsClosed(t *testing.T) {
	for _, mode := range []accountServingMode{
		accountServingModeV1Serving,
		accountServingModeHaltingV1,
		accountServingModeV1Stopped,
		accountServingModeV1Drained,
		accountServingModeV2Validating,
		accountServingModeV2Serving,
		accountServingModeHaltingV2,
		accountServingModeV2Halted,
	} {
		t.Run(string(mode), func(t *testing.T) {
			gate := validAccountServingGate(mode)
			if err := gate.validateStored(); err != nil {
				t.Fatalf("validateStored(%s): %v", mode, err)
			}
		})
	}

	tests := []struct {
		name   string
		mutate func(*accountServingGate)
	}{
		{name: "noncanonical gate ID", mutate: func(g *accountServingGate) { g.id = " gate" }},
		{name: "noncanonical scope", mutate: func(g *accountServingGate) { g.servingScopeDigest = strings.ToUpper(accountGateTestScope) }},
		{name: "noncanonical account", mutate: func(g *accountServingGate) { g.connectorAccountID = "account\x00" }},
		{name: "negative fence", mutate: func(g *accountServingGate) { g.fenceToken = -1 }},
		{name: "unknown mode", mutate: func(g *accountServingGate) { g.mode = "v2_maybe" }},
		{name: "zero transitioned fence", mutate: func(g *accountServingGate) { g.fenceToken = 0 }},
		{name: "partial epoch ID", mutate: func(g *accountServingGate) { g.servingEpochDigest = nil }},
		{name: "partial epoch digest", mutate: func(g *accountServingGate) { g.servingEpochID = nil }},
		{name: "missing active epoch", mutate: func(g *accountServingGate) { g.servingEpochID, g.servingEpochDigest = nil, nil }},
		{name: "invalid epoch ID", mutate: func(g *accountServingGate) { value := " epoch"; g.servingEpochID = &value }},
		{name: "invalid epoch digest", mutate: func(g *accountServingGate) {
			value := "sha256:" + strings.Repeat("B", 64)
			g.servingEpochDigest = &value
		}},
		{name: "epoch on V1 mode", mutate: func(g *accountServingGate) { g.mode = accountServingModeV1Serving }},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			gate := validAccountServingGate(accountServingModeV2Serving)
			test.mutate(&gate)
			if err := gate.validateStored(); !errors.Is(err, errStoredAccountServingGateInvalid) {
				t.Fatalf("validateStored() error = %v, want stored gate invalid", err)
			}
		})
	}
}

func TestLoadAccountServingGateExactUsesIndexedBinaryCoordinate(t *testing.T) {
	db, mock := newServingSchemaMock(t)
	gate := validAccountServingGate(accountServingModeV2Serving)
	expectAccountServingGateRow(mock, gate, false)

	loaded, found, err := loadAccountServingGateExact(
		t.Context(), db, accountGateTestScope, accountGateTestAccount, false,
	)
	if err != nil || !found || loaded.id != gate.id || loaded.fenceToken != gate.fenceToken {
		t.Fatalf("load exact gate = %#v, found=%t, err=%v", loaded, found, err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatalf("sql expectations: %v", err)
	}
}

func TestLoadAccountServingGateExactRejectsInvalidAndCollationAliasedRows(t *testing.T) {
	t.Run("input before SQL", func(t *testing.T) {
		db, mock := newServingSchemaMock(t)
		if _, _, err := loadAccountServingGateExact(t.Context(), db, "scope", accountGateTestAccount, false); !errors.Is(err, errInvalidAccountServingGateInput) {
			t.Fatalf("invalid scope error = %v", err)
		}
		if _, _, err := loadAccountServingGateExact(t.Context(), db, accountGateTestScope, " account", false); !errors.Is(err, errInvalidAccountServingGateInput) {
			t.Fatalf("invalid account error = %v", err)
		}
		if err := mock.ExpectationsWereMet(); err != nil {
			t.Fatalf("unexpected SQL: %v", err)
		}
	})

	t.Run("returned coordinate must be byte exact", func(t *testing.T) {
		db, mock := newServingSchemaMock(t)
		gate := validAccountServingGate(accountServingModeV2Serving)
		gate.connectorAccountID = "Connector-Account-1"
		expectAccountServingGateRowWithArgs(mock, gate, false, accountGateTestScope, accountGateTestAccount)
		if _, _, err := loadAccountServingGateExact(t.Context(), db, accountGateTestScope, accountGateTestAccount, false); !errors.Is(err, errStoredAccountServingGateInvalid) {
			t.Fatalf("aliased row error = %v, want stored gate invalid", err)
		}
	})
}

func TestLockExpectedV2ServingLocksGateThenExactActiveEpoch(t *testing.T) {
	db, mock := newServingSchemaMock(t)
	mock.ExpectBegin()
	tx, err := db.BeginTx(t.Context(), nil)
	if err != nil {
		t.Fatalf("begin: %v", err)
	}
	gate := validAccountServingGate(accountServingModeV2Serving)
	header := validAccountServingEpochHeader()
	expectAccountServingGateRow(mock, gate, true)
	expectAccountServingEpochHeaderRow(mock, header, true)
	mock.ExpectRollback()

	locked, err := lockExpectedV2Serving(
		t.Context(), tx, accountGateTestScope, accountGateTestAccount,
		gate.fenceToken, accountGateTestEpochID, accountGateTestEpochDigest,
	)
	if err != nil || locked.id != gate.id {
		t.Fatalf("lockExpectedV2Serving() = %#v, %v", locked, err)
	}
	if err := tx.Rollback(); err != nil {
		t.Fatalf("rollback: %v", err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatalf("sql expectations: %v", err)
	}
}

func TestLockExpectedV2ServingRejectsStaleOrNonServingObservations(t *testing.T) {
	tests := []struct {
		name        string
		mutateGate  func(*accountServingGate)
		mutateEpoch func(*accountServingEpochHeader)
		expectEpoch bool
		want        error
	}{
		{name: "stale fence", mutateGate: func(g *accountServingGate) { g.fenceToken++ }, want: errAccountServingGateFenceConflict},
		{name: "wrong mode", mutateGate: func(g *accountServingGate) { g.mode = accountServingModeHaltingV2 }, want: errAccountServingGateModeConflict},
		{name: "wrong gate epoch", mutateGate: func(g *accountServingGate) { value := "other-epoch"; g.servingEpochID = &value }, want: errAccountServingGateEpochConflict},
		{name: "inactive epoch", expectEpoch: true, mutateEpoch: func(h *accountServingEpochHeader) { h.state = verifiedState; h.activeScopeDigest = nil }, want: errAccountServingGateEpochConflict},
		{name: "wrong epoch digest", expectEpoch: true, mutateEpoch: func(h *accountServingEpochHeader) { h.epochDigest = "sha256:" + strings.Repeat("c", 64) }, want: errAccountServingGateEpochConflict},
		{name: "wrong active scope", expectEpoch: true, mutateEpoch: func(h *accountServingEpochHeader) {
			value := strings.Repeat("d", 64)
			h.servingScopeDigest, h.activeScopeDigest = value, &value
		}, want: errAccountServingGateEpochConflict},
		{name: "corrupt epoch state", expectEpoch: true, mutateEpoch: func(h *accountServingEpochHeader) { h.state = "activated" }, want: ErrStoredCandidateInvalid},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			db, mock := newServingSchemaMock(t)
			mock.ExpectBegin()
			tx, err := db.BeginTx(t.Context(), nil)
			if err != nil {
				t.Fatalf("begin: %v", err)
			}
			gate := validAccountServingGate(accountServingModeV2Serving)
			gate.fenceToken = 41
			if test.mutateGate != nil {
				test.mutateGate(&gate)
			}
			expectAccountServingGateRow(mock, gate, true)
			if test.expectEpoch {
				header := validAccountServingEpochHeader()
				if test.mutateEpoch != nil {
					test.mutateEpoch(&header)
				}
				expectAccountServingEpochHeaderRow(mock, header, true)
			}
			mock.ExpectRollback()

			_, err = lockExpectedV2Serving(
				t.Context(), tx, accountGateTestScope, accountGateTestAccount,
				41, accountGateTestEpochID, accountGateTestEpochDigest,
			)
			if !errors.Is(err, test.want) {
				t.Fatalf("lock error = %v, want %v", err, test.want)
			}
			if err := tx.Rollback(); err != nil {
				t.Fatalf("rollback: %v", err)
			}
			if err := mock.ExpectationsWereMet(); err != nil {
				t.Fatalf("sql expectations: %v", err)
			}
		})
	}
}

func TestLockExpectedV2ServingRejectsMissingGateAndInvalidExpectedInput(t *testing.T) {
	t.Run("missing", func(t *testing.T) {
		db, mock := newServingSchemaMock(t)
		mock.ExpectBegin()
		tx, err := db.BeginTx(t.Context(), nil)
		if err != nil {
			t.Fatalf("begin: %v", err)
		}
		query := accountServingGateSelect + " FOR UPDATE"
		mock.ExpectQuery(regexp.QuoteMeta(query)).
			WithArgs(accountGateTestScope, accountGateTestScope, accountGateTestAccount, accountGateTestAccount).
			WillReturnRows(sqlmock.NewRows(accountServingGateColumnsForTest()))
		mock.ExpectRollback()
		_, err = lockExpectedV2Serving(t.Context(), tx, accountGateTestScope, accountGateTestAccount, 41, accountGateTestEpochID, accountGateTestEpochDigest)
		if !errors.Is(err, errAccountServingGateNotFound) {
			t.Fatalf("missing gate error = %v", err)
		}
		_ = tx.Rollback()
		if err := mock.ExpectationsWereMet(); err != nil {
			t.Fatalf("sql expectations: %v", err)
		}
	})

	for _, args := range []struct {
		ctx     bool
		tx      bool
		fence   int64
		epochID string
		digest  string
	}{
		{fence: 41, epochID: accountGateTestEpochID, digest: accountGateTestEpochDigest},
		{ctx: true, tx: true, fence: -1, epochID: accountGateTestEpochID, digest: accountGateTestEpochDigest},
		{ctx: true, tx: true, fence: 41, epochID: " epoch", digest: accountGateTestEpochDigest},
		{ctx: true, tx: true, fence: 41, epochID: accountGateTestEpochID, digest: "sha256:bad"},
	} {
		var ctx = t.Context()
		if !args.ctx {
			ctx = nil
		}
		var tx *sql.Tx
		if args.tx {
			tx = &sql.Tx{}
		}
		if _, err := lockExpectedV2Serving(ctx, tx, accountGateTestScope, accountGateTestAccount, args.fence, args.epochID, args.digest); !errors.Is(err, errInvalidAccountServingGateInput) {
			t.Fatalf("invalid expected input error = %v", err)
		}
	}
}

func validAccountServingGate(mode accountServingMode) accountServingGate {
	gate := accountServingGate{
		id:                 accountGateTestID,
		servingScopeDigest: accountGateTestScope,
		connectorAccountID: accountGateTestAccount,
		mode:               mode,
		fenceToken:         1,
	}
	if mode == accountServingModeV1Serving {
		gate.fenceToken = 0
	}
	if accountServingModeRequiresEpoch(mode) {
		gate.servingEpochID = stringPointer(accountGateTestEpochID)
		gate.servingEpochDigest = stringPointer(accountGateTestEpochDigest)
	}
	return gate
}

func validAccountServingEpochHeader() accountServingEpochHeader {
	return accountServingEpochHeader{
		id:                 accountGateTestEpochID,
		epochDigest:        accountGateTestEpochDigest,
		servingScopeDigest: accountGateTestScope,
		activeScopeDigest:  stringPointer(accountGateTestScope),
		state:              accountServingEpochStateActive,
	}
}

func expectAccountServingGateRow(mock sqlmock.Sqlmock, gate accountServingGate, forUpdate bool) {
	expectAccountServingGateRowWithArgs(mock, gate, forUpdate, accountGateTestScope, accountGateTestAccount)
}

func expectAccountServingGateRowWithArgs(
	mock sqlmock.Sqlmock,
	gate accountServingGate,
	forUpdate bool,
	scope, account string,
) {
	query := accountServingGateSelect
	if forUpdate {
		query += " FOR UPDATE"
	}
	var epochID, epochDigest any
	if gate.servingEpochID != nil {
		epochID = *gate.servingEpochID
	}
	if gate.servingEpochDigest != nil {
		epochDigest = *gate.servingEpochDigest
	}
	mock.ExpectQuery(regexp.QuoteMeta(query)).
		WithArgs(scope, scope, account, account).
		WillReturnRows(sqlmock.NewRows(accountServingGateColumnsForTest()).AddRow(
			gate.id, gate.servingScopeDigest, gate.connectorAccountID, string(gate.mode),
			gate.fenceToken, epochID, epochDigest,
		))
}

func expectAccountServingEpochHeaderRow(mock sqlmock.Sqlmock, header accountServingEpochHeader, forUpdate bool) {
	query := accountServingEpochHeaderSelect
	if forUpdate {
		query += " FOR UPDATE"
	}
	var activeScope any
	if header.activeScopeDigest != nil {
		activeScope = *header.activeScopeDigest
	}
	mock.ExpectQuery(regexp.QuoteMeta(query)).
		WithArgs(accountGateTestEpochID, accountGateTestEpochID).
		WillReturnRows(sqlmock.NewRows([]string{
			"Id", "EpochDigest", "ServingScopeDigest", "ActiveScopeDigest", "State",
		}).AddRow(
			header.id, header.epochDigest, header.servingScopeDigest, activeScope, header.state,
		))
}

func accountServingGateColumnsForTest() []string {
	return []string{
		"Id", "ServingScopeDigest", "ConnectorAccountId", "Mode", "FenceToken",
		"ServingEpochId", "ServingEpochDigest",
	}
}