package worker

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

	"google.golang.org/protobuf/types/known/timestamppb"
	controlplanev1 "proxy-pool/gen/controlplane/v1"
	"proxy-pool/internal/controlplane/snapshotwire"
	ownershipDomain "proxy-pool/internal/domain/ownership"
	proxyDomain "proxy-pool/internal/domain/proxy"
	platformCredentials "proxy-pool/internal/platform/credentials"
)

func TestInitialSnapshotSourceIssuesNextFullSnapshot(t *testing.T) {
	now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
	source, err := NewInitialSnapshotSource(epochReaderStub{epoch: 9}, time.Minute, func() time.Time { return now })
	if err != nil {
		t.Fatalf("NewInitialSnapshotSource(): %v", err)
	}
	updates, err := source.Watch(context.Background(), SnapshotWatchRequest{
		WorkerID: "worker-a", SessionID: "session-a", LastAppliedVersion: 4,
	})
	if err != nil {
		t.Fatalf("Watch(): %v", err)
	}
	full := <-updates
	if full.GetVersion() != 5 || full.GetOwnershipEpoch() != 9 || !full.GetValidUntil().AsTime().Equal(now.Add(time.Minute)) {
		t.Fatalf("snapshot = %+v", full)
	}
	checksum, err := snapshotwire.Checksum(full)
	if err != nil || string(checksum[:]) != string(full.GetChecksum()) {
		t.Fatalf("snapshot checksum = %x, %v; want %x", full.GetChecksum(), err, checksum)
	}
}

