package database_test

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"
	"fmt"
	"runtime"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"git.yuntuo.hk/k2safe/go-omni/database"
)

// The driver can persist a write and lose the commit acknowledgement, or fail
// before persisting it. Both must be indistinguishable to the generic wrapper.
type txFixture struct {
	beginErr, commitErr, rollbackErr  error
	persistBeforeError                bool
	onBegin                           func(context.Context, driver.TxOptions)
	onCommit                          func()
	onRollback                        func()
	begins, commits, rollbacks, saved atomic.Int64
}

func (f *txFixture) Connect(context.Context) (driver.Conn, error) { return &txConn{owner: f}, nil }
func (f *txFixture) Driver() driver.Driver                        { return f }
func (f *txFixture) Open(string) (driver.Conn, error)             { return f.Connect(context.Background()) }

type txConn struct {
	owner   *txFixture
	pending int64
}

func (c *txConn) Prepare(string) (driver.Stmt, error) { return nil, errors.New("not used") }
func (c *txConn) Close() error                        { return nil }
func (c *txConn) Begin() (driver.Tx, error) {
	return c.BeginTx(context.Background(), driver.TxOptions{})
}
func (c *txConn) BeginTx(ctx context.Context, options driver.TxOptions) (driver.Tx, error) {
	c.owner.begins.Add(1)
	if c.owner.onBegin != nil {
		c.owner.onBegin(ctx, options)
	}
	if c.owner.beginErr != nil {
		return nil, c.owner.beginErr
	}
	c.pending = 0
	return &txDriver{conn: c}, nil
}
func (c *txConn) ExecContext(context.Context, string, []driver.NamedValue) (driver.Result, error) {
	c.pending++
	return driver.RowsAffected(1), nil
}

type txDriver struct{ conn *txConn }

func (tx *txDriver) Commit() error {
	f := tx.conn.owner
	f.commits.Add(1)
	if f.commitErr == nil || f.persistBeforeError {
		f.saved.Add(tx.conn.pending)
	}
	if f.onCommit != nil {
		f.onCommit()
	}
	return f.commitErr
}
func (tx *txDriver) Rollback() error {
	f := tx.conn.owner
	f.rollbacks.Add(1)
	if f.onRollback != nil {
		f.onRollback()
	}
	return f.rollbackErr
}

func txPool(t testing.TB, f *txFixture) *database.DB {
	t.Helper()
	pool := sql.OpenDB(f)
	pool.SetMaxOpenConns(8)
	t.Cleanup(func() { _ = pool.Close() })
	return &database.DB{DB: pool}
}

type txCause struct{ name string }

func (e *txCause) Error() string { return e.name }

func checkTxError(t *testing.T, err error, stage database.TxStage, outcome database.TxOutcome, cause, rollback error) *database.TxError {
	t.Helper()
	var txErr *database.TxError
	if !errors.As(fmt.Errorf("repository: %w", err), &txErr) {
		t.Fatalf("missing typed transaction error: %v", err)
	}
	if txErr.Stage != stage || txErr.Outcome != outcome || txErr.Err != cause || txErr.RollbackErr != rollback {
		t.Fatalf("unexpected transaction error: %+v", txErr)
	}
	if cause != nil && !errors.Is(err, cause) || rollback != nil && !errors.Is(err, rollback) {
		t.Fatalf("cause chain lost: %v", err)
	}
	return txErr
}

