package dispatch

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

	proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy"
	"github.com/proxy-pool/proxy-pool/internal/gateway/snapshot"
)

func TestAcquireFiltersAndReservesLocalCapacity(t *testing.T) {
	now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
	expiresSoon := now.Add(5 * time.Second)
	expiresLater := now.Add(time.Minute)
	store := snapshot.NewStore("cluster-a", "worker-a")
	proxies := []proxyDomain.Proxy{
		{ID: "wrong-upstream", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "b", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresLater},
		{ID: "expiring", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "a", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresSoon},
		{ID: "selected", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "a", State: proxyDomain.StateAvailable, MaxConcurrency: 1, ExpiresAt: &expiresLater, Tags: map[string]string{"region": "cn-east"}},
	}
	envelope := snapshot.Envelope{ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true, Proxies: proxies}
	envelope.Checksum = snapshot.Checksum(proxies)
	if err := store.Apply(envelope); err != nil {
		t.Fatalf("Apply(): %v", err)
	}

	dispatcher := New(store)
	lease, err := dispatcher.Acquire(Request{
		Now:          now,
		Scheme:       proxyDomain.SchemeHTTP,
		Upstreams:    []string{"a"},
		RequiredTags: map[string]string{"region": "cn-east"},
		SafetyMargin: 10 * time.Second,
	})
	if err != nil {
		t.Fatalf("Acquire(): %v", err)
	}
	if lease.Proxy.ID != "selected" {
		t.Fatalf("selected proxy = %q, want selected", lease.Proxy.ID)
	}
	if err := lease.Commit(); err != nil {
		t.Fatalf("Commit(): %v", err)
	}
	if err := lease.Release(); err != nil {
		t.Fatalf("Release(): %v", err)
	}
}

func TestAcquireNeverOversubscribesSnapshotProxy(t *testing.T) {
	store := snapshot.NewStore("cluster-a", "worker-a")
	proxies := []proxyDomain.Proxy{{ID: "p1", Scheme: proxyDomain.SchemeHTTP, State: proxyDomain.StateAvailable, MaxConcurrency: 8}}
	envelope := snapshot.Envelope{ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true, Proxies: proxies}
	envelope.Checksum = snapshot.Checksum(proxies)
	if err := store.Apply(envelope); err != nil {
		t.Fatalf("Apply(): %v", err)
	}
	dispatcher := New(store)

	var wg sync.WaitGroup
	leases := make(chan *Lease, 1000)
	for range 1000 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			lease, err := dispatcher.Acquire(Request{Now: time.Now(), Scheme: proxyDomain.SchemeHTTP})
			if err == nil {
				leases <- lease
				return
			}
			if !errors.Is(err, ErrNoCandidate) {
				t.Errorf("Acquire(): %v", err)
			}
		}()
	}
	wg.Wait()
	close(leases)

	count := 0
	for lease := range leases {
		count++
		if err := lease.Cancel(); err != nil {
			t.Errorf("Cancel(): %v", err)
		}
	}
	if count != 8 {
		t.Fatalf("reserved = %d, want 8", count)
	}
}

