package activitypool

import (
	"context"
	"errors"
	"sync"
	"testing"
	"time"

	extractionDomain "proxy-pool/internal/domain/extraction"
	ownershipDomain "proxy-pool/internal/domain/ownership"
	proxyDomain "proxy-pool/internal/domain/proxy"
)

func TestMemoryPoolUpsertAppliesProviderTTLAndRefreshesWithoutGrowth(t *testing.T) {
	now := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC)
	pool := NewMemoryPool()
	batch := FetchedBatch{
		ObservedAt:             now,
		ConfiguredTTL:          30 * time.Second,
		AllocationSafetyMargin: 3 * time.Second,
		MaxSize:                100,
		Proxies: []proxyDomain.Proxy{{
			Scheme: proxyDomain.SchemeHTTP,
			Host:   "192.0.2.10",
			Port:   8080,
			State:  proxyDomain.StateAvailable,
		}},
	}

	first, err := pool.UpsertFetched(context.Background(), "provider-a", batch)
	if err != nil {
		t.Fatalf("UpsertFetched(first): %v", err)
	}
	if first.Inserted != 1 || first.Refreshed != 0 || first.Dropped != 0 {
		t.Fatalf("first result = %+v", first)
	}

	entries := pool.Snapshot(now)
	if len(entries) != 1 {
		t.Fatalf("Snapshot() entries = %d, want 1", len(entries))
	}
	if !entries[0].Proxy.ExpiresAt.Equal(now.Add(30 * time.Second)) {
		t.Fatalf("ExpiresAt = %v, want %v", entries[0].Proxy.ExpiresAt, now.Add(30*time.Second))
	}
	if !entries[0].UsableUntil.Equal(now.Add(27 * time.Second)) {
		t.Fatalf("UsableUntil = %v, want %v", entries[0].UsableUntil, now.Add(27*time.Second))
	}
	if entries[0].Proxy.UsableUntil == nil || !entries[0].Proxy.UsableUntil.Equal(now.Add(27*time.Second)) {
		t.Fatalf("Proxy.UsableUntil = %v, want %v", entries[0].Proxy.UsableUntil, now.Add(27*time.Second))
	}

	batch.ObservedAt = now.Add(10 * time.Second)
	second, err := pool.UpsertFetched(context.Background(), "provider-a", batch)
	if err != nil {
		t.Fatalf("UpsertFetched(refresh): %v", err)
	}
	if second.Inserted != 0 || second.Refreshed != 1 || len(pool.Snapshot(batch.ObservedAt)) != 1 {
		t.Fatalf("refresh result = %+v, entries=%d", second, len(pool.Snapshot(batch.ObservedAt)))
	}
	if got := pool.Snapshot(batch.ObservedAt)[0].Proxy.ExpiresAt; !got.Equal(now.Add(40 * time.Second)) {
		t.Fatalf("refreshed ExpiresAt = %v, want %v", got, now.Add(40*time.Second))
	}
}

