package mysqlbackend

import (
	"context"
	"database/sql"
)

type connContextKey struct{}
type connLeaseProviderContextKey struct{}
type heldConnLeaseContextKey struct{}
type preferredConnLeaseKindContextKey struct{}

// ConnLeaseKind identifies the request class for a transient MySQL lease.
type ConnLeaseKind string

const (
	// ConnLeaseKindStandard is used by regular request work such as SELECT or DML.
	ConnLeaseKindStandard ConnLeaseKind = "standard"
	// ConnLeaseKindBulk is used by UNLOGGED BATCH worker fan-out.
	ConnLeaseKindBulk ConnLeaseKind = "bulk"
)

// ConnLeaseProvider acquires a transient MySQL lease for on-demand pool.Conn usage.
type ConnLeaseProvider interface {
	AcquireConnLease(kind ConnLeaseKind) (func(), error)
}

// WithConn stores a leased MySQL connection in the context.
func WithConn(ctx context.Context, conn *sql.Conn) context.Context {
	if ctx == nil || conn == nil {
		return ctx
	}
	return context.WithValue(ctx, connContextKey{}, conn)
}

// ConnFromContext returns a leased MySQL connection if present.
func ConnFromContext(ctx context.Context) (*sql.Conn, bool) {
	if ctx == nil {
		return nil, false
	}
	value := ctx.Value(connContextKey{})
	if value == nil {
		return nil, false
	}
	conn, ok := value.(*sql.Conn)
	return conn, ok
}

// WithConnLeaseProvider stores a lease provider in the context.
func WithConnLeaseProvider(ctx context.Context, provider ConnLeaseProvider) context.Context {
	if ctx == nil || provider == nil {
		return ctx
	}
	return context.WithValue(ctx, connLeaseProviderContextKey{}, provider)
}

// ConnLeaseProviderFromContext returns the lease provider stored in the context.
func ConnLeaseProviderFromContext(ctx context.Context) (ConnLeaseProvider, bool) {
	if ctx == nil {
		return nil, false
	}
	value := ctx.Value(connLeaseProviderContextKey{})
	if value == nil {
		return nil, false
	}
	provider, ok := value.(ConnLeaseProvider)
	return provider, ok
}

// WithHeldConnLease marks that the current request already holds a lease of the given kind.
func WithHeldConnLease(ctx context.Context, kind ConnLeaseKind) context.Context {
	if ctx == nil {
		return ctx
	}
	return context.WithValue(ctx, heldConnLeaseContextKey{}, kind)
}

// WithPreferredConnLeaseKind marks the preferred transient lease kind for
// on-demand pool.Conn acquisition on this request path.
func WithPreferredConnLeaseKind(ctx context.Context, kind ConnLeaseKind) context.Context {
	if ctx == nil {
		return ctx
	}
	return context.WithValue(ctx, preferredConnLeaseKindContextKey{}, kind)
}

// PreferredConnLeaseKindFromContext returns the preferred transient lease kind
// stored in the context.
func PreferredConnLeaseKindFromContext(ctx context.Context) (ConnLeaseKind, bool) {
	if ctx == nil {
		return "", false
	}
	value := ctx.Value(preferredConnLeaseKindContextKey{})
	if value == nil {
		return "", false
	}
	kind, ok := value.(ConnLeaseKind)
	return kind, ok
}

// HasHeldConnLease reports whether the context already holds a lease of the given kind.
func HasHeldConnLease(ctx context.Context, kind ConnLeaseKind) bool {
	if ctx == nil {
		return false
	}
	value := ctx.Value(heldConnLeaseContextKey{})
	if value == nil {
		return false
	}
	heldKind, ok := value.(ConnLeaseKind)
	return ok && heldKind == kind
}

// AcquireConnLeaseFromContext acquires a transient lease from the context
// provider, honoring the preferred request kind and skipping kinds already held
// on the request context.
func AcquireConnLeaseFromContext(ctx context.Context, defaultKind ConnLeaseKind) (func(), error) {
	provider, ok := ConnLeaseProviderFromContext(ctx)
	if !ok || provider == nil {
		return nil, nil
	}
	kind := defaultKind
	if preferredKind, ok := PreferredConnLeaseKindFromContext(ctx); ok {
		kind = preferredKind
	}
	if HasHeldConnLease(ctx, kind) {
		return nil, nil
	}
	return provider.AcquireConnLease(kind)
}