package source

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

	merrors "git.sendico.io/tech/sxchnge/shared/errors"
	mlogger "git.sendico.io/tech/sxchnge/shared/logging"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
	"go.mongodb.org/mongo-driver/v2/mongo/readpref"
	"go.uber.org/zap"
)

// Lease owns a credential and its revocation capability. Neither may be logged.
type Lease struct {
	URI       string
	ExpiresAt time.Time
	Revoke    func(context.Context) error
}
type Issuer interface {
	Issue(context.Context) (Lease, error)
}

type generation struct {
	client  *mongo.Client
	lease   Lease
	refs    int
	retired bool
	drained chan struct{}
}
type pin struct {
	manager    *Manager
	generation *generation
	active     atomic.Bool
}
type pinKey struct{}

// Manager publishes a verified replacement before retiring the previous pool.
// Every cursor, operation and transaction retains its originating generation.
type Manager struct {
	mu         sync.Mutex
	current    *generation
	closed     bool
	issuer     Issuer
	database   string
	logger     mlogger.Logger
	cancel     context.CancelFunc
	workerDone chan struct{}
	draining   sync.WaitGroup
}

func New(ctx context.Context, logger mlogger.Logger, database string, issuer Issuer) (*Manager, error) {
	if database == "" || issuer == nil {
		return nil, merrors.InvalidArgument("database and credential issuer are required")
	}
	life, cancel := context.WithCancel(context.Background())
	m := &Manager{issuer: issuer, database: database, logger: logger, cancel: cancel, workerDone: make(chan struct{})}
	g, err := m.prepare(ctx)
	if err != nil {
		cancel()
		return nil, err
	}
	m.current = g
	go m.run(life)
	return m, nil
}

func (m *Manager) Acquire(ctx context.Context) (*mongo.Database, context.Context, func(), error) {
	if err := ctx.Err(); err != nil {
		return nil, ctx, nil, err
	}
	if p, ok := ctx.Value(pinKey{}).(*pin); ok && p.manager == m {
		if !p.active.Load() {
			return nil, ctx, nil, merrors.FailedPrecondition("database operation scope has ended")
		}
		return p.generation.client.Database(m.database), ctx, func() {}, nil
	}
	m.mu.Lock()
	g := m.current
	if m.closed || g == nil || !time.Now().Before(g.lease.ExpiresAt) {
		m.mu.Unlock()
		return nil, ctx, nil, merrors.ServiceUnavailable("no valid database credential is available")
	}
	if session := mongo.SessionFromContext(ctx); session != nil && session.Client() != g.client {
		m.mu.Unlock()
		return nil, ctx, nil, merrors.FailedPrecondition("transaction has no matching database operation scope")
	}
	g.refs++
	m.mu.Unlock()
	p := &pin{manager: m, generation: g}
	p.active.Store(true)
	bounded, cancel := context.WithDeadline(ctx, g.lease.ExpiresAt)
	scoped := context.WithValue(bounded, pinKey{}, p)
	release := func() {
		if !p.active.CompareAndSwap(true, false) {
			return
		}
		cancel()
		m.mu.Lock()
		g.refs--
		if g.retired && g.refs == 0 {
			close(g.drained)
		}
		m.mu.Unlock()
	}
	return g.client.Database(m.database), scoped, release, nil
}

func (m *Manager) prepare(ctx context.Context) (*generation, error) {
	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	lease, err := m.issuer.Issue(ctx)
	if err != nil {
		return nil, err
	}
	var client *mongo.Client
	success := false
	defer func() {
		if !success {
			if client != nil {
				cleanup, cancel := context.WithTimeout(context.Background(), 5*time.Second)
				_ = client.Disconnect(cleanup)
				cancel()
			}
			m.draining.Add(1)
			go func() { defer m.draining.Done(); m.revoke(lease) }()
		}
	}()
	if lease.Revoke == nil || time.Until(lease.ExpiresAt) < 2*time.Second {
		return nil, merrors.FailedPrecondition("database credential lifetime is too short")
	}
	client, err = mongo.Connect(options.Client().ApplyURI(lease.URI).SetServerSelectionTimeout(5 * time.Second))
	if err != nil {
		return nil, merrors.OpaqueWrap(merrors.ErrServiceUnavailable, err, "create replacement database pool")
	}
	if err = client.Ping(ctx, readpref.Primary()); err != nil {
		return nil, merrors.OpaqueWrap(merrors.ErrServiceUnavailable, err, "validate replacement database pool")
	}
	if time.Until(lease.ExpiresAt) < time.Second {
		return nil, merrors.FailedPrecondition("replacement database credential expired during validation")
	}
	success = true
	return &generation{client: client, lease: lease, drained: make(chan struct{})}, nil
}

func (m *Manager) retireLocked(g *generation) {
	g.retired = true
	if g.refs == 0 {
		close(g.drained)
	}
	m.draining.Add(1)
	go func() {
		defer m.draining.Done()
		<-g.drained
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		if err := g.client.Disconnect(ctx); err != nil {
			m.logger.Warn("Retired database pool close failed")
		}
		m.revoke(g.lease)
		m.logger.Info("Retired database pool drained")
	}()
}

func (m *Manager) revoke(lease Lease) {
	if lease.Revoke == nil {
		return
	}
	// Expiration is the backstop when Vault is unavailable during bounded cleanup.
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	if err := lease.Revoke(ctx); err != nil {
		m.logger.Warn("Database credential cleanup failed; Vault expiry will revoke it")
	}

}

func (m *Manager) run(ctx context.Context) {
	defer close(m.workerDone)
	m.mu.Lock()
	wait := time.Until(m.current.lease.ExpiresAt) / 2
	m.mu.Unlock()
	backoff := 250 * time.Millisecond
	failureReported, expiryReported := false, false
	for {
		timer := time.NewTimer(wait)
		select {
		case <-ctx.Done():
			timer.Stop()
			return
		case <-timer.C:
		}
		next, err := m.prepare(ctx)
		if err != nil {
			if ctx.Err() != nil {
				return
			}
			m.mu.Lock()
			remaining := time.Until(m.current.lease.ExpiresAt)
			m.mu.Unlock()
			if !failureReported {
				m.logger.Warn("Database credential refresh failed; retaining current pool", zap.Duration("valid_for", max(remaining, 0)))
				failureReported = true
			}
			if remaining <= 0 && !expiryReported {
				m.logger.Error("Database credentials expired; rejecting new operations until refresh succeeds")
				expiryReported = true
			}
			wait = backoff
			if remaining > 0 {
				wait = max(100*time.Millisecond, min(wait, remaining/4))
			}
			backoff = min(5*time.Second, backoff*2)

			continue
		}
		m.mu.Lock()
		if m.closed {
			m.retireLocked(next)
			m.mu.Unlock()
			return
		}
		old := m.current
		m.current = next
		m.retireLocked(old)
		m.mu.Unlock()
		failureReported, expiryReported = false, false
		backoff = 250 * time.Millisecond
		m.logger.Info("Database pool replaced without restarting service")
		wait = time.Until(next.lease.ExpiresAt) / 2
	}
}

func (m *Manager) Close(ctx context.Context) error {
	m.mu.Lock()
	if !m.closed {
		m.closed = true
		m.cancel()
		m.retireLocked(m.current)
	}
	m.mu.Unlock()
	done := make(chan struct{})
	go func() { <-m.workerDone; m.draining.Wait(); close(done) }()
	select {
	case <-done:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}