func TestMemoryPoolCrossProviderDuplicateDoesNotReplaceSourceLifecycle(t *testing.T) {
	now := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC)
	pool := NewMemoryPool()
	proxy := proxyDomain.Proxy{
		Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080,
		State: proxyDomain.StateAvailable,
	}
	first, err := pool.UpsertFetched(context.Background(), "provider-a", FetchedBatch{
		ObservedAt: now, ConfiguredTTL: 30 * time.Second, MaxSize: 100,
		AllocationSafetyMargin: 3 * time.Second, Proxies: []proxyDomain.Proxy{proxy},
	})
	if err != nil || first.Inserted != 1 {
		t.Fatalf("UpsertFetched(provider-a) = %+v, %v", first, err)
	}

	duplicate, err := pool.UpsertFetched(context.Background(), "provider-b", FetchedBatch{
		ObservedAt: now.Add(time.Second), ConfiguredTTL: 5 * time.Minute, MaxSize: 100,
		AllocationSafetyMargin: 10 * time.Second, Proxies: []proxyDomain.Proxy{proxy},
	})
	if err != nil || duplicate.Inserted != 0 || duplicate.Refreshed != 1 {
		t.Fatalf("UpsertFetched(provider-b) = %+v, %v", duplicate, err)
	}

	entries := pool.Snapshot(now.Add(time.Second))
	if len(entries) != 1 {
		t.Fatalf("Snapshot() entries = %d, want 1", len(entries))
	}
	entry := entries[0]
	if entry.Proxy.SourceUpstream != "provider-a" {
		t.Fatalf("SourceUpstream = %q, want provider-a", entry.Proxy.SourceUpstream)
	}
	if !entry.Proxy.ExpiresAt.Equal(now.Add(30*time.Second)) ||
		!entry.UsableUntil.Equal(now.Add(27*time.Second)) {
		t.Fatalf("lifecycle = expires %v usable %v, want provider-a lifecycle", entry.Proxy.ExpiresAt, entry.UsableUntil)
	}

	afterExpiry, err := pool.UpsertFetched(context.Background(), "provider-b", FetchedBatch{
		ObservedAt: now.Add(31 * time.Second), ConfiguredTTL: 5 * time.Minute, MaxSize: 100,
		AllocationSafetyMargin: 10 * time.Second, Proxies: []proxyDomain.Proxy{proxy},
	})
	if err != nil || afterExpiry.Inserted != 1 {
		t.Fatalf("UpsertFetched(provider-b after expiry) = %+v, %v", afterExpiry, err)
	}
	replacement := pool.Snapshot(now.Add(31 * time.Second))[0]
	if replacement.Proxy.SourceUpstream != "provider-b" {
		t.Fatalf("replacement SourceUpstream = %q, want provider-b", replacement.Proxy.SourceUpstream)
	}
}

func TestMemoryPoolRefreshPreservesRuntimeStateAndHealth(t *testing.T) {
	now := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC)
	checkedAt := now.Add(-time.Second)
	pool := NewMemoryPool()
	first, err := pool.UpsertFetched(context.Background(), "provider-a", FetchedBatch{
		ObservedAt: now, ConfiguredTTL: 30 * time.Second, MaxSize: 100,
		Proxies: []proxyDomain.Proxy{{
			Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080,
			State: proxyDomain.StateAvailable, LastCheckedAt: &checkedAt,
		}},
	})
	if err != nil || first.Inserted != 1 {
		t.Fatalf("UpsertFetched(first) = %+v, %v", first, err)
	}
	initial := pool.Snapshot(now)[0]

	second, err := pool.UpsertFetched(context.Background(), "provider-a", FetchedBatch{
		ObservedAt: now.Add(10 * time.Second), ConfiguredTTL: 30 * time.Second, MaxSize: 100,
		Proxies: []proxyDomain.Proxy{{
			Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080,
			State: proxyDomain.StateFetched,
		}},
	})
	if err != nil || second.Refreshed != 1 {
		t.Fatalf("UpsertFetched(refresh) = %+v, %v", second, err)
	}
	refreshed := pool.Snapshot(now.Add(10 * time.Second))[0]
	if refreshed.Proxy.ID != initial.Proxy.ID || refreshed.State != proxyDomain.StateAvailable ||
		refreshed.Proxy.LastCheckedAt == nil || !refreshed.Proxy.LastCheckedAt.Equal(checkedAt) ||
		!refreshed.Proxy.CreatedAt.Equal(initial.Proxy.CreatedAt) {
		t.Fatalf("refreshed entry lost runtime state: initial=%+v refreshed=%+v", initial, refreshed)
	}
}

