ResolveSpec/pkg/dbmanager/connection.go

package dbmanager

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/uptrace/bun"
	"github.com/uptrace/bun/schema"
	"go.mongodb.org/mongo-driver/mongo"
	"gorm.io/gorm"

	"github.com/bitechdev/ResolveSpec/pkg/common"
	"github.com/bitechdev/ResolveSpec/pkg/common/adapters/database"
	"github.com/bitechdev/ResolveSpec/pkg/dbmanager/providers"
)

// Connection represents a single named database connection
type Connection interface {
	// Metadata
	Name() string
	Type() DatabaseType

	// ORM Access (SQL databases only)
	Bun() (*bun.DB, error)
	GORM() (*gorm.DB, error)
	Native() (*sql.DB, error)
	DB() (*sql.DB, error)

	// Common Database interface (for SQL databases)
	Database() (common.Database, error)

	// MongoDB Access (MongoDB only)
	MongoDB() (*mongo.Client, error)

	// Lifecycle
	Connect(ctx context.Context) error
	Close() error
	HealthCheck(ctx context.Context) error
	Reconnect(ctx context.Context) error

	// Stats
	Stats() *ConnectionStats
}

// ConnectionStats contains statistics about a database connection
type ConnectionStats struct {
	Name              string
	Type              DatabaseType
	Connected         bool
	LastHealthCheck   time.Time
	HealthCheckStatus string

	// SQL connection pool stats
	OpenConnections    int
	MaxOpenConnections int
	TotalOpened        int64 // physical connections ever opened
	InUse              int
	Idle               int
	WaitCount          int64
	WaitDuration       time.Duration
	MaxIdleClosed      int64
	MaxLifetimeClosed  int64
}

// sqlConnection implements Connection for SQL databases (PostgreSQL, SQLite, MSSQL)
type sqlConnection struct {
	name     string
	dbType   DatabaseType
	config   ConnectionConfig
	provider Provider

	// Lazy-initialized ORM instances (all wrap the same sql.DB)
	nativeDB *sql.DB
	bunDB    *bun.DB
	gormDB   *gorm.DB

	// Adapters for common.Database interface
	bunAdapter    *database.BunAdapter
	gormAdapter   *database.GormAdapter
	nativeAdapter common.Database

	// State
	connected bool
	mu        sync.RWMutex
	// lifecycleMu serialises Connect/Close/Reconnect against health-check pings.
	// Lock order: lifecycleMu before mu.
	lifecycleMu sync.RWMutex

	// Health check
	lastHealthCheck   time.Time
	healthCheckStatus string
}

// newSQLConnection creates a new SQL connection
func newSQLConnection(name string, dbType DatabaseType, config ConnectionConfig, provider Provider) *sqlConnection {
	return &sqlConnection{
		name:     name,
		dbType:   dbType,
		config:   config,
		provider: provider,
	}
}

// Name returns the connection name
func (c *sqlConnection) Name() string {
	return c.name
}

// Type returns the database type
func (c *sqlConnection) Type() DatabaseType {
	return c.dbType
}

// Connect establishes the database connection
func (c *sqlConnection) Connect(ctx context.Context) error {
	c.lifecycleMu.Lock()
	defer c.lifecycleMu.Unlock()
	c.mu.Lock()
	defer c.mu.Unlock()

	return c.connectLocked(ctx)
}

// connectLocked requires lifecycleMu and mu held for writing.
func (c *sqlConnection) connectLocked(ctx context.Context) error {
	if c.connected {
		return ErrAlreadyConnected
	}

	if err := c.provider.Connect(ctx, &c.config); err != nil {
		return NewConnectionError(c.name, "connect", err)
	}

	c.connected = true
	return nil
}

// Close closes the database connection and all ORM instances
func (c *sqlConnection) Close() error {
	c.lifecycleMu.Lock()
	defer c.lifecycleMu.Unlock()
	c.mu.Lock()
	defer c.mu.Unlock()

	return c.closeLocked()
}

// closeLocked requires lifecycleMu and mu held for writing. The connection is
// always marked disconnected and its cached handles dropped, even when closing
// fails, so accessors never hand out handles over a half-closed pool.
func (c *sqlConnection) closeLocked() error {
	if !c.connected {
		return nil
	}

	var errs []error

	// Close Bun if initialized. bun.DB.Close closes the underlying *sql.DB, so
	// skip it when the pool belongs to the caller.
	if o, ok := c.provider.(interface{ OwnsDB() bool }); c.bunDB != nil && (!ok || o.OwnsDB()) {
		if err := c.bunDB.Close(); err != nil {
			errs = append(errs, NewConnectionError(c.name, "close bun", err))
		}
	}

	// GORM doesn't have a separate close - it uses the underlying sql.DB

	// Close the provider (which closes the underlying sql.DB)
	if err := c.provider.Close(); err != nil {
		errs = append(errs, NewConnectionError(c.name, "close", err))
	}

	c.connected = false
	c.nativeDB = nil
	c.bunDB = nil
	c.gormDB = nil
	c.bunAdapter = nil
	c.gormAdapter = nil
	c.nativeAdapter = nil

	return errors.Join(errs...)
}

