package pool

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
	"testing"
	"time"

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

func TestOwnershipManagerPreventsDualAssignment(t *testing.T) {
	manager := newTestOwnershipManager(t, "proxy-1")
	now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
	var succeeded atomic.Int64
	var wg sync.WaitGroup

	for index := range 100 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			_, err := manager.Assign(context.Background(), now, "proxy-1", fmt.Sprintf("worker-%d", index), time.Minute)
			if err == nil {
				succeeded.Add(1)
				return
			}
			if !errors.Is(err, ErrAlreadyOwned) {
				t.Errorf("Assign(): %v", err)
			}
		}()
	}
	wg.Wait()

	if got := succeeded.Load(); got != 1 {
		t.Fatalf("successful assignments = %d, want 1", got)
	}
}

func TestOwnershipManagerRenewsOnlyCurrentAssignment(t *testing.T) {
	manager := newTestOwnershipManager(t, "proxy-1")
	now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
	ctx := context.Background()
	assigned, err := manager.Assign(ctx, now, "proxy-1", "worker-1", time.Minute)
	if err != nil {
		t.Fatalf("Assign(): %v", err)
	}
	renewed, err := manager.Renew(ctx, now.Add(30*time.Second), "proxy-1", "worker-1", assigned.Epoch, time.Minute)
	if err != nil {
		t.Fatalf("Renew(): %v", err)
	}
	if renewed.Version != assigned.Version+1 || !renewed.ExpiresAt.Equal(now.Add(90*time.Second)) {
		t.Fatalf("renewed assignment = %+v", renewed)
	}
	if _, err := manager.Renew(ctx, now, "proxy-1", "worker-2", assigned.Epoch, time.Minute); !errors.Is(err, ErrStaleAssignment) {
		t.Fatalf("Renew(stale) error = %v, want ErrStaleAssignment", err)
	}
	if expired, err := manager.Expire(ctx, now.Add(time.Minute), 32); err != nil || len(expired) != 0 {
		t.Fatalf("renewed assignment expired at old deadline: %+v", expired)
	}
}

