package server

import (
	"context"
	"errors"
	"fmt"
	"time"

	"go.uber.org/zap"

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

// outboxKicker is the interface the poller uses to dispatch sharded work
// (semantic/file_gc/quota) to the unified worker. The worker deduplicates
// kicks by tenant and OR-accumulates the work_mask.
type outboxKicker interface {
	KickWithOrg(tenantID, tidbCloudOrgID string, workMask int)
}

type warmStorageTierScanKicker interface {
	KickWarmStorageTierScan(tenantID string)
}

type outboxTenantForgetter interface {
	ForgetTenant(tenantID string)
}

const (
	// defaultTenantOutboxPollInterval is the poller's tick interval. At 200ms,
	// cross-pod work delivery is bounded to ~200ms. The target is the
	// always-provisioned meta DB, so 5 QPS per pod is negligible.
	defaultTenantOutboxPollInterval = 200 * time.Millisecond
	// tenantOutboxBatchSize is the max rows read per ListTenantNotifySince call.
	// If a batch is full, the poller immediately reads the next batch (drain
	// mode) so a burst of writes is drained quickly.
	tenantOutboxBatchSize = 1000
	// defaultTenantOutboxCursorFlushInterval is how often the poller persists
	// its in-memory cursor to tenant_outbox_cursor. At 5s, a restart loses at
	// most 5s of work signals — the safety-net scan recovers any expired
	// leases that were missed.
	defaultTenantOutboxCursorFlushInterval = 5 * time.Second
)

// tenantOutboxPoller is the single per-pod goroutine that reads the unified
// tenant_notify_outbox table (in the always-provisioned meta DB) and
// dispatches work by work_mask:
//   - SSE bit → wake the local EventBus (broadcast: all pods with subscribers)
//   - Metrics cleanup bit → delete this pod's tenant-scoped metrics
//   - Semantic/FileGC/Quota bits → kick the unified worker if this pod's shard
//     resolver owns the tenant (sharded: shard owner only)
//
// This replaces notify_poller.go (SSE-only) and eliminates all periodic
// per-tenant TiDB scanning. Idle tenant TiDBs can scale to zero because the
// poller never touches them — it only reads the central meta DB.
type tenantOutboxPoller struct {
	metaStore                  *meta.Store
	buses                      *eventBuses
	worker                     outboxKicker
	shardFn                    func(string) bool
	tenantFilter               func(string) bool
	podID                      string
	interval                   time.Duration
	cursorFlushEvery           time.Duration
	embeddingConfigInvalidator func(string)

	lastID uint64
}

func (p *tenantOutboxPoller) withEmbeddingConfigInvalidator(invalidate func(string)) *tenantOutboxPoller {
	p.embeddingConfigInvalidator = invalidate
	return p
}

// newTenantOutboxPoller creates a tenantOutboxPoller. metaStore must be non-nil.
// shardFn returns true when this pod owns sharded work for the tenant. When
// shardFn is nil, this pod owns all sharded work (single-pod mode).
// tenantFilter restricts the poller to a subset of tenants; when nil, all
// tenants are processed. Single-tenant local mode sets it to the local tenant
// so the process never writes durable state or kicks for other tenants.
func newTenantOutboxPoller(
	metaStore *meta.Store,
	buses *eventBuses,
	worker outboxKicker,
	shardFn func(string) bool,
	podID string,
	interval time.Duration,
	cursorFlushEvery time.Duration,
) *tenantOutboxPoller {
	if interval <= 0 {
		interval = defaultTenantOutboxPollInterval
	}
	if cursorFlushEvery <= 0 {
		cursorFlushEvery = defaultTenantOutboxCursorFlushInterval
	}
	if shardFn == nil {
		shardFn = func(string) bool { return true }
	}
	return &tenantOutboxPoller{
		metaStore:        metaStore,
		buses:            buses,
		worker:           worker,
		shardFn:          shardFn,
		podID:            podID,
		interval:         interval,
		cursorFlushEvery: cursorFlushEvery,
	}
}

// withTenantFilter returns the poller with the tenant filter set. A nil filter
// processes all tenants; a non-nil filter skips any tenant for which it returns
// false (no durable pending write, no kick, no scan). This is used by
// single-tenant local mode to isolate the process from other tenants that may
// share the same meta DB.
func (p *tenantOutboxPoller) withTenantFilter(f func(string) bool) *tenantOutboxPoller {
	p.tenantFilter = f
	return p
}

// initCursor synchronously initializes the poller's lastID before run starts
// and before the server accepts live traffic. On restart, the cursor is
// recovered from tenant_outbox_cursor so no work is skipped. On first launch
// (no cursor row), the cursor is set to MAX(id) so historical rows are skipped
// — the pod never owned work before its first start.
func (p *tenantOutboxPoller) initCursor(ctx context.Context) {
	if p.podID != "" {
		cursor, err := p.metaStore.GetTenantOutboxCursor(ctx, p.podID)
		if err == nil && cursor != nil {
			p.lastID = cursor.LastID
			logger.Info(ctx, "tenant_outbox_poller_cursor_recovered",
				zap.String("pod_id", p.podID),
				zap.Uint64("cursor", p.lastID))
			return
		}
		if errors.Is(err, meta.ErrNotFound) {
			// First launch — no prior cursor. Fall through to MAX(id) below.
		} else if err != nil {
			// Non-ErrNotFound error: do NOT fall through to MAX(id), which would
			// skip historical rows. Start from 0 instead so all unpruned rows are
			// re-read; duplicate kicks are harmless (work is idempotent/deduped).
			logger.Warn(ctx, "tenant_outbox_poller_cursor_read_failed",
				zap.String("pod_id", p.podID),
				zap.Error(err))
			p.lastID = 0
			logger.Info(ctx, "tenant_outbox_poller_cursor_initialized",
				zap.Uint64("cursor", p.lastID))
			return
		}
	}
	maxID, err := p.metaStore.MaxTenantNotifyID(ctx)
	if err != nil {
		logger.Warn(ctx, "tenant_outbox_poller_init_cursor_failed", zap.Error(err))
		maxID = 0
	}
	p.lastID = maxID
	logger.Info(ctx, "tenant_outbox_poller_cursor_initialized",
		zap.Uint64("cursor", p.lastID))
}

// run starts the poller goroutine. It blocks until ctx is cancelled. The cursor
// (lastID) must already be initialized by initCursor before run is called.
func (p *tenantOutboxPoller) run(ctx context.Context) {
	logger.Info(ctx, "tenant_outbox_poller_started",
		zap.Duration("interval", p.interval),
		zap.Uint64("cursor", p.lastID),
		zap.Duration("cursor_flush_every", p.cursorFlushEvery))

	// Recover due delete-reconcile requests immediately on startup (a pending
	// row that became due while this pod was down must not wait for the next
	// cursor-flush tick).
	p.scanDeleteReconcilePending(ctx)
	// A page commit intentionally leaves the scan job pending until the next
	// kick. Rediscover incomplete jobs on startup as well, so a crash between
	// those two operations cannot strand a warm tenant's scan.
	p.scanStorageTierPending(ctx)

	ticker := time.NewTicker(p.interval)
	defer ticker.Stop()
	flushTicker := time.NewTicker(p.cursorFlushEvery)
	defer flushTicker.Stop()
	for {
		select {
		case <-ctx.Done():
			// Flush cursor on shutdown using a fresh non-cancelled context so
			// the final cursor position is persisted. Using the already-cancelled
			// ctx would cause UpsertTenantOutboxCursor to fail with
			// context.Canceled and the cursor would be lost — a pod that shuts
			// down before its first periodic flush would restart at MAX(id),
			// skipping unprocessed rows.
			shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
			p.flushCursor(shutdownCtx)
			shutdownCancel()
			return
		case <-ticker.C:
			p.pollOnce(ctx)
		case <-flushTicker.C:
			p.flushCursor(ctx)
			// Durable delete-reconcile recovery: pending rows survive pod
			// restarts and are re-dispatched here until the sweep succeeds.
			p.scanDeleteReconcilePending(ctx)
			p.scanStorageTierPending(ctx)
		}
	}
}

// pollOnce reads new outbox rows and dispatches by work_mask. If a full batch
// is returned, it immediately reads the next batch (drain mode) so a burst of
// events is processed without waiting for the next tick.
func (p *tenantOutboxPoller) pollOnce(ctx context.Context) {
	for {
		if ctx.Err() != nil {
			return
		}
		pollStart := time.Now()
		rows, err := p.metaStore.ListTenantNotifySince(ctx, p.lastID, tenantOutboxBatchSize)
		if err != nil {
			metrics.RecordTenantOutboxPoll(metrics.ResultForError(err), time.Since(pollStart), 0, 0, false)
			logger.Warn(ctx, "tenant_outbox_poller_list_failed",
				zap.Uint64("after_id", p.lastID),
				zap.Error(err))
			// List failure: the meta-DB read that drives kick dispatch failed.
			// Record as a user_db_access path-level signal (the poller itself
			// only reads the meta DB, but a list failure delays all downstream
			// tenant-DB work). Avoids the previous abuse of putting the pod
			// name in the tenant_id label of drive9_event_bus_poll_failures_total.
			metrics.RecordOperation("user_db_access", "outbox_poll_list", "error", 0)
			return
		}
		for _, row := range rows {
			if err := p.dispatch(ctx, row); err != nil {
				// The durable pending write failed for a delete-reconcile
				// signal. Do NOT advance the cursor past this row: consuming
				// the once-read outbox row without a durable pending marker
				// would lose the request on restart. The row is re-read on the
				// next poll.
				logger.Warn(ctx, "tenant_outbox_poller_dispatch_failed",
					zap.Uint64("row_id", row.ID),
					zap.String("tenant_id", row.TenantID),
					zap.Error(err))
				return
			}
			p.lastID = row.ID
		}
		oldestAge := time.Duration(0)
		if len(rows) > 0 {
			oldestAge = time.Since(rows[0].CreatedAt)
			if oldestAge < 0 {
				oldestAge = 0
			}
		}
		metrics.RecordTenantOutboxPoll("ok", time.Since(pollStart), len(rows), oldestAge, len(rows) == tenantOutboxBatchSize)
		// If we got fewer rows than the batch size, we're caught up — wait for
		// the next tick. Otherwise drain the remaining rows immediately.
		if len(rows) < tenantOutboxBatchSize {
			return
		}
	}
}

// dispatch sends a single outbox row to the right consumer based on work_mask.
// It returns an error only when a WorkDeleteReconcile row could not be made
// durable in delete_reconcile_pending; the caller must not advance the cursor
// past such a row.
func (p *tenantOutboxPoller) dispatch(ctx context.Context, row meta.TenantNotifyRow) error {
	if p.tenantFilter != nil && !p.tenantFilter(row.TenantID) {
		// Single-tenant local mode: ignore rows for any tenant other than the
		// local tenant. No durable pending write, no kick, no scan.
		return nil
	}
	if row.WorkMask&WorkAPIKeyCacheCleanup != 0 && p.metaStore != nil {
		p.metaStore.InvalidateAPIKeyResolveCacheForTenant(row.TenantID)
	}
	if row.WorkMask&WorkMetricsCleanup != 0 {
		metrics.DeleteTenantCounters(row.TenantID)
		if p.metaStore != nil {
			p.metaStore.InvalidateAPIKeyResolveCacheForTenant(row.TenantID)
		}
		if forgetter, ok := p.worker.(outboxTenantForgetter); ok {
			forgetter.ForgetTenant(row.TenantID)
		}
	}
	if row.WorkMask&(WorkMetricsCleanup|WorkEmbeddingConfigCleanup) != 0 && p.embeddingConfigInvalidator != nil {
		p.embeddingConfigInvalidator(row.TenantID)
	}
	if row.WorkMask&WorkSSE != 0 {
		// SSE is broadcast: wake any local bus with subscribers for this tenant.
		// If no bus exists, skip — we never touch the tenant's TiDB.
		if bus := p.buses.getIfExists(row.TenantID); bus != nil {
			bus.Publish()
		}
	}
	shardedMask := row.WorkMask & (WorkSemantic | WorkFileGC | WorkDeleteReconcile)
	if row.WorkMask&WorkStorageTierScan != 0 && p.worker != nil {
		// Broadcast delivery is safe because each worker is warm-only and the
		// Central DB lease permits only one warm pod to scan at a time.
		p.worker.KickWithOrg(row.TenantID, row.TiDBCloudOrgID, WorkStorageTierScan)
	}
	// Make a WorkDeleteReconcile signal durable on EVERY pod that reads the
	// row, owner or not, BEFORE any cursor advances. Shard ownership gates the
	// kick, never whether durability is recorded: if the current owner dies
	// before reading the row (or before its poller advances), the next owner
	// still finds the pending marker via scanDeleteReconcilePending. A failed
	// write returns an error so the caller does not advance past the row.
	// A lagging pod may re-create a marker the owner already swept and cleared
	// (INSERT IGNORE with due=now); that yields at most one redundant no-op
	// sweep per delete, which is harmless and self-clearing.
	if row.WorkMask&WorkDeleteReconcile != 0 && p.metaStore != nil {
		if err := p.metaStore.InsertDeleteReconcilePendingIfAbsent(ctx, row.TenantID, time.Now()); err != nil {
			return fmt.Errorf("persist delete-reconcile pending for %s: %w", row.TenantID, err)
		}
	}
	if shardedMask != 0 && p.worker != nil {
		// Sharded work: only the shard owner processes it. Other pods skip;
		// the safety-net scan recovers any work whose kick was lost.
		if p.shardFn(row.TenantID) {
			p.worker.KickWithOrg(row.TenantID, row.TiDBCloudOrgID, shardedMask)
			// Record the kick that will trigger a tenant-DB access via the
			// worker. This gives a baseline rate of "kicks dispatched" to
			// compare against actual tenant_worker_acquire — a divergence
			// (kicks without acquires) points to lost work.
			metrics.RecordTenantOperationCountWithOrg(row.TenantID, row.TiDBCloudOrgID, "user_db_access", "outbox_dispatch_kick", "ok")
		}
	}
	return nil
}

// scanDeleteReconcilePending re-dispatches durable delete-reconcile requests
// whose due time has passed. This is the restart-safe recovery path for
// WorkDeleteReconcile: the original outbox row is consumed once, but the
// pending row survives restarts and is re-kicked here until the worker
// completes a successful sweep (which deletes the row). Runs once at poller
// startup and then on the cursor-flush cadence. Recovered kicks pass an empty
// org ID (the pending table does not persist it), so their kick metric is
// attributed to the guest org; the worker still resolves the tenant's real
// org for tenant-scoped operations.
func (p *tenantOutboxPoller) scanDeleteReconcilePending(ctx context.Context) {
	if p.metaStore == nil || p.worker == nil {
		return
	}
	tenants, err := p.metaStore.ListDeleteReconcilePendingDue(ctx, time.Now(), tenantOutboxBatchSize)
	if err != nil {
		if !isContextDoneErr(err) {
			logger.Warn(ctx, "tenant_outbox_poller_delete_reconcile_scan_failed", zap.Error(err))
		}
		return
	}
	for _, tenantID := range tenants {
		if p.tenantFilter != nil && !p.tenantFilter(tenantID) {
			continue
		}
		if p.shardFn(tenantID) {
			p.worker.KickWithOrg(tenantID, "", WorkDeleteReconcile)
		}
	}
}

// scanStorageTierPending re-dispatches incomplete storage-tier scans. This
// reads only the central meta DB; WorkStorageTierScan consumers use the
// worker's AcquireCached path and therefore do not activate cold tenant DBs.
func (p *tenantOutboxPoller) scanStorageTierPending(ctx context.Context) {
	if p.metaStore == nil || p.worker == nil {
		return
	}
	tenantIDs, err := p.metaStore.ListIncompleteStorageTierScanTenants(ctx, tenantOutboxBatchSize)
	if err != nil {
		if !isContextDoneErr(err) {
			logger.Warn(ctx, "tenant_outbox_poller_storage_tier_scan_failed", zap.Error(err))
		}
		return
	}
	for _, tenantID := range tenantIDs {
		if p.tenantFilter != nil && !p.tenantFilter(tenantID) {
			continue
		}
		if warmKicker, ok := p.worker.(warmStorageTierScanKicker); ok {
			warmKicker.KickWarmStorageTierScan(tenantID)
		} else {
			p.worker.KickWithOrg(tenantID, "", WorkStorageTierScan)
		}
	}
}

// flushCursor persists the current in-memory cursor to tenant_outbox_cursor
// so a restart resumes from the last processed row. Best-effort: errors are
// logged and the in-memory cursor remains valid.
func (p *tenantOutboxPoller) flushCursor(ctx context.Context) {
	if p.podID == "" {
		return
	}
	if err := p.metaStore.UpsertTenantOutboxCursor(ctx, p.podID, p.lastID); err != nil {
		metrics.RecordTenantOutboxCursorFlush(metrics.ResultForError(err))
		if ctx.Err() == nil {
			logger.Warn(ctx, "tenant_outbox_poller_cursor_flush_failed",
				zap.String("pod_id", p.podID),
				zap.Uint64("cursor", p.lastID),
				zap.Error(err))
		}
		return
	}
	metrics.RecordTenantOutboxCursorFlush("ok")
}