func TestMemoryPoolDropsCandidatesWithoutUsableTTL(t *testing.T) {
	now := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC)
	expired := now.Add(-time.Second)
	tooShort := now.Add(2 * time.Second)
	pool := NewMemoryPool()

	result, err := pool.UpsertFetched(context.Background(), "provider-a", FetchedBatch{
		ObservedAt:             now,
		AllocationSafetyMargin: 3 * time.Second,
		MaxSize:                100,
		Proxies: []proxyDomain.Proxy{
			{Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8001},
			{Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.11", Port: 8002, ExpiresAt: &expired},
			{Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.12", Port: 8003, ExpiresAt: &tooShort},
		},
	})
	if err != nil {
		t.Fatalf("UpsertFetched(): %v", err)
	}
	if result.Dropped != 3 || result.Accepted != 0 || len(pool.Snapshot(now)) != 0 {
		t.Fatalf("result = %+v, entries=%d", result, len(pool.Snapshot(now)))
	}
}

func TestMemoryPoolRejectsInvalidBatchAtomically(t *testing.T) {
	now := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC)
	pool := NewMemoryPool()
	result, err := pool.UpsertFetched(context.Background(), "provider-a", FetchedBatch{
		ObservedAt:    now,
		ConfiguredTTL: time.Minute,
		MaxSize:       100,
		Proxies: []proxyDomain.Proxy{
			{Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8001},
			{Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.11", Port: 8002, SourceUpstream: "provider-b"},
		},
	})
	if !errors.Is(err, ErrInvalidBatch) {
		t.Fatalf("UpsertFetched() error = %v, want ErrInvalidBatch", err)
	}
	if result != (UpsertResult{}) || len(pool.Snapshot(now)) != 0 {
		t.Fatalf("failed batch retained state: result=%+v entries=%+v", result, pool.Snapshot(now))
	}
}

func TestMemoryPoolNeverExtractsProxyTwiceAndDoesNotWriteAuditRecords(t *testing.T) {
	now := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC)
	pool := poolWithOneProxy(t, now)

	var wait sync.WaitGroup
	results := make(chan string, 1000)
	for range 1000 {
		wait.Add(1)
		go func() {
			defer wait.Done()
			result, err := pool.Extract(context.Background(), extractionDomain.Command{
				Requested:         1,
				Fulfillment:       extractionDomain.Partial,
				Now:               now,
				MinRemainingTTL:   10 * time.Second,
				MaxHealthCheckAge: time.Minute,
			})
			if err != nil {
				t.Errorf("Extract(): %v", err)
				return
			}
			for _, item := range result.Items {
				results <- item.ID
			}
		}()
	}
	wait.Wait()
	close(results)

	count := 0
	for range results {
		count++
	}
	if count != 1 {
		t.Fatalf("proxy extracted %d times, want exactly once", count)
	}
	if got := pool.Snapshot(now); len(got) != 1 || got[0].State != proxyDomain.StateExtracted {
		t.Fatalf("post-extraction snapshot = %+v", got)
	}
}

func TestMemoryPoolIdempotencyIsBoundedByProxyExpiry(t *testing.T) {
	now := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC)
	pool := poolWithOneProxy(t, now)
	command := extractionDomain.Command{
		ClientID:          "client-a",
		IdempotencyKey:    "idem-a",
		IdempotencyTTL:    5 * time.Minute,
		Requested:         1,
		Fulfillment:       extractionDomain.Partial,
		Now:               now,
		MinRemainingTTL:   10 * time.Second,
		MaxHealthCheckAge: time.Minute,
	}

	first, err := pool.Extract(context.Background(), command)
	if err != nil || first.Returned != 1 {
		t.Fatalf("first Extract() = %+v, %v", first, err)
	}
	command.Now = now.Add(20 * time.Second)
	replayed, err := pool.Extract(context.Background(), command)
	if err != nil || replayed.Returned != 1 || replayed.Items[0].ID != first.Items[0].ID {
		t.Fatalf("replayed Extract() = %+v, %v", replayed, err)
	}

	if purged := pool.PurgeExpired(now.Add(31 * time.Second)); purged != 1 {
		t.Fatalf("PurgeExpired() = %d, want 1", purged)
	}
	command.Now = now.Add(31 * time.Second)
	afterExpiry, err := pool.Extract(context.Background(), command)
	if err != nil || afterExpiry.Returned != 0 {
		t.Fatalf("Extract(after expiry) = %+v, %v", afterExpiry, err)
	}
}