func TestInTxStagesAndCauses(t *testing.T) {
	primary, rollback := &txCause{"primary"}, &txCause{"rollback"}
	for _, tc := range []struct {
		name                             string
		beginErr, workErr, commitErr, rb error
		stage                            database.TxStage
		outcome                          database.TxOutcome
		wantWork, wantCommit, wantRB     int64
	}{
		{"begin", primary, nil, nil, nil, database.TxStageBegin, database.TxOutcomeNotStarted, 0, 0, 0},
		{"work", nil, primary, nil, nil, database.TxStageWork, database.TxOutcomeRolledBack, 1, 0, 1},
		{"workAndRollback", nil, primary, nil, rollback, database.TxStageWork, database.TxOutcomeUnknown, 1, 0, 1},
		{"commit", nil, nil, primary, nil, database.TxStageCommit, database.TxOutcomeUnknown, 1, 1, 0},
	} {
		t.Run(tc.name, func(t *testing.T) {
			f := &txFixture{beginErr: tc.beginErr, commitErr: tc.commitErr, rollbackErr: tc.rb}
			pool := txPool(t, f)
			var calls int64
			err := pool.InTx(context.Background(), nil, func(tx *sql.Tx) error {
				calls++
				_, err := tx.ExecContext(context.Background(), "write")
				if err != nil {
					t.Fatal(err)
				}
				return tc.workErr
			})
			rb := tc.rb
			if tc.stage == database.TxStageCommit {
				rb = sql.ErrTxDone // No driver rollback happened after Commit.
			}
			checkTxError(t, err, tc.stage, tc.outcome, primary, rb)
			var typedCause *txCause
			if !errors.As(err, &typedCause) || typedCause != primary {
				t.Fatal("primary typed cause not first in chain")
			}
			if calls != tc.wantWork || f.commits.Load() != tc.wantCommit || f.rollbacks.Load() != tc.wantRB || f.saved.Load() != 0 {
				t.Fatal("unexpected callback/commit/rollback calls or persisted write")
			}
		})
	}
	t.Run("rollbackTypedCause", func(t *testing.T) {
		pool := txPool(t, &txFixture{rollbackErr: rollback})
		cause := errors.New("work failed")
		err := pool.InTx(nil, nil, func(*sql.Tx) error { return cause })
		var found *txCause
		if !errors.As(err, &found) || found != rollback {
			t.Fatal("typed rollback cause lost")
		}
	})
}

func TestInTxCommitOutcomeUnknown(t *testing.T) {
	for _, cause := range []error{&txCause{"lost acknowledgement"}, context.DeadlineExceeded, driver.ErrBadConn} {
		for _, persisted := range []bool{false, true} {
			t.Run(fmt.Sprintf("%s/persisted=%t", cause, persisted), func(t *testing.T) {
				f := &txFixture{commitErr: cause, persistBeforeError: persisted}
				pool := txPool(t, f)
				calls := 0
				err := pool.InTx(nil, nil, func(tx *sql.Tx) error {
					calls++
					_, err := tx.ExecContext(context.Background(), "register object")
					return err
				})
				checkTxError(t, err, database.TxStageCommit, database.TxOutcomeUnknown, cause, sql.ErrTxDone)
				wantSaved := int64(0)
				if persisted {
					wantSaved = 1
				}
				if f.saved.Load() != wantSaved || calls != 1 || f.commits.Load() != 1 || f.rollbacks.Load() != 0 {
					t.Fatal("commit was retried, rolled back, or incorrectly modeled")
				}
			})
		}
	}
}

func TestInTxCancellationAndOwnership(t *testing.T) {
	t.Run("beforeBegin", func(t *testing.T) {
		f := &txFixture{}
		pool := txPool(t, f)
		ctx, cancel := context.WithCancel(context.Background())
		cancel()
		err := pool.InTx(ctx, nil, func(*sql.Tx) error { t.Fatal("work called"); return nil })
		checkTxError(t, err, database.TxStageBegin, database.TxOutcomeNotStarted, context.Canceled, nil)
		if f.begins.Load() != 0 {
			t.Fatal("driver began a canceled transaction")
		}
	})
	for _, returnedWorkError := range []bool{false, true} {
		t.Run(fmt.Sprintf("automaticRollback/workError=%t", returnedWorkError), func(t *testing.T) {
			rolledBack := make(chan struct{})
			rollbackCause := errors.New("rollback connection lost")
			f := &txFixture{rollbackErr: rollbackCause, onRollback: func() { close(rolledBack) }}
			pool := txPool(t, f)
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()
			err := pool.InTx(ctx, nil, func(*sql.Tx) error {
				cancel()
				select {
				case <-rolledBack:
				case <-time.After(3 * time.Second):
					t.Fatal("automatic rollback did not run")
				}
				if returnedWorkError {
					return context.Canceled
				}
				return nil
			})
			stage, primary := database.TxStageCommit, error(sql.ErrTxDone)
			if returnedWorkError {
				stage, primary = database.TxStageWork, context.Canceled
			}
			checkTxError(t, err, stage, database.TxOutcomeUnknown, primary, sql.ErrTxDone)
			// database/sql owns this asynchronous rollback and does not expose
			// its driver error. Do not pretend ErrTxDone means successful rollback.
			if errors.Is(err, rollbackCause) || f.rollbacks.Load() != 1 || f.commits.Load() != 0 {
				t.Fatal("unexpected automatic rollback accounting")
			}
		})
	}
	t.Run("cancelDuringCommit", func(t *testing.T) {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		f := &txFixture{commitErr: context.Canceled, persistBeforeError: true, onCommit: cancel}
		pool := txPool(t, f)
		err := pool.InTx(ctx, nil, func(tx *sql.Tx) error { _, err := tx.ExecContext(ctx, "write"); return err })
		checkTxError(t, err, database.TxStageCommit, database.TxOutcomeUnknown, context.Canceled, sql.ErrTxDone)
		if f.saved.Load() != 1 || f.rollbacks.Load() != 0 {
			t.Fatal("cancellation hid a committed write")
		}
	})
	for _, action := range []string{"commit", "rollback"} {
		t.Run("manual/"+action, func(t *testing.T) {
			f := &txFixture{}
			pool := txPool(t, f)
			cause := errors.New("callback failed")
			err := pool.InTx(nil, nil, func(tx *sql.Tx) error {
				var err error
				if action == "commit" {
					err = tx.Commit()
				} else {
					err = tx.Rollback()
				}
				if err != nil {
					t.Fatal(err)
				}
				return cause
			})
			checkTxError(t, err, database.TxStageWork, database.TxOutcomeUnknown, cause, sql.ErrTxDone)
		})
	}
}

