package dispatch

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

	proxyDomain "proxy-pool/internal/domain/proxy"
	"proxy-pool/internal/domain/routing"
	"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 TestAcquireRejectsSnapshotAfterOverallValidityDeadline(t *testing.T) {
	now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
	expiresAt := now.Add(time.Hour)
	store := snapshot.NewStore("cluster-a", "worker-a")
	envelope := snapshot.Envelope{
		ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true,
		ValidUntil: now.Add(time.Second),
		Proxies: []proxyDomain.Proxy{{
			ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, State: proxyDomain.StateAvailable,
			MaxConcurrency: 1, ExpiresAt: &expiresAt,
		}},
	}
	envelope.Checksum = snapshot.Checksum(envelope.Proxies)
	if err := store.Apply(envelope); err != nil {
		t.Fatalf("Apply(): %v", err)
	}
	if _, err := New(store).Acquire(Request{Now: now.Add(time.Second), Scheme: proxyDomain.SchemeHTTP}); !errors.Is(err, ErrNoCandidate) {
		t.Fatalf("Acquire(at validity deadline) error = %v, want ErrNoCandidate", 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 TestAcquireUsesSequentialCurrentUpstream(t *testing.T) {
	store := routingStrategyStore(t)
	dispatcher := New(store)
	lease, err := dispatcher.Acquire(Request{
		Now: time.Now(), RoutingName: "sequential", Upstreams: []string{"upstream-a", "upstream-b"},
		Strategy: routing.Strategy{Type: routing.StrategySequential, CurrentUpstream: "upstream-b"},
	})
	if err != nil {
		t.Fatalf("Acquire(): %v", err)
	}
	if lease.Proxy.SourceUpstream != "upstream-b" {
		t.Fatalf("selected upstream = %q, want upstream-b", lease.Proxy.SourceUpstream)
	}
	if err := lease.Cancel(); err != nil {
		t.Fatalf("Cancel(): %v", err)
	}
}

func TestAcquireRoundRobinCyclesRoutingUpstreams(t *testing.T) {
	store := routingStrategyStore(t)
	dispatcher := New(store)
	request := Request{
		Now: time.Now(), RoutingName: "round-robin", Upstreams: []string{"upstream-a", "upstream-b"},
		Strategy: routing.Strategy{Type: routing.StrategyRoundRobin},
	}
	for index, want := range []string{"upstream-a", "upstream-b", "upstream-a"} {
		lease, err := dispatcher.Acquire(request)
		if err != nil {
			t.Fatalf("Acquire(%d): %v", index, err)
		}
		if lease.Proxy.SourceUpstream != want {
			t.Fatalf("Acquire(%d) upstream = %q, want %q", index, lease.Proxy.SourceUpstream, want)
		}
		if err := lease.Cancel(); err != nil {
			t.Fatalf("Cancel(%d): %v", index, err)
		}
	}
}

func TestAcquireResetsRoundRobinStateWhenSnapshotVersionChanges(t *testing.T) {
	store := routingStrategyStore(t)
	dispatcher := New(store)
	request := Request{
		Now: time.Now(), RoutingName: "round-robin", Upstreams: []string{"upstream-a", "upstream-b"},
		Strategy: routing.Strategy{Type: routing.StrategyRoundRobin},
	}
	first, err := dispatcher.Acquire(request)
	if err != nil {
		t.Fatalf("Acquire(first): %v", err)
	}
	if first.Proxy.SourceUpstream != "upstream-a" {
		t.Fatalf("first upstream = %q, want upstream-a", first.Proxy.SourceUpstream)
	}
	if err := first.Cancel(); err != nil {
		t.Fatalf("Cancel(first): %v", err)
	}

	next := snapshot.Envelope{
		ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 2, Full: true,
		Proxies: []proxyDomain.Proxy{
			{ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "upstream-a", State: proxyDomain.StateAvailable, MaxConcurrency: 1},
			{ID: "proxy-b", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "upstream-b", State: proxyDomain.StateAvailable, MaxConcurrency: 1},
		},
	}
	next.Checksum = snapshot.Checksum(next.Proxies)
	if err := store.Apply(next); err != nil {
		t.Fatalf("Apply(next): %v", err)
	}
	second, err := dispatcher.Acquire(request)
	if err != nil {
		t.Fatalf("Acquire(second): %v", err)
	}
	if second.Proxy.SourceUpstream != "upstream-a" {
		t.Fatalf("second upstream = %q, want reset upstream-a", second.Proxy.SourceUpstream)
	}
	if err := second.Cancel(); err != nil {
		t.Fatalf("Cancel(second): %v", err)
	}
}

func TestAcquireFallsBackWhenSelectedRoutingUpstreamIsFull(t *testing.T) {
	store := routingStrategyStore(t)
	var held *proxyDomain.Reservation
	for _, entry := range store.Current().Entries {
		if entry.Proxy.SourceUpstream != "upstream-a" {
			continue
		}
		reservation, ok := entry.Runtime.Reserve()
		if !ok {
			t.Fatal("Reserve(upstream-a) = false, want true")
		}
		if err := reservation.Commit(); err != nil {
			t.Fatalf("Commit(upstream-a): %v", err)
		}
		held = reservation
		break
	}
	if held == nil {
		t.Fatal("upstream-a reservation was not created")
	}
	defer func() {
		if err := held.Release(); err != nil {
			t.Errorf("Release(upstream-a): %v", err)
		}
	}()

	lease, err := New(store).Acquire(Request{
		Now: time.Now(), RoutingName: "round-robin", Upstreams: []string{"upstream-a", "upstream-b"},
		Strategy: routing.Strategy{Type: routing.StrategyRoundRobin},
	})
	if err != nil {
		t.Fatalf("Acquire(): %v", err)
	}
	if lease.Proxy.SourceUpstream != "upstream-b" {
		t.Fatalf("selected upstream = %q, want upstream-b", lease.Proxy.SourceUpstream)
	}
	if err := lease.Cancel(); err != nil {
		t.Fatalf("Cancel(): %v", err)
	}
}

func TestAcquireUsesWeightedRoutingStrategy(t *testing.T) {
	store := routingStrategyStore(t)
	dispatcher := New(store, &fixedRoutingRandom{values: []int{0, 2, 3}})
	request := Request{
		Now: time.Now(), RoutingName: "weighted", Upstreams: []string{"upstream-a", "upstream-b"},
		Strategy: routing.Strategy{Type: routing.StrategyWeighted, Weights: map[string]uint32{"upstream-a": 3, "upstream-b": 1}},
	}
	for index, want := range []string{"upstream-a", "upstream-a", "upstream-b"} {
		lease, err := dispatcher.Acquire(request)
		if err != nil {
			t.Fatalf("Acquire(%d): %v", index, err)
		}
		if lease.Proxy.SourceUpstream != want {
			t.Fatalf("Acquire(%d) upstream = %q, want %q", index, lease.Proxy.SourceUpstream, want)
		}
		if err := lease.Cancel(); err != nil {
			t.Fatalf("Cancel(%d): %v", index, err)
		}
	}
}

func TestAcquireUsesLeastConnectionsAndSkipsFullUpstream(t *testing.T) {
	store := routingStrategyStore(t)
	view := store.Current()
	for _, entry := range view.Entries {
		if entry.Proxy.SourceUpstream != "upstream-a" {
			continue
		}
		reservation, ok := entry.Runtime.Reserve()
		if !ok {
			t.Fatal("Reserve(upstream-a) = false, want true")
		}
		if err := reservation.Commit(); err != nil {
			t.Fatalf("Commit(upstream-a): %v", err)
		}
		defer func() {
			if err := reservation.Release(); err != nil {
				t.Errorf("Release(upstream-a): %v", err)
			}
		}()
		break
	}

	lease, err := New(store).Acquire(Request{
		Now: time.Now(), RoutingName: "least", Upstreams: []string{"upstream-a", "upstream-b"},
		Strategy: routing.Strategy{Type: routing.StrategyLeastConnections},
	})
	if err != nil {
		t.Fatalf("Acquire(): %v", err)
	}
	if lease.Proxy.SourceUpstream != "upstream-b" {
		t.Fatalf("selected upstream = %q, want upstream-b", lease.Proxy.SourceUpstream)
	}
	if err := lease.Cancel(); err != nil {
		t.Fatalf("Cancel(): %v", err)
	}
}

func TestAcquireWaitsForReleasedLocalCapacity(t *testing.T) {
	store := routingStrategyStore(t)
	dispatcher := New(store)
	held, err := dispatcher.Acquire(Request{Now: time.Now(), Upstreams: []string{"upstream-a"}})
	if err != nil {
		t.Fatalf("Acquire(held): %v", err)
	}
	if err := held.Commit(); err != nil {
		t.Fatalf("Commit(held): %v", err)
	}
	released := make(chan struct{})
	go func() {
		time.Sleep(10 * time.Millisecond)
		_ = held.Release()
		close(released)
	}()

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	lease, err := dispatcher.AcquireWait(ctx, Request{Upstreams: []string{"upstream-a"}}, 100*time.Millisecond)
	if err != nil {
		t.Fatalf("AcquireWait(): %v", err)
	}
	if err := lease.Cancel(); err != nil {
		t.Fatalf("Cancel(): %v", err)
	}
	<-released
}

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: 5},
			{epoch: 2, version: 6},
			{epoch: 2, version: 7},
			{epoch: 3, version: 8},
			{epoch: 3, version: 9},
		}
		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 BenchmarkAcquire100kRoutingRoundRobin(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), RoutingName: "benchmark-round-robin",
		Upstreams: []string{"upstream-a", "upstream-b"}, Strategy: routing.Strategy{Type: routing.StrategyRoundRobin},
		Scheme: proxyDomain.SchemeHTTP, RequiredTags: map[string]string{"region": "cn-east", "tier": "gold"},
		SafetyMargin: 5 * time.Second,
	}

	b.ReportAllocs()
	b.ResetTimer()
	b.RunParallel(func(parallel *testing.PB) {
		for parallel.Next() {
			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 routingStrategyStore(t *testing.T) *snapshot.Store {
	t.Helper()
	store := snapshot.NewStore("cluster-a", "worker-a")
	envelope := snapshot.Envelope{
		ClusterID: "cluster-a", WorkerID: "worker-a", Epoch: 1, Version: 1, Full: true,
		Proxies: []proxyDomain.Proxy{
			{ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "upstream-a", State: proxyDomain.StateAvailable, MaxConcurrency: 1},
			{ID: "proxy-b", Scheme: proxyDomain.SchemeHTTP, SourceUpstream: "upstream-b", State: proxyDomain.StateAvailable, MaxConcurrency: 1},
		},
	}
	envelope.Checksum = snapshot.Checksum(envelope.Proxies)
	if err := store.Apply(envelope); err != nil {
		t.Fatalf("Apply(): %v", err)
	}
	return store
}

type fixedRoutingRandom struct {
	values []int
	index  int
}

func (source *fixedRoutingRandom) Intn(limit int) int {
	value := source.values[source.index]
	source.index++
	return value % limit
}

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"
}