func TestSharedRepositoryMakesOwnershipAndExtractionMutuallyExclusive(t *testing.T) {
	for iteration := range 100 {
		now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
		store := newTestActivityPool(t, now, "proxy-1")
		manager, err := NewOwnershipManager(store)
		if err != nil {
			t.Fatalf("iteration %d NewOwnershipManager(): %v", iteration, err)
		}

		start := make(chan struct{})
		assigned := make(chan bool, 1)
		extracted := make(chan bool, 1)
		go func() {
			<-start
			_, assignErr := manager.Assign(context.Background(), now, "proxy-1", "worker-1", time.Minute)
			if assignErr != nil && !errors.Is(assignErr, ErrOwnershipUnavailable) {
				t.Errorf("iteration %d Assign(): %v", iteration, assignErr)
			}
			assigned <- assignErr == nil
		}()
		go func() {
			<-start
			result, extractErr := store.Extract(context.Background(), extractionDomain.Command{
				ClientID: "client-1", Requested: 1, Fulfillment: extractionDomain.Partial, Now: now,
			})
			if extractErr != nil {
				t.Errorf("iteration %d Extract(): %v", iteration, extractErr)
			}
			extracted <- result.Returned == 1
		}()
		close(start)

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

func TestOwnershipManagerRequiresDrainAckAtZeroRuntime(t *testing.T) {
	manager := newTestOwnershipManager(t, "proxy-1")
	now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
	ctx := context.Background()
	assignment, err := manager.Assign(ctx, now, "proxy-1", "worker-1", time.Minute)
	if err != nil {
		t.Fatalf("Assign(): %v", err)
	}
	draining, err := manager.BeginDrain(ctx, "proxy-1", "worker-1", assignment.Epoch)
	if err != nil {
		t.Fatalf("BeginDrain(): %v", err)
	}
	if !draining.Draining || draining.Version != assignment.Version+1 {
		t.Fatalf("draining assignment = %+v", draining)
	}

	if err := manager.AcknowledgeDrain(ctx, "proxy-1", "worker-1", assignment.Epoch, 1, 0); !errors.Is(err, ErrDrainNotReady) {
		t.Fatalf("AcknowledgeDrain(active) error = %v, want ErrDrainNotReady", err)
	}
	if err := manager.AcknowledgeDrain(ctx, "proxy-1", "worker-1", assignment.Epoch, 0, 0); err != nil {
		t.Fatalf("AcknowledgeDrain(zero): %v", err)
	}
	if _, ok, err := manager.Get(ctx, "proxy-1"); err != nil || ok {
		t.Fatal("assignment still exists after drain acknowledgement")
	}
}

func TestOwnershipManagerExpiresCrashedWorkerAssignment(t *testing.T) {
	manager := newTestOwnershipManager(t, "proxy-1")
	now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
	ctx := context.Background()
	first, err := manager.Assign(ctx, now, "proxy-1", "worker-1", time.Minute)
	if err != nil {
		t.Fatalf("Assign(first): %v", err)
	}
	if expired, err := manager.Expire(ctx, now.Add(59*time.Second), 32); err != nil || len(expired) != 0 {
		t.Fatalf("expired early: %+v", expired)
	}
	if expired, err := manager.Expire(ctx, now.Add(time.Minute), 32); err != nil || len(expired) != 1 || expired[0].ProxyID != "proxy-1" {
		t.Fatalf("Expire() = %+v, want proxy-1", expired)
	}

	second, err := manager.Assign(ctx, now.Add(time.Minute), "proxy-1", "worker-2", time.Minute)
	if err != nil {
		t.Fatalf("Assign(second): %v", err)
	}
	if second.Epoch <= first.Epoch {
		t.Fatalf("second epoch = %d, want greater than %d", second.Epoch, first.Epoch)
	}
}

func TestOwnershipManagerPropagatesCanceledContextWithoutCallingRepository(t *testing.T) {
	repository := &ownershipRepositoryStub{}
	manager, err := NewOwnershipManager(repository)
	if err != nil {
		t.Fatalf("NewOwnershipManager(): %v", err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	_, err = manager.Assign(ctx, time.Now(), "proxy-1", "worker-1", time.Minute)
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("Assign() error = %v, want context.Canceled", err)
	}
	if repository.assignCalled {
		t.Fatal("repository Assign called with canceled context")
	}
}

func TestOwnershipManagerPassesContextAndRepositoryError(t *testing.T) {
	storageErr := errors.New("redis unavailable")
	ctx := context.WithValue(context.Background(), ownershipContextKey{}, "request-1")
	repository := &ownershipRepositoryStub{
		assign: func(got context.Context, _ time.Time, _, _ string, _ time.Duration) (Assignment, error) {
			if got != ctx {
				t.Fatal("Assign() did not pass the original context")
			}
			return Assignment{}, storageErr
		},
	}
	manager, err := NewOwnershipManager(repository)
	if err != nil {
		t.Fatalf("NewOwnershipManager(): %v", err)
	}

	_, err = manager.Assign(ctx, time.Now(), "proxy-1", "worker-1", time.Minute)
	if !errors.Is(err, storageErr) {
		t.Fatalf("Assign() error = %v, want repository error", err)
	}
}

func TestOwnershipManagerPassesExpireLimit(t *testing.T) {
	ctx := context.Background()
	now := time.Date(2026, 7, 29, 12, 0, 0, 0, time.UTC)
	want := []Assignment{{ProxyID: "proxy-1"}}
	repository := &ownershipRepositoryStub{
		expire: func(got context.Context, gotNow time.Time, gotLimit int) ([]Assignment, error) {
			if got != ctx || !gotNow.Equal(now) || gotLimit != 32 {
				t.Fatalf("Expire() arguments = (%v, %v, %d)", got, gotNow, gotLimit)
			}
			return want, nil
		},
	}
	manager, err := NewOwnershipManager(repository)
	if err != nil {
		t.Fatalf("NewOwnershipManager(): %v", err)
	}

	got, err := manager.Expire(ctx, now, 32)
	if err != nil || len(got) != 1 || got[0].ProxyID != want[0].ProxyID {
		t.Fatalf("Expire() = %+v, %v", got, err)
	}
}

func TestOwnershipManagerRejectsInvalidContextAndExpireLimit(t *testing.T) {
	manager, err := NewOwnershipManager(&ownershipRepositoryStub{})
	if err != nil {
		t.Fatalf("NewOwnershipManager(): %v", err)
	}
	now := time.Date(2026, 7, 29, 12, 0, 0, 0, time.UTC)

	if _, err := manager.Assign(nil, now, "proxy-1", "worker-1", time.Minute); !errors.Is(err, ErrInvalidOwnership) {
		t.Fatalf("Assign(nil context) error = %v, want ErrInvalidOwnership", err)
	}
	if _, err := manager.Expire(context.Background(), now, 0); !errors.Is(err, ErrInvalidOwnership) {
		t.Fatalf("Expire(zero limit) error = %v, want ErrInvalidOwnership", err)
	}
}

type ownershipContextKey struct{}

type ownershipRepositoryStub struct {
	assign       func(context.Context, time.Time, string, string, time.Duration) (Assignment, error)
	expire       func(context.Context, time.Time, int) ([]Assignment, error)
	assignCalled bool
}

func (r *ownershipRepositoryStub) Assign(ctx context.Context, now time.Time, proxyID, workerID string, ttl time.Duration) (Assignment, error) {
	r.assignCalled = true
	if r.assign != nil {
		return r.assign(ctx, now, proxyID, workerID, ttl)
	}
	return Assignment{}, nil
}

func (r *ownershipRepositoryStub) Renew(context.Context, time.Time, string, string, uint64, time.Duration) (Assignment, error) {
	return Assignment{}, nil
}

func (r *ownershipRepositoryStub) BeginDrain(context.Context, string, string, uint64) (Assignment, error) {
	return Assignment{}, nil
}

func (r *ownershipRepositoryStub) AcknowledgeDrain(context.Context, string, string, uint64, int64, int64) error {
	return nil
}

func (r *ownershipRepositoryStub) Get(context.Context, string) (Assignment, bool, error) {
	return Assignment{}, false, nil
}

func (r *ownershipRepositoryStub) Expire(ctx context.Context, now time.Time, limit int) ([]Assignment, error) {
	if r.expire != nil {
		return r.expire(ctx, now, limit)
	}
	return nil, nil
}

var _ ownershipDomain.Repository = (*ownershipRepositoryStub)(nil)

func newTestOwnershipManager(t *testing.T, proxyIDs ...string) *OwnershipManager {
	t.Helper()
	now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
	manager, err := NewOwnershipManager(newTestActivityPool(t, now, proxyIDs...))
	if err != nil {
		t.Fatalf("NewOwnershipManager(): %v", err)
	}
	return manager
}

func newTestActivityPool(t *testing.T, now time.Time, proxyIDs ...string) *activitypool.MemoryPool {
	t.Helper()
	proxies := make([]proxyDomain.Proxy, 0, len(proxyIDs))
	for index, proxyID := range proxyIDs {
		proxies = append(proxies, proxyDomain.Proxy{
			ID: proxyID, Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.1",
			Port: uint16(8000 + index), State: proxyDomain.StateAvailable,
		})
	}
	store := activitypool.NewMemoryPool()
	result, err := store.UpsertFetched(context.Background(), "provider-a", activitypool.FetchedBatch{
		ObservedAt: now, ConfiguredTTL: 10 * time.Minute, Proxies: proxies,
	})
	if err != nil || result.Inserted != len(proxyIDs) {
		t.Fatalf("UpsertFetched() = %+v, %v", result, err)
	}
	return store
}