func TestInTxSuccessPanicAndConcurrency(t *testing.T) {
	t.Run("success", func(t *testing.T) {
		type key struct{}
		ctx := context.WithValue(context.Background(), key{}, "value")
		f := &txFixture{onBegin: func(received context.Context, opts driver.TxOptions) {
			if received.Value(key{}) != "value" || !opts.ReadOnly || opts.Isolation != driver.IsolationLevel(sql.LevelSerializable) {
				t.Error("context or options changed")
			}
		}}
		pool := txPool(t, f)
		if err := pool.InTx(ctx, &sql.TxOptions{Isolation: sql.LevelSerializable, ReadOnly: true}, func(*sql.Tx) error { return nil }); err != nil {
			t.Fatal(err)
		}
		if f.commits.Load() != 1 || f.rollbacks.Load() != 0 {
			t.Fatal("unexpected successful transaction calls")
		}
	})
	for _, rbErr := range []error{nil, errors.New("rollback failed during unwind")} {
		t.Run(fmt.Sprintf("panic/%v", rbErr), func(t *testing.T) {
			f := &txFixture{rollbackErr: rbErr}
			pool := txPool(t, f)
			value := &struct{ text string }{"original panic"}
			func() {
				defer func() {
					if recover() != value {
						t.Error("panic identity lost")
					}
				}()
				_ = pool.InTx(nil, nil, func(*sql.Tx) error { panic(value) })
			}()
			if f.rollbacks.Load() != 1 || f.commits.Load() != 0 {
				t.Fatal("panic did not roll back exactly once")
			}
		})
	}
	t.Run("goexit", func(t *testing.T) {
		f := &txFixture{}
		pool := txPool(t, f)
		done := make(chan struct{})
		go func() {
			defer close(done)
			_ = pool.InTx(nil, nil, func(*sql.Tx) error { runtime.Goexit(); return nil })
		}()
		select {
		case <-done:
		case <-time.After(3 * time.Second):
			t.Fatal("Goexit did not finish")
		}
		if f.rollbacks.Load() != 1 || f.commits.Load() != 0 {
			t.Fatal("Goexit did not roll back exactly once")
		}
	})
	t.Run("concurrent", func(t *testing.T) {
		f := &txFixture{}
		pool := txPool(t, f)
		var group sync.WaitGroup
		for i := 0; i < 64; i++ {
			group.Add(1)
			go func(i int) {
				defer group.Done()
				cause := &txCause{fmt.Sprint(i)}
				err := pool.InTx(nil, nil, func(*sql.Tx) error { return cause })
				var txErr *database.TxError
				if !errors.As(err, &txErr) || txErr.Err != cause || txErr.Outcome != database.TxOutcomeRolledBack || !errors.Is(err, cause) {
					t.Error("cross-request error state")
				}
			}(i)
		}
		group.Wait()
		if f.rollbacks.Load() != 64 || f.commits.Load() != 0 {
			t.Fatal("incorrect concurrent calls")
		}
	})
}

