package packs

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"
	"fmt"
	"strings"
	"time"
)

type managedLifecycleSemanticTransaction struct {
	conn     *sql.Conn
	tx       *sql.Tx
	previous managedLifecycleSessionState
}

type managedLifecycleSessionState struct {
	txnMode                string
	constraintCheckInPlace int
	isolation              string
}

func beginManagedLifecycleSemanticTransaction(ctx context.Context, db *sql.DB) (*managedLifecycleSemanticTransaction, error) {
	conn, err := db.Conn(ctx)
	if err != nil {
		return nil, fmt.Errorf("acquire managed lifecycle transaction connection: %w", err)
	}
	var previous managedLifecycleSessionState
	if err := conn.QueryRowContext(ctx, "SELECT @@SESSION.tidb_txn_mode, @@SESSION.tidb_constraint_check_in_place_pessimistic, @@SESSION.transaction_isolation").Scan(
		&previous.txnMode, &previous.constraintCheckInPlace, &previous.isolation,
	); err != nil {
		_ = conn.Close()
		return nil, fmt.Errorf("snapshot managed lifecycle transaction session: %w", err)
	}
	if _, err := managedLifecycleIsolationStatement("SET SESSION TRANSACTION ISOLATION LEVEL ", previous.isolation); err != nil {
		_ = conn.Close()
		return nil, err
	}
	transaction := &managedLifecycleSemanticTransaction{conn: conn, previous: previous}
	closeOnError := func(cause error) (*managedLifecycleSemanticTransaction, error) {
		return nil, errors.Join(cause, transaction.close())
	}
	if _, err := conn.ExecContext(ctx, "SET SESSION tidb_txn_mode = 'pessimistic'"); err != nil {
		return closeOnError(fmt.Errorf("select managed lifecycle pessimistic mode: %w", err))
	}
	if _, err := conn.ExecContext(ctx, "SET SESSION tidb_constraint_check_in_place_pessimistic = ON"); err != nil {
		return closeOnError(fmt.Errorf("select managed lifecycle in-place uniqueness checks: %w", err))
	}
	if _, err := conn.ExecContext(ctx, "SET SESSION TRANSACTION ISOLATION LEVEL READ COMMITTED"); err != nil {
		return closeOnError(fmt.Errorf("select managed lifecycle read-committed mode: %w", err))
	}
	tx, err := conn.BeginTx(ctx, nil)
	if err != nil {
		return closeOnError(fmt.Errorf("begin managed lifecycle transaction: %w", err))
	}
	transaction.tx = tx
	return transaction, nil
}

func (t *managedLifecycleSemanticTransaction) commit() (error, error) {
	commitErr := t.tx.Commit()
	return commitErr, t.close()
}

func (t *managedLifecycleSemanticTransaction) rollback() bool {
	rollbackErr := t.tx.Rollback()
	closeErr := t.close()
	return rollbackErr == nil && closeErr == nil
}

func (t *managedLifecycleSemanticTransaction) close() error {
	isolationStatement, isolationErr := managedLifecycleIsolationStatement("SET SESSION TRANSACTION ISOLATION LEVEL ", t.previous.isolation)
	if isolationErr != nil {
		_ = t.conn.Close()
		return isolationErr
	}
	resetCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, constraintResetErr := t.conn.ExecContext(resetCtx, "SET SESSION tidb_constraint_check_in_place_pessimistic = ?", t.previous.constraintCheckInPlace)
	_, modeResetErr := t.conn.ExecContext(resetCtx, "SET SESSION tidb_txn_mode = ?", t.previous.txnMode)
	_, isolationResetErr := t.conn.ExecContext(resetCtx, isolationStatement)
	resetErr := errors.Join(constraintResetErr, modeResetErr, isolationResetErr)
	if resetErr != nil {
		_ = t.conn.Raw(func(any) error { return driver.ErrBadConn })
		_ = t.conn.Close()
		return resetErr
	}
	return t.conn.Close()
}

func managedLifecycleIsolationStatement(prefix, isolation string) (string, error) {
	switch strings.ToUpper(strings.ReplaceAll(strings.TrimSpace(isolation), "_", "-")) {
	case "READ-COMMITTED":
		return prefix + "READ COMMITTED", nil
	case "REPEATABLE-READ":
		return prefix + "REPEATABLE READ", nil
	case "READ-UNCOMMITTED":
		return prefix + "READ UNCOMMITTED", nil
	case "SERIALIZABLE":
		return prefix + "SERIALIZABLE", nil
	default:
		return "", fmt.Errorf("managed lifecycle cannot restore unsupported transaction isolation %q", isolation)
	}
}