package sqltx

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"
	"fmt"
	"io"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

type cleanupResetMode int

const (
	cleanupResetSucceeds cleanupResetMode = iota
	cleanupResetBlocks
	cleanupResetFails
)

type cleanupTestState struct {
	mu                     sync.Mutex
	resetMode              cleanupResetMode
	constraintResetMode    cleanupResetMode
	initialConstraintValue string
	connections            []*cleanupTestConn
	resetDeadlines         []time.Time
}

func (s *cleanupTestState) connect() *cleanupTestConn {
	s.mu.Lock()
	defer s.mu.Unlock()
	constraintValue := s.initialConstraintValue
	if constraintValue == "" {
		constraintValue = "OFF"
	}
	conn := &cleanupTestConn{
		id:              len(s.connections) + 1,
		state:           s,
		constraintValue: constraintValue,
	}
	s.connections = append(s.connections, conn)
	return conn
}

func (s *cleanupTestState) connectionCount() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.connections)
}

func (s *cleanupTestState) recordResetDeadline(deadline time.Time) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.resetDeadlines = append(s.resetDeadlines, deadline)
}

func (s *cleanupTestState) lastResetDeadline() (time.Time, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if len(s.resetDeadlines) == 0 {
		return time.Time{}, false
	}
	return s.resetDeadlines[len(s.resetDeadlines)-1], true
}

type cleanupTestConnector struct {
	state *cleanupTestState
}

func (c *cleanupTestConnector) Connect(context.Context) (driver.Conn, error) {
	return c.state.connect(), nil
}

func (c *cleanupTestConnector) Driver() driver.Driver {
	return cleanupTestDriver{state: c.state}
}

type cleanupTestDriver struct {
	state *cleanupTestState
}

func (d cleanupTestDriver) Open(string) (driver.Conn, error) {
	return d.state.connect(), nil
}

type cleanupTestConn struct {
	id              int
	state           *cleanupTestState
	constraintValue string
	closeCount      atomic.Int32
}

func (c *cleanupTestConn) Prepare(string) (driver.Stmt, error) {
	return nil, errors.New("prepare is not supported")
}

func (c *cleanupTestConn) Close() error {
	c.closeCount.Add(1)
	return nil
}

func (c *cleanupTestConn) Begin() (driver.Tx, error) {
	return &cleanupTestTx{}, nil
}

func (c *cleanupTestConn) BeginTx(context.Context, driver.TxOptions) (driver.Tx, error) {
	return &cleanupTestTx{}, nil
}

func (c *cleanupTestConn) QueryContext(_ context.Context, query string, _ []driver.NamedValue) (driver.Rows, error) {
	if strings.Contains(query, "@@SESSION.tidb_constraint_check_in_place_pessimistic") {
		return &cleanupTestRows{value: c.constraintValue}, nil
	}
	return nil, fmt.Errorf("unsupported query %q", query)
}

func (c *cleanupTestConn) ExecContext(ctx context.Context, query string, _ []driver.NamedValue) (driver.Result, error) {
	switch {
	case strings.Contains(query, "tidb_constraint_check_in_place_pessimistic = ON"):
		c.constraintValue = "ON"
		return driver.RowsAffected(0), nil
	case strings.Contains(query, "tidb_constraint_check_in_place_pessimistic = OFF"):
		deadline, ok := ctx.Deadline()
		if !ok {
			return nil, errors.New("constraint reset context has no deadline")
		}
		c.state.recordResetDeadline(deadline)
		switch c.state.constraintResetMode {
		case cleanupResetSucceeds:
			c.constraintValue = "OFF"
			return driver.RowsAffected(0), nil
		case cleanupResetBlocks:
			<-ctx.Done()
			return nil, ctx.Err()
		case cleanupResetFails:
			return nil, errors.New("constraint reset failed")
		default:
			return nil, fmt.Errorf("unknown constraint reset mode %d", c.state.constraintResetMode)
		}
	case strings.Contains(query, "tidb_txn_mode = 'pessimistic'"):
		return driver.RowsAffected(0), nil
	case strings.Contains(query, "tidb_txn_mode = DEFAULT"):
		deadline, ok := ctx.Deadline()
		if !ok {
			return nil, errors.New("reset context has no deadline")
		}
		c.state.recordResetDeadline(deadline)
		switch c.state.resetMode {
		case cleanupResetSucceeds:
			return driver.RowsAffected(0), nil
		case cleanupResetBlocks:
			<-ctx.Done()
			return nil, ctx.Err()
		case cleanupResetFails:
			return nil, errors.New("reset failed")
		default:
			return nil, fmt.Errorf("unknown reset mode %d", c.state.resetMode)
		}
	default:
		return driver.RowsAffected(0), nil
	}
}