func TestAcquireSurvivesConcurrentSnapshotApply(t *testing.T) {
	now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
	store := snapshot.NewStore("cluster-a", "worker-a")
	base := makeTestProxies(256)
	envelope := snapshot.Envelope{
		ClusterID: "cluster-a",
		WorkerID:  "worker-a",
		Epoch:     1,
		Version:   1,
		Full:      true,
		Proxies:   base,
	}
	envelope.Checksum = snapshot.Checksum(envelope.Proxies)
	if err := store.Apply(envelope); err != nil {
		t.Fatalf("Apply(base): %v", err)
	}

	dispatcher := New(store)
	var successes atomic.Int64
	var applyErr atomic.Value
	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		defer wg.Done()
		steps := []struct {
			epoch   uint64
			version uint64
		}{
			{epoch: 1, version: 2},
			{epoch: 1, version: 3},
			{epoch: 1, version: 4},
			{epoch: 2, version: 1},
			{epoch: 2, version: 2},
			{epoch: 2, version: 3},
			{epoch: 3, version: 1},
			{epoch: 3, version: 2},
		}
		for _, step := range steps {
			next := snapshot.Envelope{
				ClusterID: "cluster-a",
				WorkerID:  "worker-a",
				Epoch:     step.epoch,
				Version:   step.version,
				Full:      true,
				Proxies:   makeTestProxies(256),
			}
			next.Checksum = snapshot.Checksum(next.Proxies)
			if err := store.Apply(next); err != nil {
				applyErr.Store(err)
				return
			}
		}
	}()

	for range 8 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for range 512 {
				lease, err := dispatcher.Acquire(Request{
					Now:       now,
					Scheme:    proxyDomain.SchemeHTTP,
					Upstreams: []string{"upstream-a"},
					RequiredTags: map[string]string{
						"region": "cn-east",
					},
				})
				if err != nil {
					if !errors.Is(err, ErrNoCandidate) {
						t.Errorf("Acquire(): %v", err)
						return
					}
					continue
				}
				successes.Add(1)
				if err := lease.Cancel(); err != nil {
					t.Errorf("Cancel(): %v", err)
					return
				}
			}
		}()
	}
	wg.Wait()

	if err, _ := applyErr.Load().(error); err != nil {
		t.Fatalf("Apply(): %v", err)
	}
	if successes.Load() == 0 {
		t.Fatal("Acquire() never succeeded during concurrent Apply")
	}
}

func BenchmarkAcquire100kIndexed(b *testing.B) {
	store := snapshot.NewStore("cluster-a", "worker-a")
	proxies := makeTestProxies(100_000)
	envelope := snapshot.Envelope{
		ClusterID: "cluster-a",
		WorkerID:  "worker-a",
		Epoch:     1,
		Version:   1,
		Full:      true,
		Proxies:   proxies,
	}
	envelope.Checksum = snapshot.Checksum(envelope.Proxies)
	if err := store.Apply(envelope); err != nil {
		b.Fatalf("Apply(): %v", err)
	}

	dispatcher := New(store)
	request := Request{
		Now:       time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC),
		Scheme:    proxyDomain.SchemeHTTP,
		Upstreams: []string{"upstream-a"},
		RequiredTags: map[string]string{
			"region": "cn-east",
			"tier":   "gold",
		},
		SafetyMargin: 5 * time.Second,
	}

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		lease, err := dispatcher.Acquire(request)
		if err != nil {
			b.Fatalf("Acquire(): %v", err)
		}
		if err := lease.Cancel(); err != nil {
			b.Fatalf("Cancel(): %v", err)
		}
	}
}

func makeTestProxies(count int) []proxyDomain.Proxy {
	now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
	expiresAt := now.Add(10 * time.Minute)
	proxies := make([]proxyDomain.Proxy, 0, count)
	for index := range count {
		proxies = append(proxies, proxyDomain.Proxy{
			ID:             fmt.Sprintf("proxy-%06d", index),
			Scheme:         schemeForIndex(index),
			Host:           fmt.Sprintf("10.0.%d.%d", index/256, index%256),
			Port:           uint16(20000 + index%1000),
			SourceUpstream: upstreamForIndex(index),
			State:          proxyDomain.StateAvailable,
			MaxConcurrency: 8,
			ExpiresAt:      &expiresAt,
			Tags: map[string]string{
				"region": regionForIndex(index),
				"tier":   tierForIndex(index),
			},
		})
	}
	return proxies
}

func schemeForIndex(index int) proxyDomain.Scheme {
	switch index % 3 {
	case 0:
		return proxyDomain.SchemeHTTP
	case 1:
		return proxyDomain.SchemeHTTPS
	default:
		return proxyDomain.SchemeSOCKS5
	}
}

func upstreamForIndex(index int) string {
	if index%2 == 0 {
		return "upstream-a"
	}
	return "upstream-b"
}

func regionForIndex(index int) string {
	switch index % 4 {
	case 0:
		return "cn-east"
	case 1:
		return "us-west"
	case 2:
		return "eu-central"
	default:
		return "ap-south"
	}
}

func tierForIndex(index int) string {
	if index%5 == 0 {
		return "gold"
	}
	return "silver"
}