package server

import (
	"context"

	"go.uber.org/zap"

	"github.com/mem9-ai/drive9/pkg/logger"
	"github.com/mem9-ai/drive9/pkg/metrics"
)

func (m *tenantWorkerManager) setOwnsTenant(fn func(string) bool) {
	if m == nil {
		return
	}
	m.ownsTenant = fn
}

func (m *tenantWorkerManager) ownsFileGC(tenantID string) bool {
	return m.ownsShard(tenantID)
}

// ownsShard reports whether this pod currently owns the tenant's shard for
// background (non-request) work. A nil manager or nil resolver fn means
// single-pod mode: owned.
func (m *tenantWorkerManager) ownsShard(tenantID string) bool {
	if m == nil || m.ownsTenant == nil {
		return true
	}
	return m.ownsTenant(tenantID)
}

// dropUnownedFileGCGauges deletes this process's FileGC backlog series for
// tenants the current shard ring no longer assigns here. It does not touch
// still-owned tenants (including an in-flight drain). Mismatch re-kick
// timers are included so a claim/COUNT stall with omitted gauges still
// loses its backoff when ownership moves.
func (m *tenantWorkerManager) dropUnownedFileGCGauges() {
	if m == nil {
		return
	}
	seen := make(map[string]struct{})
	var tenants []string
	add := func(tenantID string) {
		if tenantID == "" {
			return
		}
		if _, ok := seen[tenantID]; ok {
			return
		}
		seen[tenantID] = struct{}{}
		tenants = append(tenants, tenantID)
	}
	for _, tenantID := range metrics.FileGCBacklogTenantIDs() {
		add(tenantID)
	}
	m.mu.Lock()
	for tenantID := range m.fileGCMismatchTimers {
		add(tenantID)
	}
	for tenantID := range m.fileGCMismatchAttempts {
		add(tenantID)
	}
	m.mu.Unlock()
	dropped := 0
	for _, tenantID := range tenants {
		if m.ownsFileGC(tenantID) {
			continue
		}
		metrics.ClearFileGCClaimableBacklog(tenantID)
		m.clearFileGCMismatchBackoff(tenantID)
		dropped++
	}
	if dropped > 0 {
		logger.Info(context.Background(), "file_gc_backlog_gauges_dropped_unowned",
			zap.Int("tenants", dropped))
	}
}

// withOwnedFileGCGauges owner-fences FileGC gauge writes on this process.
// Non-owners delete leftover series instead of publishing; a ring change
// between the two ownership checks also drops the just-written series.
func (m *tenantWorkerManager) withOwnedFileGCGauges(tenantID string, publish func()) {
	if m == nil || tenantID == "" || publish == nil {
		return
	}
	if !m.ownsFileGC(tenantID) {
		metrics.ClearFileGCClaimableBacklog(tenantID)
		return
	}
	defer func() {
		if !m.ownsFileGC(tenantID) {
			metrics.ClearFileGCClaimableBacklog(tenantID)
		}
	}()
	publish()
}