package server

import (
	"sync"
	"testing"
)

// TestBumpAndLoadKeySemantics is a table over the counter keyspace. The keys
// are the only thing that binds a counter to a transport, and two callers can
// disagree about the spelling of a transport: bump() deliberately normalises
// an empty transport to "udp" (the pre-transport default) while load() does
// not normalise at all, so a caller that asks for the wrong spelling silently
// reads 0 rather than failing. Pinning both directions here is what stops that
// from turning into a counter that is written but never shown.
func TestBumpAndLoadKeySemantics(t *testing.T) {
	cases := []struct {
		name      string
		transport string
		field     string
		wantKey   string
	}{
		{name: "udp queries", transport: "udp", field: "queries", wantKey: "udp.queries"},
		{name: "tcp sent", transport: "tcp", field: "sent", wantKey: "tcp.sent"},
		{name: "dot truncated", transport: "dot", field: "truncated", wantKey: "dot.truncated"},
		{name: "doh dropped", transport: "doh", field: "dropped", wantKey: "doh.dropped"},
		// An empty transport is the historical spelling of "the default
		// transport", and classify()/servePacket() predate the per-transport
		// split, so an empty one must not create a ".queries" key that
		// ByTransport can never find.
		{name: "empty transport falls back to udp", transport: "", field: "queries", wantKey: "udp.queries"},
		{name: "unknown transport is kept verbatim", transport: "quic", field: "queries", wantKey: "quic.queries"},
		{name: "unknown field is kept verbatim", transport: "udp", field: "nonsense", wantKey: "udp.nonsense"},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			s := new(Server)
			s.bump(tc.transport, tc.field)

			// The key is the whole lookup path, so assert on what was
			// actually stored rather than inferring it from a read.
			if _, ok := s.counters.Load(tc.wantKey); !ok {
				keys := make([]string, 0, 4)
				s.counters.Range(func(k, _ any) bool {
					keys = append(keys, k.(string))
					return true
				})
				t.Fatalf("bump(%q, %q) stored %v, want key %q", tc.transport, tc.field, keys, tc.wantKey)
			}
			// The normalised key is what a reader of the default transport
			// sees. load() does not normalise, which is deliberate: it is
			// always called with a literal transport from the four-transport
			// list, so a caller that invented a spelling reads 0 rather than
			// silently getting another transport's numbers.
			if got := s.load(tc.wantKey[:len(tc.wantKey)-len(tc.field)-1], tc.field); got != 1 {
				t.Errorf("load(key transport, %q) = %d, want 1", tc.field, got)
			}
			// Counters accumulate, they do not reset.
			s.bump(tc.transport, tc.field)
			if got := s.load(tc.wantKey[:len(tc.wantKey)-len(tc.field)-1], tc.field); got != 2 {
				t.Errorf("after two bumps load(%q) = %d, want 2", tc.field, got)
			}
			// A field nobody wrote reads 0 instead of aliasing a sibling.
			if got := s.load(tc.wantKey[:len(tc.wantKey)-len(tc.field)-1], tc.field+"x"); got != 0 {
				t.Errorf("load(%q) = %d, want 0", tc.field+"x", got)
			}
		})
	}
}