type cleanupTestRows struct {
	value string
	done  bool
}

func (*cleanupTestRows) Columns() []string {
	return []string{"tidb_constraint_check_in_place_pessimistic"}
}
func (*cleanupTestRows) Close() error { return nil }

func (r *cleanupTestRows) Next(destination []driver.Value) error {
	if r.done {
		return io.EOF
	}
	destination[0] = r.value
	r.done = true
	return nil
}

type cleanupTestTx struct{}

func (*cleanupTestTx) Commit() error   { return nil }
func (*cleanupTestTx) Rollback() error { return nil }

func openCleanupTestDB(t *testing.T, mode cleanupResetMode) (*sql.DB, *cleanupTestState) {
	t.Helper()
	state := &cleanupTestState{
		resetMode:           mode,
		constraintResetMode: cleanupResetSucceeds,
	}
	db := sql.OpenDB(&cleanupTestConnector{state: state})
	db.SetMaxOpenConns(1)
	db.SetMaxIdleConns(1)
	t.Cleanup(func() { _ = db.Close() })
	return db, state
}

func openConstraintCleanupTestDB(t *testing.T, mode cleanupResetMode) (*sql.DB, *cleanupTestState) {
	t.Helper()
	state := &cleanupTestState{
		resetMode:           cleanupResetSucceeds,
		constraintResetMode: mode,
	}
	db := sql.OpenDB(&cleanupTestConnector{state: state})
	db.SetMaxOpenConns(1)
	db.SetMaxIdleConns(1)
	t.Cleanup(func() { _ = db.Close() })
	return db, state
}

func beginCleanupTestTx(t *testing.T, db *sql.DB) func() error {
	t.Helper()
	tx, cleanup, err := BeginPessimisticTx(context.Background(), db)
	if err != nil {
		t.Fatalf("begin pessimistic tx: %v", err)
	}
	if err := tx.Rollback(); err != nil {
		t.Fatalf("rollback: %v", err)
	}
	return cleanup
}

func beginConstraintCleanupTestTx(t *testing.T, db *sql.DB) func() error {
	t.Helper()
	tx, cleanup, err := BeginPessimisticConstraintCheckingTx(context.Background(), db)
	if err != nil {
		t.Fatalf("begin constraint-checking pessimistic tx: %v", err)
	}
	if err := tx.Rollback(); err != nil {
		t.Fatalf("rollback: %v", err)
	}
	return cleanup
}

func TestPessimisticTxCleanupBoundsBlockedResetAndDiscardsConnection(t *testing.T) {
	db, state := openCleanupTestDB(t, cleanupResetBlocks)
	cleanup := beginCleanupTestTx(t, db)

	started := time.Now()
	err := cleanup()
	elapsed := time.Since(started)
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("cleanup error = %v, want context deadline exceeded", err)
	}
	if elapsed < pessimisticTxCleanupTimeout-50*time.Millisecond {
		t.Fatalf("cleanup returned too early after %s", elapsed)
	}
	if elapsed > pessimisticTxCleanupTimeout+500*time.Millisecond {
		t.Fatalf("cleanup exceeded bounded reset budget: %s", elapsed)
	}

	deadline, ok := state.lastResetDeadline()
	if !ok {
		t.Fatal("reset did not receive a deadline")
	}
	budget := deadline.Sub(started)
	if budget < pessimisticTxCleanupTimeout-50*time.Millisecond || budget > pessimisticTxCleanupTimeout+50*time.Millisecond {
		t.Fatalf("reset deadline budget = %s, want about %s", budget, pessimisticTxCleanupTimeout)
	}

	conn, err := db.Conn(context.Background())
	if err != nil {
		t.Fatalf("acquire replacement conn: %v", err)
	}
	defer conn.Close()
	if got := state.connectionCount(); got != 2 {
		t.Fatalf("connection count = %d, want 2 after timed-out reset", got)
	}
	if got := state.connections[0].closeCount.Load(); got != 1 {
		t.Fatalf("bad connection close count = %d, want 1", got)
	}
}