func TestInTxPreconditionsAndNestedErrors(t *testing.T) {
	for _, pool := range []*database.DB{nil, {}} {
		err := pool.InTx(nil, nil, func(*sql.Tx) error { t.Fatal("unexpected callback"); return nil })
		var txErr *database.TxError
		if !errors.As(err, &txErr) || txErr.Stage != database.TxStageBegin || txErr.Outcome != database.TxOutcomeNotStarted || txErr.Err == nil {
			t.Fatal("uninitialized database was not classified")
		}
	}
	f := &txFixture{}
	pool := txPool(t, f)
	err := pool.InTx(nil, nil, nil)
	var txErr *database.TxError
	if !errors.As(err, &txErr) || txErr.Stage != database.TxStageBegin || f.begins.Load() != 0 {
		t.Fatal("nil callback began a transaction")
	}
	inner := &database.TxError{Stage: database.TxStageCommit, Outcome: database.TxOutcomeUnknown, Err: errors.New("other transaction")}
	err = pool.InTx(nil, nil, func(*sql.Tx) error { return inner })
	outer := checkTxError(t, err, database.TxStageWork, database.TxOutcomeRolledBack, inner, nil)
	if outer == inner || !errors.As(outer.Err, &txErr) || txErr != inner {
		t.Fatal("nested error was flattened")
	}
	// Empty or typed-nil error values must not panic or expose nil unwrap entries.
	for _, value := range []*database.TxError{nil, {}, {RollbackErr: inner}} {
		_ = value.Error()
		for _, child := range value.Unwrap() {
			if child == nil {
				t.Fatal("nil unwrap child")
			}
		}
	}
}

func TestInTxCancellationRace(t *testing.T) {
	const count = 128
	rollbackCause := errors.New("rollback acknowledgement unavailable")
	var rollbacks sync.WaitGroup
	rollbacks.Add(count)
	f := &txFixture{rollbackErr: rollbackCause, onRollback: rollbacks.Done}
	pool := txPool(t, f)
	for i := 0; i < count; i++ {
		ctx, cancel := context.WithCancel(context.Background())
		workFailure := i%2 == 0
		err := pool.InTx(ctx, nil, func(*sql.Tx) error {
			cancel()
			if i%3 == 0 {
				runtime.Gosched() // Let the automatic rollback race explicit cleanup.
			}
			if workFailure {
				return context.Canceled
			}
			return nil
		})
		cancel()
		var txErr *database.TxError
		if !errors.As(err, &txErr) || txErr.Outcome != database.TxOutcomeUnknown || f.commits.Load() != 0 {
			t.Fatal("cancellation was treated as confirmed rollback", err)
		}
		wantStage := database.TxStageCommit
		if workFailure {
			wantStage = database.TxStageWork
		}
		if txErr.Stage != wantStage || txErr.RollbackErr != rollbackCause && txErr.RollbackErr != sql.ErrTxDone {
			t.Fatal("unexpected cancellation diagnostics", err)
		}
		if txErr.RollbackErr == rollbackCause && !errors.Is(err, rollbackCause) {
			t.Fatal("observed rollback failure lost")
		}
	}
	done := make(chan struct{})
	go func() { rollbacks.Wait(); close(done) }()
	select {
	case <-done:
	case <-time.After(3 * time.Second):
		t.Fatal("canceled transactions did not finish rollback")
	}
}

func BenchmarkInTx(b *testing.B) {
	cause := errors.New("example failure")
	for _, name := range []string{"success", "beginFailure", "workFailure", "rollbackFailure", "commitUnknown"} {
		b.Run(name, func(b *testing.B) {
			f := &txFixture{}
			var workErr error
			switch name {
			case "beginFailure":
				f.beginErr = cause
			case "workFailure":
				workErr = cause
			case "rollbackFailure":
				workErr, f.rollbackErr = cause, cause
			case "commitUnknown":
				f.commitErr = cause
			}
			pool := txPool(b, f)
			work := func(*sql.Tx) error { return workErr }
			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				err := pool.InTx(context.Background(), nil, work)
				if (err == nil) != (name == "success") {
					b.Fatal(err)
				}
			}
		})
	}
}