// HealthCheck verifies the connection is alive. The network ping runs without
// holding mu, so handle accessors are never blocked behind a slow ping.
func (c *sqlConnection) HealthCheck(ctx context.Context) error {
	if c == nil {
		return fmt.Errorf("connection is nil")
	}

	// lifecycleMu (read) keeps Close/Reconnect from tearing the provider down
	// mid-ping without blocking the accessors that only need mu.
	c.lifecycleMu.RLock()
	defer c.lifecycleMu.RUnlock()

	c.mu.RLock()
	connected := c.connected
	provider := c.provider
	c.mu.RUnlock()

	if !connected {
		c.setHealth("disconnected")
		return ErrConnectionClosed
	}

	if err := provider.HealthCheck(ctx); err != nil {
		c.setHealth("unhealthy: " + err.Error())
		return NewConnectionError(c.name, "health check", err)
	}

	c.setHealth("healthy")
	return nil
}

func (c *sqlConnection) setHealth(status string) {
	c.mu.Lock()
	c.lastHealthCheck = time.Now()
	c.healthCheckStatus = status
	c.mu.Unlock()
}

// Reconnect refreshes the connection as a single critical section.
//
// Providers that support it (PostgreSQL) retire their pooled connections and
// dial fresh ones without closing the *sql.DB, so handles handed out earlier
// keep working. Other providers fall back to Close+Connect, which invalidates
// earlier handles; that is meant for explicit operator use only, since
// *sql.DB already replaces broken connections by itself.
func (c *sqlConnection) Reconnect(ctx context.Context) (err error) {
	c.lifecycleMu.Lock()
	defer c.lifecycleMu.Unlock()
	c.mu.Lock()
	defer c.mu.Unlock()

	defer func() { RecordReconnectAttempt(c.name, c.dbType, err == nil) }()

	if c.connected {
		if r, ok := c.provider.(providers.Refresher); ok {
			if err := r.Refresh(ctx); err != nil {
				return NewConnectionError(c.name, "reconnect", err)
			}
			return nil
		}
	}

	if err := c.closeLocked(); err != nil {
		return err
	}
	return c.connectLocked(ctx)
}

// Native returns the native *sql.DB connection
func (c *sqlConnection) Native() (*sql.DB, error) {
	if c == nil {
		return nil, fmt.Errorf("connection is nil")
	}
	c.mu.RLock()
	if c.nativeDB != nil {
		defer c.mu.RUnlock()
		return c.nativeDB, nil
	}
	c.mu.RUnlock()

	c.mu.Lock()
	defer c.mu.Unlock()

	// Double-check after acquiring write lock
	if c.nativeDB != nil {
		return c.nativeDB, nil
	}

	if !c.connected {
		return nil, ErrConnectionClosed
	}

	// Get native connection from provider
	db, err := c.provider.GetNative()
	if err != nil {
		return nil, NewConnectionError(c.name, "get native", err)
	}

	c.nativeDB = db
	return c.nativeDB, nil
}

// DB returns the underlying *sql.DB connection
func (c *sqlConnection) DB() (*sql.DB, error) {
	return c.Native()
}

// Bun returns a Bun ORM instance wrapping the native connection
func (c *sqlConnection) Bun() (*bun.DB, error) {
	if c == nil {
		return nil, fmt.Errorf("connection is nil")
	}
	c.mu.RLock()
	if c.bunDB != nil {
		defer c.mu.RUnlock()
		return c.bunDB, nil
	}
	c.mu.RUnlock()

	c.mu.Lock()
	defer c.mu.Unlock()

	// Double-check after acquiring write lock
	if c.bunDB != nil {
		return c.bunDB, nil
	}

	if !c.connected {
		return nil, ErrConnectionClosed
	}

	// Get native connection first
	native, err := c.provider.GetNative()
	if err != nil {
		return nil, NewConnectionError(c.name, "get bun", err)
	}

	// Create Bun DB wrapping the same sql.DB
	dialect := c.getBunDialect()
	c.bunDB = bun.NewDB(native, dialect)

	return c.bunDB, nil
}

