package server

import (
	"context"
	"time"

	"go.uber.org/zap"

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

// tenantSemanticObserveInterval is the throttle interval for the per-tenant
// semantic queue observation (queue_lag_seconds / dead_lettered gauges). It is
// deliberately decoupled from MaintenanceInterval: those gauges feed lag
// alerting, so they must stay fresh even when an operator stretches the
// maintenance cadence (fs_events cleanup etc.) to hours or days. A
// package-level variable so tests can shrink it.
var tenantSemanticObserveInterval = time.Minute

// observeTenantNow refreshes the tenant's semantic queue observation
// immediately. The throttle timestamp advances only on a successful
// observation. On failure any previous stamp is deleted, so the trailing
// observeTenantThrottled call in the same kick (or the next kick) retries
// instead of skipping: a drained queue must not keep a frozen lag gauge
// because the "clear it now" query hit a transient error.
func (m *tenantWorkerManager) observeTenantNow(ctx context.Context, target *tenantTarget) {
	now := time.Now()
	if err := m.observeTenant(ctx, target, now); err != nil {
		m.mu.Lock()
		delete(m.lastSemanticObserve, target.tenantID)
		m.mu.Unlock()
		return
	}
	m.mu.Lock()
	if m.lastSemanticObserve == nil {
		m.lastSemanticObserve = make(map[string]time.Time)
	}
	m.lastSemanticObserve[target.tenantID] = now
	m.mu.Unlock()
}

// observeTenantThrottled samples the tenant's semantic queue state at most
// once per tenantSemanticObserveInterval. It runs on every processed kick
// (not only WorkSemantic ones) so a warm tenant's lag gauge reflects reality
// within one interval even when MaintenanceInterval is configured to hours.
func (m *tenantWorkerManager) observeTenantThrottled(ctx context.Context, target *tenantTarget) {
	if ctx.Err() != nil || target == nil || target.store == nil {
		return
	}
	m.mu.Lock()
	last := m.lastSemanticObserve[target.tenantID]
	due := time.Since(last) >= tenantSemanticObserveInterval
	m.mu.Unlock()
	if !due {
		return
	}
	m.observeTenantNow(ctx, target)
}

// observeTenant samples the tenant's semantic queue state and publishes the
// lag/dead-letter gauges. It returns the observation error (if any) so the
// caller can decide whether to advance the throttle timestamp.
func (m *tenantWorkerManager) observeTenant(ctx context.Context, target *tenantTarget, now time.Time) error {
	obs, err := target.store.ObserveSemanticTasks(ctx, now.UTC())
	if err != nil {
		if ctx.Err() == nil {
			logger.Warn(ctx, "tenant_worker_observe_failed",
				zap.String("tenant_id", target.tenantID), zap.Error(err))
		}
		return err
	}
	m.recordSemanticWorkerObservation(target.tenantID, target.metricOrgID(), obs, now)
	return nil
}

func (m *tenantWorkerManager) recordSemanticWorkerObservation(tenantID, tidbCloudOrgID string, obs *datastore.SemanticTaskObservation, now time.Time) {
	tidbCloudOrgID = normalizeTenantMetricTiDBCloudOrgID(tidbCloudOrgID)
	m.mu.Lock()
	previousOrgID := m.semanticMetricOrg[tenantID]
	if m.semanticMetricOrg == nil {
		m.semanticMetricOrg = make(map[string]string)
	}
	m.semanticMetricOrg[tenantID] = tidbCloudOrgID
	m.mu.Unlock()
	if previousOrgID != "" && previousOrgID != tidbCloudOrgID {
		metrics.DeleteTenantGaugeWithOrg(tenantID, previousOrgID, "semantic_worker", "dead_lettered")
		metrics.DeleteTenantGaugeWithOrg(tenantID, previousOrgID, "semantic_worker", "queue_lag_seconds")
	}
	recordSemanticWorkerObservation(tenantID, tidbCloudOrgID, obs, now)
}

func recordSemanticWorkerObservation(tenantID, tidbCloudOrgID string, obs *datastore.SemanticTaskObservation, now time.Time) {
	if obs == nil {
		return
	}
	if obs.DeadLettered > 0 {
		metrics.RecordTenantGaugeWithOrg(tenantID, tidbCloudOrgID, "semantic_worker", "dead_lettered", float64(obs.DeadLettered))
	} else {
		metrics.DeleteTenantGaugeWithOrg(tenantID, tidbCloudOrgID, "semantic_worker", "dead_lettered")
	}
	tenantLag := float64(0)
	if obs.OldestClaimableAvailableAt != nil {
		tenantLag = now.UTC().Sub(obs.OldestClaimableAvailableAt.UTC()).Seconds()
		if tenantLag < 0 {
			tenantLag = 0
		}
	}
	if tenantLag > 0 {
		metrics.RecordTenantGaugeWithOrg(tenantID, tidbCloudOrgID, "semantic_worker", "queue_lag_seconds", tenantLag)
	} else {
		metrics.DeleteTenantGaugeWithOrg(tenantID, tidbCloudOrgID, "semantic_worker", "queue_lag_seconds")
	}
}