package server

import (
	"context"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

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

// newObserveTestWorker builds a fallback-backed worker whose claimable task
// type is img_extract_text (capability-only, like resolver-backed shared
// tenants), plus a target pointing at the same store.
func newObserveTestWorker(t *testing.T) (*tenantWorkerManager, *tenantTarget) {
	t.Helper()
	initWorkerReconcileSchema(t)
	store, err := datastore.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = store.Close() })
	b, err := backend.NewWithOptions(store, backend.Options{
		AsyncImageExtract: backend.AsyncImageExtractOptions{Capability: true},
	})
	if err != nil {
		t.Fatal(err)
	}
	m := newTenantWorkerManager(b, nil, nil, nil, TenantWorkerOptions{}, 0)
	if m == nil {
		t.Fatal("expected viable worker manager")
	}
	target := &tenantTarget{
		tenantID:         tenantLocalID,
		backend:          b,
		store:            store,
		allowedTaskTypes: []semantic.TaskType{semantic.TaskTypeImgExtractText},
		release:          func() {},
	}
	t.Cleanup(func() {
		metrics.DeleteTenantGaugeWithOrg(tenantLocalID, "", "semantic_worker", "queue_lag_seconds")
		metrics.DeleteTenantGaugeWithOrg(tenantLocalID, "", "semantic_worker", "dead_lettered")
		metrics.DeleteTenantGaugeWithOrg(tenantLocalID, defaultTenantMetricTiDBCloudOrgID, "semantic_worker", "queue_lag_seconds")
		metrics.DeleteTenantGaugeWithOrg(tenantLocalID, defaultTenantMetricTiDBCloudOrgID, "semantic_worker", "dead_lettered")
	})
	return m, target
}

func enqueueOldImgTask(t *testing.T, store *datastore.Store, taskID string, age time.Duration) {
	t.Helper()
	now := time.Now().UTC()
	created, err := store.EnqueueSemanticTask(context.Background(), &semantic.Task{
		TaskID:          taskID,
		TaskType:        semantic.TaskTypeImgExtractText,
		ResourceID:      "file-" + taskID,
		ResourceVersion: 1,
		Status:          semantic.TaskQueued,
		AvailableAt:     now.Add(-age),
		CreatedAt:       now.Add(-age),
		UpdatedAt:       now.Add(-age),
	})
	if err != nil {
		t.Fatal(err)
	}
	if !created {
		t.Fatal("expected task to be created")
	}
}

func semanticLagGaugePresent(t *testing.T) bool {
	t.Helper()
	rec := httptest.NewRecorder()
	metrics.WritePrometheus(rec)
	return strings.Contains(rec.Body.String(),
		`drive9_service_gauge{component="semantic_worker",name="queue_lag_seconds",tenant_id="`+tenantLocalID+`"`)
}

func TestClaimAndProcessOneReportsQueueEmpty(t *testing.T) {
	m, target := newObserveTestWorker(t)
	ctx := context.Background()

	claimed, empty := m.claimAndProcessOne(ctx, target)
	if claimed || !empty {
		t.Fatalf("empty queue: claimed=%v empty=%v, want false/true", claimed, empty)
	}

	enqueueOldImgTask(t, target.store, "task-observe-claim-signal", time.Minute)
	claimed, empty = m.claimAndProcessOne(ctx, target)
	if !claimed || empty {
		t.Fatalf("queued task: claimed=%v empty=%v, want true/false", claimed, empty)
	}

	// The claimed task is acked (not_configured); the queue is empty again.
	claimed, empty = m.claimAndProcessOne(ctx, target)
	if claimed || !empty {
		t.Fatalf("drained queue: claimed=%v empty=%v, want false/true", claimed, empty)
	}
}

func TestProcessKickedRefreshesSemanticObservationOnDrain(t *testing.T) {
	m, target := newObserveTestWorker(t)
	ctx := context.Background()

	// Backlog with an old claimable task: observation exports a lag gauge.
	enqueueOldImgTask(t, target.store, "task-observe-drain", 10*time.Minute)
	m.observeTenantNow(ctx, target)
	if !semanticLagGaugePresent(t) {
		t.Fatal("queue_lag_seconds gauge missing with claimable backlog")
	}

	// A semantic kick drains the queue; the empty claim at the end must
	// refresh the observation immediately, clearing the lag gauge without
	// waiting for any throttle interval.
	m.processKicked(ctx, tenantLocalID, "", WorkSemantic)

	if semanticLagGaugePresent(t) {
		t.Fatal("queue_lag_seconds gauge still exported after the queue drained")
	}
}

func TestObserveTenantNowDoesNotAdvanceThrottleOnFailure(t *testing.T) {
	m, target := newObserveTestWorker(t)
	ctx := context.Background()

	// A failed observation (closed store) must not advance the throttle: the
	// same-kick trailing observeTenantThrottled (or the next kick) is the
	// retry, so a drained queue cannot be left with a frozen lag gauge.
	badStore, err := datastore.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	if err := badStore.Close(); err != nil {
		t.Fatal(err)
	}
	badTarget := *target
	badTarget.store = badStore

	// Seed a fresh stamp from a prior successful observation: the realistic
	// drain has one, and the failure must clear it so the throttled trailing
	// call does not skip the retry.
	m.lastSemanticObserve[target.tenantID] = time.Now()

	m.observeTenantNow(ctx, &badTarget)
	if _, ok := m.lastSemanticObserve[target.tenantID]; ok {
		t.Fatal("throttle timestamp survived a failed observation")
	}

	// The trailing throttled call in the same kick retries immediately.
	enqueueOldImgTask(t, target.store, "task-observe-retry", 10*time.Minute)
	m.observeTenantThrottled(ctx, target)
	if _, ok := m.lastSemanticObserve[target.tenantID]; !ok {
		t.Fatal("throttled observation did not retry after a failed attempt")
	}
	if !semanticLagGaugePresent(t) {
		t.Fatal("queue_lag_seconds gauge missing after retried observation")
	}
}

func TestObserveTenantThrottledOwnsCadenceIndependentOfMaintenance(t *testing.T) {
	// No t.Parallel: shrinks the package-level observation interval.
	oldInterval := tenantSemanticObserveInterval
	t.Cleanup(func() { tenantSemanticObserveInterval = oldInterval })

	m, target := newObserveTestWorker(t)
	ctx := context.Background()

	// Piggyback maintenance cadence must not drive semantic observation:
	// simulate a long maintenance interval by stamping lastMaintenance now.
	m.lastMaintenance[target.tenantID] = time.Now()

	tenantSemanticObserveInterval = time.Minute
	enqueueOldImgTask(t, target.store, "task-observe-throttle", 10*time.Minute)

	// First observation runs even though piggyback maintenance is not due.
	m.observeTenantThrottled(ctx, target)
	if !semanticLagGaugePresent(t) {
		t.Fatal("throttled observation must publish lag while maintenance is not due")
	}

	// Clear the backlog out-of-band; a throttled second call must NOT
	// re-observe (the gauge stays stale within the interval).
	if _, err := target.store.DB().ExecContext(ctx, `DELETE FROM semantic_tasks`); err != nil {
		t.Fatal(err)
	}
	m.observeTenantThrottled(ctx, target)
	if !semanticLagGaugePresent(t) {
		t.Fatal("observation within the throttle interval must be skipped")
	}

	// Once the interval elapses the gauge refreshes and clears.
	tenantSemanticObserveInterval = 0
	m.observeTenantThrottled(ctx, target)
	if semanticLagGaugePresent(t) {
		t.Fatal("queue_lag_seconds gauge must clear after the throttle interval")
	}
}