// TestStatsAggregatesOnlyKnownFields is the other half of the key contract: a
// key whose suffix is not a HandlerStats field (sent, truncated, or garbage)
// must be ignored by the aggregate instead of being miscounted into it.
func TestStatsAggregatesOnlyKnownFields(t *testing.T) {
	cases := []struct {
		field  string
		weight func(*HandlerStats) uint64
	}{
		{field: "queries", weight: func(s *HandlerStats) uint64 { return s.Queries }},
		{field: "blocked", weight: func(s *HandlerStats) uint64 { return s.Blocked }},
		{field: "cached", weight: func(s *HandlerStats) uint64 { return s.Cached }},
		{field: "errors", weight: func(s *HandlerStats) uint64 { return s.Errors }},
		{field: "dropped", weight: func(s *HandlerStats) uint64 { return s.Dropped }},
	}
	for _, tc := range cases {
		t.Run(tc.field, func(t *testing.T) {
			s := new(Server)
			s.bump("tcp", tc.field)
			s.bump("tcp", tc.field)
			got := s.Stats()
			if v := tc.weight(&got); v != 2 {
				t.Errorf("Stats().%s = %d, want 2", tc.field, v)
			}
			// The per-transport view is the transport-scoped mirror of the
			// same bumps and must agree.
			ts := s.ByTransport()["tcp"]
			switch tc.field {
			case "queries":
				if ts.Queries != 2 {
					t.Errorf("ByTransport tcp.Queries = %d, want 2", ts.Queries)
				}
			case "blocked":
				if ts.Blocked != 2 {
					t.Errorf("ByTransport tcp.Blocked = %d, want 2", ts.Blocked)
				}
			case "cached":
				if ts.Cached != 2 {
					t.Errorf("ByTransport tcp.Cached = %d, want 2", ts.Cached)
				}
			case "errors":
				if ts.Errors != 2 {
					t.Errorf("ByTransport tcp.Errors = %d, want 2", ts.Errors)
				}
			case "dropped":
				if ts.Dropped != 2 {
					t.Errorf("ByTransport tcp.Dropped = %d, want 2", ts.Dropped)
				}
			}
		})
	}

	// Fields the aggregate does not know about stay out of it but do show up
	// per transport, which is how the console's transport tab works.
	s := new(Server)
	s.bump("udp", "sent")
	s.bump("udp", "truncated")
	s.bump("udp", "nonsense")
	if got := s.Stats(); got.Queries+got.Blocked+got.Cached+got.Errors+got.Dropped != 0 {
		t.Errorf("Stats() = %+v, want every aggregate field 0", got)
	}
	ts := s.ByTransport()["udp"]
	if ts.Sent != 1 || ts.Truncated != 1 {
		t.Errorf("ByTransport udp = %+v, want Sent 1 Truncated 1", ts)
	}
}

// TestStatsSnapshotIsSafeWhileCountersMove hammers bump() from many
// goroutines while Stats() and ByTransport() are read. Under -race this is the
// real assertion; the final check is an invariant (the aggregate equals the
// sum over transports) rather than a claim about what a reader happened to
// observe mid-flight, so it cannot flake.
func TestStatsSnapshotIsSafeWhileCountersMove(t *testing.T) {
	const (
		writers        = 8
		bumpsPerWriter = 200
		readers        = 4
	)
	s := new(Server)
	transports := []string{"udp", "tcp", "dot", "doh"}

	var start sync.WaitGroup
	start.Add(1)
	done := make(chan struct{})

	// Writers and readers wait on separate groups: the readers stop only when
	// `done` is closed, and `done` is closed once the writers are finished, so
	// folding both into one WaitGroup would wait for readers that are waiting
	// for the very Wait that waits for them.
	var writersWG, readersWG sync.WaitGroup
	for w := 0; w < writers; w++ {
		writersWG.Add(1)
		go func(w int) {
			defer writersWG.Done()
			transport := transports[w%len(transports)]
			start.Wait()
			for i := 0; i < bumpsPerWriter; i++ {
				s.bump(transport, "queries")
				s.bump(transport, "sent")
			}
		}(w)
	}
	for r := 0; r < readers; r++ {
		readersWG.Add(1)
		go func() {
			defer readersWG.Done()
			start.Wait()
			for {
				select {
				case <-done:
					return
				default:
				}
				_ = s.Stats()
				_ = s.ByTransport()
			}
		}()
	}
	start.Done()
	writersWG.Wait()
	close(done)
	readersWG.Wait()

	total := s.Stats()
	var sum uint64
	for _, tr := range transports {
		ts := s.ByTransport()[tr]
		sum += ts.Queries
		if ts.Queries != ts.Sent {
			t.Errorf("transport %q: Queries %d != Sent %d", tr, ts.Queries, ts.Sent)
		}
	}
	if sum != total.Queries {
		t.Errorf("sum of per-transport queries = %d, aggregate = %d", sum, total.Queries)
	}
	// One "queries" bump and one "sent" bump per iteration, so the aggregate
	// is writers*bumpsPerWriter queries and the same number of sends.
	if want := uint64(writers * bumpsPerWriter); total.Queries != want {
		t.Errorf("aggregate queries = %d, want %d", total.Queries, want)
	}
}