package data

import (
	"context"
	"sync"
	"sync/atomic"

	"github.com/redis/go-redis/v9"
	"go.mongodb.org/mongo-driver/mongo"
	"gorm.io/gorm"
)

// reloadableDB makes the pointer swap atomic. Replaced pools are retained
// until application shutdown so in-flight GORM operations remain valid.
type reloadableDB struct {
	current atomic.Pointer[gorm.DB]
	mu      sync.Mutex
	retired []*gorm.DB
}

type reloadableMongo struct {
	mu      sync.RWMutex
	current *mongo.Client
	retired []*mongo.Client
}

func newReloadableMongo(client *mongo.Client) *reloadableMongo {
	return &reloadableMongo{current: client}
}
func (r *reloadableMongo) replace(client *mongo.Client) {
	r.mu.Lock()
	old := r.current
	r.current = client
	if old != nil && old != client {
		r.retired = append(r.retired, old)
	}
	r.mu.Unlock()
}
func (r *reloadableMongo) close() {
	r.mu.Lock()
	all := append([]*mongo.Client{r.current}, r.retired...)
	r.current = nil
	r.retired = nil
	r.mu.Unlock()
	for _, client := range all {
		if client != nil {
			_ = client.Disconnect(context.Background())
		}
	}
}

func newReloadableDB(db *gorm.DB) *reloadableDB {
	r := &reloadableDB{}
	registerDataScopeCallbacks(db)
	r.current.Store(db)
	return r
}

func (r *reloadableDB) WithContext(ctx context.Context) *gorm.DB {
	return r.current.Load().WithContext(ctx)
}

func (r *reloadableDB) DB() *gorm.DB { return r.current.Load() }

func (r *reloadableDB) replace(db *gorm.DB) {
	registerDataScopeCallbacks(db)
	old := r.current.Swap(db)
	if old != nil && old != db {
		r.mu.Lock()
		r.retired = append(r.retired, old)
		r.mu.Unlock()
	}
}

func (r *reloadableDB) close() {
	current := r.current.Load()
	r.mu.Lock()
	all := append([]*gorm.DB{current}, r.retired...)
	r.retired = nil
	r.mu.Unlock()
	seen := map[*gorm.DB]struct{}{}
	for _, db := range all {
		if db == nil {
			continue
		}
		if _, ok := seen[db]; ok {
			continue
		}
		seen[db] = struct{}{}
		if sqlDB, err := db.DB(); err == nil {
			_ = sqlDB.Close()
		}
	}
}

type reloadableRedis struct {
	mu      sync.RWMutex
	current redis.UniversalClient
	retired []redis.UniversalClient
}

func newReloadableRedis(client redis.UniversalClient) *reloadableRedis {
	return &reloadableRedis{current: client}
}

func (r *reloadableRedis) load() redis.UniversalClient {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return r.current
}

func (r *reloadableRedis) replace(client redis.UniversalClient) {
	r.mu.Lock()
	old := r.current
	r.current = client
	if old != nil && old != client {
		r.retired = append(r.retired, old)
	}
	r.mu.Unlock()
}

func (r *reloadableRedis) close() {
	r.mu.Lock()
	all := append([]redis.UniversalClient{r.current}, r.retired...)
	r.current = nil
	r.retired = nil
	r.mu.Unlock()
	for _, client := range all {
		if client == nil {
			continue
		}
		_ = client.Close()
	}
}