func TestMemoryPoolAllOrNothingDoesNotConsumePartialInventory(t *testing.T) {
	now := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC)
	pool := poolWithOneProxy(t, now)
	result, err := pool.Extract(context.Background(), extractionDomain.Command{
		Requested: 2, Fulfillment: extractionDomain.AllOrNothing, Now: now,
	})
	if !errors.Is(err, extractionDomain.ErrInsufficientProxies) || result.Returned != 0 {
		t.Fatalf("all-or-nothing Extract() = %+v, %v", result, err)
	}
	partial, err := pool.Extract(context.Background(), extractionDomain.Command{
		Requested: 1, Fulfillment: extractionDomain.Partial, Now: now,
	})
	if err != nil || partial.Returned != 1 {
		t.Fatalf("partial Extract() after rollback = %+v, %v", partial, err)
	}
}

func TestMemoryPoolAssignRecoversExpiredLeaseWithoutSeparateSweep(t *testing.T) {
	now := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC)
	pool := poolWithOneProxy(t, now)
	proxyID := pool.Snapshot(now)[0].Proxy.ID
	first, err := pool.Assign(context.Background(), now, proxyID, "worker-a", 5*time.Second)
	if err != nil {
		t.Fatalf("Assign(first): %v", err)
	}
	second, err := pool.Assign(context.Background(), now.Add(5*time.Second), proxyID, "worker-b", 5*time.Second)
	if err != nil {
		t.Fatalf("Assign(after lease expiry): %v", err)
	}
	if second.Epoch <= first.Epoch || second.WorkerID != "worker-b" {
		t.Fatalf("second assignment = %+v, first=%+v", second, first)
	}
}

func TestMemoryPoolMakesOwnershipAndExtractionMutuallyExclusive(t *testing.T) {
	now := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC)
	for iteration := range 100 {
		pool := poolWithOneProxy(t, now)
		start := make(chan struct{})
		assigned := make(chan bool, 1)
		extracted := make(chan bool, 1)
		go func() {
			<-start
			_, err := pool.Assign(context.Background(), now, pool.Snapshot(now)[0].Proxy.ID, "worker-a", time.Minute)
			if err != nil && !errors.Is(err, ownershipDomain.ErrOwnershipUnavailable) {
				t.Errorf("Assign(): %v", err)
			}
			assigned <- err == nil
		}()
		go func() {
			<-start
			result, err := pool.Extract(context.Background(), extractionDomain.Command{
				ClientID: "client-a", Requested: 1, Fulfillment: extractionDomain.Partial, Now: now,
			})
			if err != nil {
				t.Errorf("Extract(): %v", err)
			}
			extracted <- result.Returned == 1
		}()
		close(start)

		wins := 0
		if <-assigned {
			wins++
		}
		if <-extracted {
			wins++
		}
		if wins != 1 {
			t.Fatalf("iteration %d winners = %d, want 1", iteration, wins)
		}
	}
}

func TestMemoryPoolOwnershipMethodsPropagateCanceledContext(t *testing.T) {
	now := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC)
	pool := poolWithOneProxy(t, now)
	proxyID := pool.Snapshot(now)[0].Proxy.ID
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	if _, err := pool.Assign(ctx, now, proxyID, "worker-a", time.Minute); !errors.Is(err, context.Canceled) {
		t.Fatalf("Assign() error = %v, want context.Canceled", err)
	}
	if _, err := pool.Renew(ctx, now, proxyID, "worker-a", 1, time.Minute); !errors.Is(err, context.Canceled) {
		t.Fatalf("Renew() error = %v, want context.Canceled", err)
	}
	if _, err := pool.BeginDrain(ctx, proxyID, "worker-a", 1); !errors.Is(err, context.Canceled) {
		t.Fatalf("BeginDrain() error = %v, want context.Canceled", err)
	}
	if err := pool.AcknowledgeDrain(ctx, proxyID, "worker-a", 1, 0, 0); !errors.Is(err, context.Canceled) {
		t.Fatalf("AcknowledgeDrain() error = %v, want context.Canceled", err)
	}
	if _, _, err := pool.Get(ctx, proxyID); !errors.Is(err, context.Canceled) {
		t.Fatalf("Get() error = %v, want context.Canceled", err)
	}
	if _, err := pool.Expire(ctx, now, 1); !errors.Is(err, context.Canceled) {
		t.Fatalf("Expire() error = %v, want context.Canceled", err)
	}
}

