package gormx

import (
	"context"
	"database/sql"
	"sync/atomic"

	"gorm.io/gorm"
)

// GORM can call DB().Close on failed initialization. Do not expose the owner's
// sql.DB until initialization has succeeded; SQL methods remain fully usable.
// Dialects using context.Background during version probing are bounded by the
// startup context, which stops affecting queries as soon as Open succeeds.
type borrowedPool struct {
	db      *sql.DB
	startup atomic.Pointer[startupContext]
}

type startupContext struct{ context.Context }

func (p *borrowedPool) operationContext(ctx context.Context) context.Context {
	if startup := p.startup.Load(); startup != nil {
		return startup.Context
	}
	return ctx
}
func (p *borrowedPool) GetDBConn() (*sql.DB, error) {
	if p.startup.Load() != nil {
		return nil, gorm.ErrInvalidDB
	}
	return p.db, nil
}
func (p *borrowedPool) PrepareContext(ctx context.Context, query string) (*sql.Stmt, error) {
	return p.db.PrepareContext(p.operationContext(ctx), query)
}
func (p *borrowedPool) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
	return p.db.ExecContext(p.operationContext(ctx), query, args...)
}
func (p *borrowedPool) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
	return p.db.QueryContext(p.operationContext(ctx), query, args...)
}
func (p *borrowedPool) QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row {
	return p.db.QueryRowContext(p.operationContext(ctx), query, args...)
}
func (p *borrowedPool) BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error) {
	return p.db.BeginTx(p.operationContext(ctx), opts)
}

var _ gorm.ConnPool = (*borrowedPool)(nil)
var _ gorm.TxBeginner = (*borrowedPool)(nil)
var _ gorm.GetDBConnector = (*borrowedPool)(nil)