func TestRefreshingSnapshotSourceIssuesSequentialFullSnapshots(t *testing.T) {
	base := &refreshingSnapshotSourceStub{}
	source, err := NewRefreshingSnapshotSource(base, 10*time.Millisecond)
	if err != nil {
		t.Fatalf("NewRefreshingSnapshotSource() = %v", err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	updates, err := source.Watch(ctx, SnapshotWatchRequest{
		WorkerID: "worker-a", SessionID: "session-a", LastAppliedVersion: 4, LastChecksum: []byte{4},
	})
	if err != nil {
		t.Fatalf("Watch() = %v", err)
	}
	first := receiveSnapshot(t, updates)
	second := receiveSnapshot(t, updates)
	if first.GetVersion() != 5 || second.GetVersion() != 6 || second.GetVersion() <= first.GetVersion() {
		t.Fatalf("refreshed versions = (%d, %d)", first.GetVersion(), second.GetVersion())
	}
	if calls := base.Calls(); len(calls) < 2 || calls[0] != 4 || calls[1] != 5 {
		t.Fatalf("base requested versions = %v", calls)
	}
	if checksums := base.Checksums(); len(checksums) < 2 || string(checksums[0]) != string([]byte{4}) ||
		string(checksums[1]) != string([]byte{5}) {
		t.Fatalf("base requested checksums = %v", checksums)
	}
}

func TestRefreshingSnapshotSourceRefreshesImmediatelyAfterManagementNotification(t *testing.T) {
	base := &refreshingSnapshotSourceStub{}
	broker := NewSnapshotRefreshBroker()
	source, err := NewRefreshingSnapshotSourceWithRefreshEvents(base, time.Hour, broker)
	if err != nil {
		t.Fatalf("NewRefreshingSnapshotSourceWithRefreshEvents() = %v", err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	updates, err := source.Watch(ctx, SnapshotWatchRequest{
		WorkerID: "worker-a", SessionID: "session-a", LastAppliedVersion: 4, LastChecksum: []byte{4},
	})
	if err != nil {
		t.Fatalf("Watch() = %v", err)
	}
	if first := receiveSnapshot(t, updates); first.GetVersion() != 5 {
		t.Fatalf("initial version = %d, want 5", first.GetVersion())
	}
	broker.NotifySnapshotRefresh()
	if refreshed := receiveSnapshot(t, updates); refreshed.GetVersion() != 6 {
		t.Fatalf("refreshed version = %d, want 6", refreshed.GetVersion())
	}
	if calls := base.Calls(); len(calls) != 2 || calls[0] != 4 || calls[1] != 5 {
		t.Fatalf("base requested versions = %v", calls)
	}
}

func TestSnapshotRefreshBrokerBroadcastsCoalescedNotifications(t *testing.T) {
	broker := NewSnapshotRefreshBroker()
	first, unsubscribeFirst := broker.SubscribeSnapshotRefresh()
	second, unsubscribeSecond := broker.SubscribeSnapshotRefresh()
	defer unsubscribeSecond()

	broker.NotifySnapshotRefresh()
	broker.NotifySnapshotRefresh()
	select {
	case <-first:
	case <-time.After(time.Second):
		t.Fatal("first subscriber did not receive notification")
	}
	select {
	case <-second:
	case <-time.After(time.Second):
		t.Fatal("second subscriber did not receive notification")
	}
	select {
	case <-first:
		t.Fatal("first subscriber received an uncoalesced notification")
	default:
	}

	unsubscribeFirst()
	broker.NotifySnapshotRefresh()
	select {
	case <-first:
		t.Fatal("unsubscribed receiver received a notification")
	default:
	}
	select {
	case <-second:
	case <-time.After(time.Second):
		t.Fatal("remaining subscriber did not receive notification")
	}
}

func TestRefreshingSnapshotSourceRefreshesBeforeShortSnapshotExpiry(t *testing.T) {
	source, err := NewRefreshingSnapshotSource(&refreshingSnapshotSourceStub{}, 2*time.Second)
	if err != nil {
		t.Fatalf("NewRefreshingSnapshotSource() = %v", err)
	}
	now := time.Now().UTC()
	if delay := source.delayFor(&controlplanev1.WorkerSnapshot{
		GeneratedAt: timestamppb.New(now), ValidUntil: timestamppb.New(now.Add(100 * time.Millisecond)),
	}); delay != 50*time.Millisecond {
		t.Fatalf("delayFor(short-lived snapshot) = %s, want 50ms", delay)
	}
}

func TestOwnedSnapshotSourceBuildsBoundedProxySnapshot(t *testing.T) {
	now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
	expiresAt := now.Add(10 * time.Minute)
	usableUntil := now.Add(5 * time.Minute)
	source, err := NewOwnedSnapshotSource(epochReaderStub{epoch: 9}, snapshotReaderStub{proxies: []ownershipDomain.SnapshotProxy{{
		Proxy: proxyDomain.Proxy{
			ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080,
			SourceUpstream: "upstream-a", ExpiresAt: &expiresAt, UsableUntil: &usableUntil,
			MaxConcurrency: 7, State: proxyDomain.StateAvailable, Tags: map[string]string{"region": "cn"},
		},
		OwnershipEpoch: 4, LeaseExpiresAt: now.Add(time.Minute),
	}}}, time.Minute*2, 10, 4096, func() time.Time { return now })
	if err != nil {
		t.Fatalf("NewOwnedSnapshotSource(): %v", err)
	}
	updates, err := source.Watch(context.Background(), SnapshotWatchRequest{WorkerID: "worker-a", SessionID: "session-a"})
	if err != nil {
		t.Fatalf("Watch(): %v", err)
	}
	full := <-updates
	if full.GetVersion() != 1 || full.GetOwnershipEpoch() != 9 || len(full.GetProxies()) != 1 {
		t.Fatalf("snapshot = %+v", full)
	}
	proxy := full.GetProxies()[0]
	if proxy.GetOwnershipEpoch() != 4 || !proxy.GetUsableUntil().AsTime().Equal(now.Add(time.Minute)) ||
		!full.GetValidUntil().AsTime().Equal(now.Add(time.Minute)) {
		t.Fatalf("wire proxy = %+v, valid until = %s", proxy, full.GetValidUntil().AsTime())
	}
	checksum, err := snapshotwire.Checksum(full)
	if err != nil || string(checksum[:]) != string(full.GetChecksum()) {
		t.Fatalf("snapshot checksum = %x, %v; want %x", full.GetChecksum(), err, checksum)
	}
}

func TestOwnedSnapshotSourceIncludesRoutingInFullSnapshotChecksum(t *testing.T) {
	now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
	expiresAt := now.Add(10 * time.Minute)
	usableUntil := now.Add(5 * time.Minute)
	routing := &controlplanev1.RoutingRule{
		Name: "gateway", Enabled: true, HostRegex: ".*", Upstreams: []string{"upstream-a"},
		Strategy:      &controlplanev1.RoutingStrategy{Type: controlplanev1.StrategyType_STRATEGY_TYPE_RANDOM},
		OnUnavailable: controlplanev1.UnavailableAction_UNAVAILABLE_ACTION_REJECT,
	}
	source, err := NewOwnedSnapshotSource(epochReaderStub{epoch: 9}, snapshotReaderStub{proxies: []ownershipDomain.SnapshotProxy{{
		Proxy: proxyDomain.Proxy{
			ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080,
			SourceUpstream: "upstream-a", ExpiresAt: &expiresAt, UsableUntil: &usableUntil,
			MaxConcurrency: 7, State: proxyDomain.StateAvailable,
		},
		OwnershipEpoch: 4, LeaseExpiresAt: now.Add(time.Minute),
	}}}, time.Minute*2, 10, 4096, func() time.Time { return now }, routingSourceStub{rules: []*controlplanev1.RoutingRule{routing}})
	if err != nil {
		t.Fatalf("NewOwnedSnapshotSource(): %v", err)
	}
	updates, err := source.Watch(context.Background(), SnapshotWatchRequest{WorkerID: "worker-a", SessionID: "session-a"})
	if err != nil {
		t.Fatalf("Watch(): %v", err)
	}
	full := <-updates
	if len(full.GetRouting()) != 1 || full.GetRouting()[0].GetName() != "gateway" || full.GetRouting()[0] == routing {
		t.Fatalf("snapshot routing = %+v", full.GetRouting())
	}
	checksum, err := snapshotwire.Checksum(full)
	if err != nil || string(checksum[:]) != string(full.GetChecksum()) {
		t.Fatalf("snapshot checksum = %x, %v; want %x", full.GetChecksum(), err, checksum)
	}
}

func TestOwnedSnapshotSourceRejectsCredentialReferenceUntilMaterialIsAvailable(t *testing.T) {
	now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
	expiresAt := now.Add(time.Minute)
	usableUntil := now.Add(time.Minute)
	source, err := NewOwnedSnapshotSource(epochReaderStub{epoch: 1}, snapshotReaderStub{proxies: []ownershipDomain.SnapshotProxy{{
		Proxy: proxyDomain.Proxy{
			ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080,
			SourceUpstream: "upstream-a", ExpiresAt: &expiresAt, UsableUntil: &usableUntil,
			MaxConcurrency: 1, State: proxyDomain.StateAvailable, CredentialVersion: "v1",
		},
		OwnershipEpoch: 1, LeaseExpiresAt: now.Add(time.Minute),
	}}}, time.Minute, 10, 4096, func() time.Time { return now })
	if err != nil {
		t.Fatalf("NewOwnedSnapshotSource(): %v", err)
	}
	_, err = source.Watch(context.Background(), SnapshotWatchRequest{WorkerID: "worker-a", SessionID: "session-a"})
	if !errors.Is(err, ErrSnapshotCredentialsUnavailable) {
		t.Fatalf("Watch() error = %v, want ErrSnapshotCredentialsUnavailable", err)
	}
}

func TestOwnedSnapshotSourceIncludesReferencedCredentialMaterial(t *testing.T) {
	now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
	expiresAt := now.Add(time.Minute)
	usableUntil := now.Add(30 * time.Second)
	store, err := platformCredentials.NewMemoryStore(2)
	if err != nil {
		t.Fatalf("NewMemoryStore(): %v", err)
	}
	reference, err := store.Put(context.Background(), "provider-a", platformCredentials.Value{Username: "upstream", Password: "secret"})
	if err != nil {
		t.Fatalf("Put(): %v", err)
	}
	source, err := NewOwnedSnapshotSourceWithCredentials(epochReaderStub{epoch: 1}, snapshotReaderStub{proxies: []ownershipDomain.SnapshotProxy{{
		Proxy: proxyDomain.Proxy{
			ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080, SourceUpstream: "upstream-a",
			Username: "upstream", SecretRef: reference.SecretRef, CredentialVersion: reference.CredentialVersion,
			ExpiresAt: &expiresAt, UsableUntil: &usableUntil, MaxConcurrency: 1, State: proxyDomain.StateAvailable,
		}, OwnershipEpoch: 1, LeaseExpiresAt: now.Add(time.Minute),
	}}}, time.Minute, 10, 4096, func() time.Time { return now }, store)
	if err != nil {
		t.Fatalf("NewOwnedSnapshotSourceWithCredentials(): %v", err)
	}
	updates, err := source.Watch(context.Background(), SnapshotWatchRequest{WorkerID: "worker-a", SessionID: "session-a"})
	if err != nil {
		t.Fatalf("Watch(): %v", err)
	}
	full := <-updates
	if len(full.GetCredentials()) != 1 || full.GetCredentials()[0].GetSecretRef() != reference.SecretRef ||
		full.GetCredentials()[0].GetUsername() != "upstream" || full.GetCredentials()[0].GetPassword() != "secret" ||
		full.GetProxies()[0].GetCredentialVersion() != reference.CredentialVersion {
		t.Fatal("snapshot credential material or proxy credential reference is invalid")
	}
	checksum, err := snapshotwire.Checksum(full)
	if err != nil || string(checksum[:]) != string(full.GetChecksum()) {
		t.Fatalf("snapshot checksum = %x, %v; want %x", full.GetChecksum(), err, checksum)
	}
}

type epochReaderStub struct {
	epoch uint64
	err   error
}

type refreshingSnapshotSourceStub struct {
	mu        sync.Mutex
	calls     []uint64
	checksums [][]byte
}

func (source *refreshingSnapshotSourceStub) Watch(
	_ context.Context,
	request SnapshotWatchRequest,
) (<-chan *controlplanev1.WorkerSnapshot, error) {
	source.mu.Lock()
	source.calls = append(source.calls, request.LastAppliedVersion)
	source.checksums = append(source.checksums, append([]byte(nil), request.LastChecksum...))
	source.mu.Unlock()
	updates := make(chan *controlplanev1.WorkerSnapshot, 1)
	updates <- &controlplanev1.WorkerSnapshot{
		Version: request.LastAppliedVersion + 1, Checksum: []byte{byte(request.LastAppliedVersion + 1)},
	}
	return updates, nil
}

func (source *refreshingSnapshotSourceStub) Calls() []uint64 {
	source.mu.Lock()
	defer source.mu.Unlock()
	return append([]uint64(nil), source.calls...)
}

func (source *refreshingSnapshotSourceStub) Checksums() [][]byte {
	source.mu.Lock()
	defer source.mu.Unlock()
	checksums := make([][]byte, len(source.checksums))
	for index := range source.checksums {
		checksums[index] = append([]byte(nil), source.checksums[index]...)
	}
	return checksums
}

func receiveSnapshot(t *testing.T, updates <-chan *controlplanev1.WorkerSnapshot) *controlplanev1.WorkerSnapshot {
	t.Helper()
	select {
	case snapshot, ok := <-updates:
		if !ok || snapshot == nil {
			t.Fatal("snapshot updates closed before next full snapshot")
		}
		return snapshot
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for snapshot")
		return nil
	}
}

func (reader epochReaderStub) CurrentOwnershipEpoch(context.Context) (uint64, error) {
	return reader.epoch, reader.err
}

type snapshotReaderStub struct {
	proxies []ownershipDomain.SnapshotProxy
	err     error
}

type routingSourceStub struct {
	rules []*controlplanev1.RoutingRule
	err   error
}

func (source routingSourceStub) Read(context.Context) ([]*controlplanev1.RoutingRule, error) {
	return source.rules, source.err
}

func (reader snapshotReaderStub) ReadWorkerSnapshot(context.Context, string, int) ([]ownershipDomain.SnapshotProxy, error) {
	return reader.proxies, reader.err
}