// GORM returns a GORM instance wrapping the native connection
func (c *sqlConnection) GORM() (*gorm.DB, error) {
	if c == nil {
		return nil, fmt.Errorf("connection is nil")
	}
	c.mu.RLock()
	if c.gormDB != nil {
		defer c.mu.RUnlock()
		return c.gormDB, nil
	}
	c.mu.RUnlock()

	c.mu.Lock()
	defer c.mu.Unlock()

	// Double-check after acquiring write lock
	if c.gormDB != nil {
		return c.gormDB, nil
	}

	if !c.connected {
		return nil, ErrConnectionClosed
	}

	// Get native connection first
	native, err := c.provider.GetNative()
	if err != nil {
		return nil, NewConnectionError(c.name, "get gorm", err)
	}

	// Create GORM DB wrapping the same sql.DB
	dialector := c.getGORMDialector(native)
	db, err := gorm.Open(dialector, &gorm.Config{})
	if err != nil {
		return nil, NewConnectionError(c.name, "initialize gorm", err)
	}

	c.gormDB = db
	return c.gormDB, nil
}

// Database returns the common.Database interface using the configured default ORM
func (c *sqlConnection) Database() (common.Database, error) {
	if c == nil {
		return nil, fmt.Errorf("connection is nil")
	}
	c.mu.RLock()
	defaultORM := c.config.DefaultORM
	c.mu.RUnlock()

	switch ORMType(defaultORM) {
	case ORMTypeBun:
		return c.getBunAdapter()
	case ORMTypeGORM:
		return c.getGORMAdapter()
	case ORMTypeNative:
		return c.getNativeAdapter()
	default:
		// Default to Bun
		return c.getBunAdapter()
	}
}

// MongoDB returns an error for SQL connections
func (c *sqlConnection) MongoDB() (*mongo.Client, error) {
	return nil, ErrNotMongoDB
}

// Stats returns connection statistics
func (c *sqlConnection) Stats() *ConnectionStats {
	if c == nil {
		return nil
	}
	c.mu.RLock()
	defer c.mu.RUnlock()

	stats := &ConnectionStats{
		Name:              c.name,
		Type:              c.dbType,
		Connected:         c.connected,
		LastHealthCheck:   c.lastHealthCheck,
		HealthCheckStatus: c.healthCheckStatus,
	}

	// Get SQL stats if connected
	if c.connected && c.provider != nil {
		if providerStats := c.provider.Stats(); providerStats != nil {
			stats.OpenConnections = providerStats.OpenConnections
			stats.MaxOpenConnections = providerStats.MaxOpenConnections
			stats.TotalOpened = providerStats.TotalOpened
			stats.InUse = providerStats.InUse
			stats.Idle = providerStats.Idle
			stats.WaitCount = providerStats.WaitCount
			stats.WaitDuration = providerStats.WaitDuration
			stats.MaxIdleClosed = providerStats.MaxIdleClosed
			stats.MaxLifetimeClosed = providerStats.MaxLifetimeClosed
		}
	}

	return stats
}

// The adapter factories only re-fetch the current handle. They must not close
// the shared pool: *sql.DB discards bad connections on its own, and closing it
// here would break every other holder of the pool.
func (c *sqlConnection) reopenNativeForAdapter() (*sql.DB, error) {
	return c.Native()
}

func (c *sqlConnection) reopenBunForAdapter() (*bun.DB, error) {
	return c.Bun()
}

func (c *sqlConnection) reopenGORMForAdapter() (*gorm.DB, error) {
	return c.GORM()
}

// getBunAdapter returns or creates the Bun adapter
func (c *sqlConnection) getBunAdapter() (common.Database, error) {
	if c == nil {
		return nil, fmt.Errorf("connection is nil")
	}
	c.mu.RLock()
	if c.bunAdapter != nil {
		defer c.mu.RUnlock()
		return c.bunAdapter, nil
	}
	c.mu.RUnlock()

	c.mu.Lock()
	defer c.mu.Unlock()

	if c.bunAdapter != nil {
		return c.bunAdapter, nil
	}

	// Double-check bunDB exists (while already holding write lock)
	if c.bunDB == nil {
		// Get native connection first
		native, err := c.provider.GetNative()
		if err != nil {
			return nil, NewConnectionError(c.name, "get bun", err)
		}

		// Create Bun DB wrapping the same sql.DB
		dialect := c.getBunDialect()
		c.bunDB = bun.NewDB(native, dialect)
	}

	c.bunAdapter = database.NewBunAdapter(c.bunDB).
		WithDBFactory(c.reopenBunForAdapter).
		SetMetricsEnabled(c.config.EnableMetrics)
	return c.bunAdapter, nil
}