func TestMemoryPoolOwnershipMethodRechecksContextAfterLock(t *testing.T) {
	now := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC)
	pool := poolWithOneProxy(t, now)
	proxyID := pool.Snapshot(now)[0].Proxy.ID
	base, cancel := context.WithCancel(context.Background())
	ctx := &firstCheckContext{Context: base, checked: make(chan struct{})}
	result := make(chan error, 1)

	pool.mu.Lock()
	go func() {
		_, err := pool.Assign(ctx, now, proxyID, "worker-a", time.Minute)
		result <- err
	}()
	<-ctx.checked
	cancel()
	pool.mu.Unlock()

	if err := <-result; !errors.Is(err, context.Canceled) {
		t.Fatalf("Assign() error = %v, want context.Canceled", err)
	}
	if _, ok, err := pool.Get(context.Background(), proxyID); err != nil || ok {
		t.Fatal("canceled Assign() mutated ownership")
	}
}

func TestMemoryPoolExpireUsesStableProxyIDOrderAndLimit(t *testing.T) {
	now := time.Date(2026, 7, 29, 10, 0, 0, 0, time.UTC)
	pool := NewMemoryPool()
	proxies := []proxyDomain.Proxy{
		{ID: "proxy-c", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.3", Port: 8003, State: proxyDomain.StateAvailable},
		{ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.1", Port: 8001, State: proxyDomain.StateAvailable},
		{ID: "proxy-b", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.2", Port: 8002, State: proxyDomain.StateAvailable},
	}
	if _, err := pool.UpsertFetched(context.Background(), "provider-a", FetchedBatch{
		ObservedAt: now, ConfiguredTTL: 10 * time.Minute, MaxSize: 100, Proxies: proxies,
	}); err != nil {
		t.Fatalf("UpsertFetched(): %v", err)
	}
	for _, proxyID := range []string{"proxy-c", "proxy-a", "proxy-b"} {
		if _, err := pool.Assign(context.Background(), now, proxyID, "worker-a", time.Minute); err != nil {
			t.Fatalf("Assign(%s): %v", proxyID, err)
		}
	}

	expired, err := pool.Expire(context.Background(), now.Add(time.Minute), 2)
	if err != nil {
		t.Fatalf("Expire(): %v", err)
	}
	if len(expired) != 2 || expired[0].ProxyID != "proxy-a" || expired[1].ProxyID != "proxy-b" {
		t.Fatalf("Expire() = %+v, want proxy-a then proxy-b", expired)
	}
	if _, ok, err := pool.Get(context.Background(), "proxy-c"); err != nil || !ok {
		t.Fatalf("Get(proxy-c) = ok %v, error %v; want remaining assignment", ok, err)
	}
}

type firstCheckContext struct {
	context.Context
	checked chan struct{}
	once    sync.Once
}

func (c *firstCheckContext) Err() error {
	err := c.Context.Err()
	c.once.Do(func() { close(c.checked) })
	return err
}

func poolWithOneProxy(t *testing.T, now time.Time) *MemoryPool {
	t.Helper()
	pool := NewMemoryPool()
	checkedAt := now
	result, err := pool.UpsertFetched(context.Background(), "provider-a", FetchedBatch{
		ObservedAt:             now,
		ConfiguredTTL:          30 * time.Second,
		AllocationSafetyMargin: 3 * time.Second,
		MaxSize:                100,
		Proxies: []proxyDomain.Proxy{{
			Scheme:        proxyDomain.SchemeHTTP,
			Host:          "192.0.2.10",
			Port:          8080,
			State:         proxyDomain.StateAvailable,
			LastCheckedAt: &checkedAt,
		}},
	})
	if err != nil || result.Inserted != 1 {
		t.Fatalf("UpsertFetched() = %+v, %v", result, err)
	}
	return pool
}