func TestPessimisticTxCleanupDiscardsConnectionAfterResetFailure(t *testing.T) {
	db, state := openCleanupTestDB(t, cleanupResetFails)
	cleanup := beginCleanupTestTx(t, db)

	err := cleanup()
	if err == nil || !strings.Contains(err.Error(), "reset failed") {
		t.Fatalf("cleanup error = %v, want reset failure", err)
	}

	conn, err := db.Conn(context.Background())
	if err != nil {
		t.Fatalf("acquire replacement conn: %v", err)
	}
	defer conn.Close()
	if got := state.connectionCount(); got != 2 {
		t.Fatalf("connection count = %d, want 2 after failed reset", got)
	}
	if got := state.connections[0].closeCount.Load(); got != 1 {
		t.Fatalf("bad connection close count = %d, want 1", got)
	}
}

func TestPessimisticTxCleanupResetsAndReturnsHealthyConnectionToPool(t *testing.T) {
	db, state := openCleanupTestDB(t, cleanupResetSucceeds)
	cleanup := beginCleanupTestTx(t, db)

	if err := cleanup(); err != nil {
		t.Fatalf("cleanup: %v", err)
	}

	conn, err := db.Conn(context.Background())
	if err != nil {
		t.Fatalf("reacquire conn: %v", err)
	}
	defer conn.Close()
	if got := state.connectionCount(); got != 1 {
		t.Fatalf("connection count = %d, want healthy connection reuse", got)
	}
	if got := state.connections[0].closeCount.Load(); got != 0 {
		t.Fatalf("physical close count = %d, want connection retained in pool", got)
	}
	if _, ok := state.lastResetDeadline(); !ok {
		t.Fatal("successful reset did not receive a bounded context")
	}
}

func TestConstraintCheckingPessimisticTxCleanupBoundsBlockedRestoreAndDiscardsConnection(t *testing.T) {
	db, state := openConstraintCleanupTestDB(t, cleanupResetBlocks)
	cleanup := beginConstraintCleanupTestTx(t, db)

	started := time.Now()
	err := cleanup()
	elapsed := time.Since(started)
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("cleanup error = %v, want context deadline exceeded", err)
	}
	if elapsed < pessimisticTxCleanupTimeout-50*time.Millisecond {
		t.Fatalf("cleanup returned too early after %s", elapsed)
	}
	if elapsed > pessimisticTxCleanupTimeout+500*time.Millisecond {
		t.Fatalf("cleanup exceeded bounded reset budget: %s", elapsed)
	}

	conn, err := db.Conn(context.Background())
	if err != nil {
		t.Fatalf("acquire replacement conn: %v", err)
	}
	defer conn.Close()
	if got := state.connectionCount(); got != 2 {
		t.Fatalf("connection count = %d, want 2 after blocked constraint restore", got)
	}
	if got := state.connections[0].closeCount.Load(); got != 1 {
		t.Fatalf("bad connection close count = %d, want 1", got)
	}
}

func TestConstraintCheckingPessimisticTxCleanupDiscardsConnectionAfterRestoreFailure(t *testing.T) {
	db, state := openConstraintCleanupTestDB(t, cleanupResetFails)
	cleanup := beginConstraintCleanupTestTx(t, db)

	err := cleanup()
	if err == nil || !strings.Contains(err.Error(), "constraint reset failed") {
		t.Fatalf("cleanup error = %v, want constraint reset failure", err)
	}

	conn, err := db.Conn(context.Background())
	if err != nil {
		t.Fatalf("acquire replacement conn: %v", err)
	}
	defer conn.Close()
	if got := state.connectionCount(); got != 2 {
		t.Fatalf("connection count = %d, want 2 after failed constraint restore", got)
	}
	if got := state.connections[0].closeCount.Load(); got != 1 {
		t.Fatalf("bad connection close count = %d, want 1", got)
	}
}