// getGORMAdapter returns or creates the GORM adapter
func (c *sqlConnection) getGORMAdapter() (common.Database, error) {
	if c == nil {
		return nil, fmt.Errorf("connection is nil")
	}
	c.mu.RLock()
	if c.gormAdapter != nil {
		defer c.mu.RUnlock()
		return c.gormAdapter, nil
	}
	c.mu.RUnlock()

	c.mu.Lock()
	defer c.mu.Unlock()

	if c.gormAdapter != nil {
		return c.gormAdapter, nil
	}

	// Double-check gormDB exists (while already holding write lock)
	if c.gormDB == nil {
		// Get native connection first
		native, err := c.provider.GetNative()
		if err != nil {
			return nil, NewConnectionError(c.name, "get gorm", err)
		}

		// Create GORM DB wrapping the same sql.DB
		dialector := c.getGORMDialector(native)
		db, err := gorm.Open(dialector, &gorm.Config{})
		if err != nil {
			return nil, NewConnectionError(c.name, "initialize gorm", err)
		}

		c.gormDB = db
	}

	c.gormAdapter = database.NewGormAdapter(c.gormDB).
		WithDBFactory(c.reopenGORMForAdapter).
		SetMetricsEnabled(c.config.EnableMetrics)
	return c.gormAdapter, nil
}

// getNativeAdapter returns or creates the native adapter
func (c *sqlConnection) getNativeAdapter() (common.Database, error) {
	if c == nil {
		return nil, fmt.Errorf("connection is nil")
	}
	c.mu.RLock()
	if c.nativeAdapter != nil {
		defer c.mu.RUnlock()
		return c.nativeAdapter, nil
	}
	c.mu.RUnlock()

	c.mu.Lock()
	defer c.mu.Unlock()

	if c.nativeAdapter != nil {
		return c.nativeAdapter, nil
	}

	// Double-check nativeDB exists (while already holding write lock)
	if c.nativeDB == nil {
		if !c.connected {
			return nil, ErrConnectionClosed
		}

		// Get native connection from provider
		db, err := c.provider.GetNative()
		if err != nil {
			return nil, NewConnectionError(c.name, "get native", err)
		}

		c.nativeDB = db
	}

	// Create a native adapter based on database type
	switch c.dbType {
	case DatabaseTypePostgreSQL, DatabaseTypeSQLite, DatabaseTypeMSSQL:
		// The adapter takes the driver name so it can adjust its dialect.
		c.nativeAdapter = database.NewPgSQLAdapter(c.nativeDB, string(c.dbType)).
			WithDBFactory(c.reopenNativeForAdapter).
			SetMetricsEnabled(c.config.EnableMetrics)
	default:
		return nil, ErrUnsupportedDatabase
	}

	return c.nativeAdapter, nil
}

// getBunDialect returns the appropriate Bun dialect for the database type
func (c *sqlConnection) getBunDialect() schema.Dialect {

	switch c.dbType {
	case DatabaseTypePostgreSQL:
		return database.GetPostgresDialect()
	case DatabaseTypeSQLite:
		return database.GetSQLiteDialect()
	case DatabaseTypeMSSQL:
		return database.GetMSSQLDialect()
	default:
		// Default to PostgreSQL
		return database.GetPostgresDialect()
	}
}

// getGORMDialector returns the appropriate GORM dialector for the database type
func (c *sqlConnection) getGORMDialector(db *sql.DB) gorm.Dialector {
	switch c.dbType {
	case DatabaseTypePostgreSQL:
		return database.GetPostgresDialector(db)
	case DatabaseTypeSQLite:
		return database.GetSQLiteDialector(db)
	case DatabaseTypeMSSQL:
		return database.GetMSSQLDialector(db)
	default:
		// Default to PostgreSQL
		return database.GetPostgresDialector(db)
	}
}

// mongoConnection implements Connection for MongoDB
type mongoConnection struct {
	name     string
	config   ConnectionConfig
	provider Provider

	// MongoDB client
	client *mongo.Client

	// State
	connected   bool
	mu          sync.RWMutex
	lifecycleMu sync.RWMutex // see sqlConnection.lifecycleMu

	// Health check
	lastHealthCheck   time.Time
	healthCheckStatus string
}

// newMongoConnection creates a new MongoDB connection
func newMongoConnection(name string, config ConnectionConfig, provider Provider) *mongoConnection {
	return &mongoConnection{
		name:     name,
		config:   config,
		provider: provider,
	}
}

// Name returns the connection name
func (c *mongoConnection) Name() string {
	return c.name
}

// Type returns the database type (MongoDB)
func (c *mongoConnection) Type() DatabaseType {
	return DatabaseTypeMongoDB
}

// Connect establishes the MongoDB connection
func (c *mongoConnection) Connect(ctx context.Context) error {
	c.lifecycleMu.Lock()
	defer c.lifecycleMu.Unlock()
	c.mu.Lock()
	defer c.mu.Unlock()

	return c.connectLocked(ctx)
}

// connectLocked requires lifecycleMu and mu held for writing.
func (c *mongoConnection) connectLocked(ctx context.Context) error {
	if c.connected {
		return ErrAlreadyConnected
	}

	if err := c.provider.Connect(ctx, &c.config); err != nil {
		return NewConnectionError(c.name, "connect", err)
	}

	// Get the mongo client
	client, err := c.provider.GetMongo()
	if err != nil {
		_ = c.provider.Close()
		return NewConnectionError(c.name, "get mongo client", err)
	}

	c.client = client
	c.connected = true
	return nil
}

// Close closes the MongoDB connection
func (c *mongoConnection) Close() error {
	c.lifecycleMu.Lock()
	defer c.lifecycleMu.Unlock()
	c.mu.Lock()
	defer c.mu.Unlock()

	return c.closeLocked()
}

// closeLocked requires lifecycleMu and mu held for writing. The connection is
// marked disconnected even when the provider fails to close cleanly.
func (c *mongoConnection) closeLocked() error {
	if !c.connected {
		return nil
	}

	err := c.provider.Close()

	c.connected = false
	c.client = nil
	if err != nil {
		return NewConnectionError(c.name, "close", err)
	}
	return nil
}

// HealthCheck verifies the MongoDB connection is alive. The ping runs without
// holding mu so handle accessors are never blocked behind it.
func (c *mongoConnection) HealthCheck(ctx context.Context) error {
	c.lifecycleMu.RLock()
	defer c.lifecycleMu.RUnlock()

	c.mu.RLock()
	connected := c.connected
	c.mu.RUnlock()

	if !connected {
		c.setHealth("disconnected")
		return ErrConnectionClosed
	}

	if err := c.provider.HealthCheck(ctx); err != nil {
		c.setHealth("unhealthy: " + err.Error())
		return NewConnectionError(c.name, "health check", err)
	}

	c.setHealth("healthy")
	return nil
}

func (c *mongoConnection) setHealth(status string) {
	c.mu.Lock()
	c.lastHealthCheck = time.Now()
	c.healthCheckStatus = status
	c.mu.Unlock()
}

// Reconnect closes and re-establishes the MongoDB connection atomically.
func (c *mongoConnection) Reconnect(ctx context.Context) (err error) {
	c.lifecycleMu.Lock()
	defer c.lifecycleMu.Unlock()
	c.mu.Lock()
	defer c.mu.Unlock()

	defer func() { RecordReconnectAttempt(c.name, DatabaseTypeMongoDB, err == nil) }()

	if err := c.closeLocked(); err != nil {
		return err
	}
	return c.connectLocked(ctx)
}

// MongoDB returns the MongoDB client
func (c *mongoConnection) MongoDB() (*mongo.Client, error) {
	c.mu.RLock()
	defer c.mu.RUnlock()

	if !c.connected || c.client == nil {
		return nil, ErrConnectionClosed
	}

	return c.client, nil
}

// Bun returns an error for MongoDB connections
func (c *mongoConnection) Bun() (*bun.DB, error) {
	return nil, ErrNotSQLDatabase
}

// GORM returns an error for MongoDB connections
func (c *mongoConnection) GORM() (*gorm.DB, error) {
	return nil, ErrNotSQLDatabase
}

// Native returns an error for MongoDB connections
func (c *mongoConnection) Native() (*sql.DB, error) {
	return nil, ErrNotSQLDatabase
}

// DB returns an error for MongoDB connections
func (c *mongoConnection) DB() (*sql.DB, error) {
	return nil, ErrNotSQLDatabase
}

// Database returns an error for MongoDB connections
func (c *mongoConnection) Database() (common.Database, error) {
	return nil, ErrNotSQLDatabase
}

// Stats returns connection statistics for MongoDB
func (c *mongoConnection) Stats() *ConnectionStats {
	c.mu.RLock()
	defer c.mu.RUnlock()

	return &ConnectionStats{
		Name:              c.name,
		Type:              DatabaseTypeMongoDB,
		Connected:         c.connected,
		LastHealthCheck:   c.lastHealthCheck,
		HealthCheckStatus: c.healthCheckStatus,
	}
}