package cache

import (
	"encoding/binary"
	"errors"
	"math/rand/v2"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"testing"
	"time"

	"github.com/miekg/dns"
	bolt "go.etcd.io/bbolt"
)

// fakeClock is a manually advanced clock so TTL, ageing and serve-stale
// behaviour is tested without sleeping.
type fakeClock struct {
	mu sync.Mutex
	t  time.Time
}

func newClock() *fakeClock { return &fakeClock{t: time.Date(2024, 5, 1, 12, 0, 0, 0, time.UTC)} }

func (c *fakeClock) Now() time.Time {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.t
}

func (c *fakeClock) Advance(d time.Duration) {
	c.mu.Lock()
	c.t = c.t.Add(d)
	c.mu.Unlock()
}

// answer builds a packed A response for name.
func answer(t *testing.T, name string, ttl uint32, ip string) *dns.Msg {
	t.Helper()
	m := new(dns.Msg)
	m.SetQuestion(dns.Fqdn(name), dns.TypeA)
	m.Response = true
	m.Answer = []dns.RR{&dns.A{
		Hdr: dns.RR_Header{Name: dns.Fqdn(name), Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: ttl},
		A:   net.ParseIP(ip).To4(),
	}}
	return m
}

func key(name string) Key { return QueryKey(name, dns.TypeA, dns.ClassINET, false, 0) }

func TestLRUGetPutAndHitKinds(t *testing.T) {
	clk := newClock()
	l := NewLRU(100, 4)
	m := answer(t, "example.com", 60, "1.2.3.4")
	wire, err := m.Pack()
	if err != nil {
		t.Fatalf("pack: %v", err)
	}
	k := key("example.com")
	now := clk.Now()
	l.Put(k, &Entry{Wire: wire, TTL: 60, CreatedAt: now, ExpiresAt: now.Add(time.Minute), StaleUntil: now.Add(90 * time.Second), Bytes: len(wire)})

	if e, kind := l.Get(k, now); e == nil || kind != Fresh {
		t.Fatalf("fresh hit: got entry=%v kind=%v, want non-nil/Fresh", e, kind)
	}
	// One second later the entry is still fresh; after the TTL it is stale but
	// answerable until StaleUntil, and past that it is a miss and dropped.
	clk.Advance(59 * time.Second)
	if _, kind := l.Get(k, clk.Now()); kind != Fresh {
		t.Errorf("at TTL-1s: kind = %v, want Fresh", kind)
	}
	clk.Advance(2 * time.Second)
	e, kind := l.Get(k, clk.Now())
	if e == nil || kind != Stale {
		t.Fatalf("just past TTL: got entry=%v kind=%v, want non-nil/Stale", e, kind)
	}
	if got := e.Hits.Load(); got != 3 {
		t.Errorf("entry Hits = %d, want 3 (every Get counts, fresh or stale)", got)
	}
	clk.Advance(31 * time.Second) // now past StaleUntil (TTL+30s)
	if e, kind := l.Get(k, clk.Now()); e != nil || kind != Miss {
		t.Fatalf("past stale window: got entry=%v kind=%v, want nil/Miss", e, kind)
	}
	if got := l.Stats().Entries; got != 0 {
		t.Errorf("entries after a dead read = %d, want 0 (the entry must be dropped)", got)
	}
	st := l.Stats()
	if st.Hits != 2 || st.StaleHits != 1 || st.Misses != 1 {
		t.Errorf("stats = %+v, want hits=2 staleHits=1 misses=1", st)
	}
}

func TestLRUEvictionKeepsRecentlyUsed(t *testing.T) {
	// One shard so the budget is exactly what we set: eviction order is then
	// fully determined and the test can assert it.
	l := NewLRU(4, 1)
	clk := newClock()
	now := clk.Now()
	put := func(name string) {
		l.Put(key(name), &Entry{TTL: 60, CreatedAt: now, ExpiresAt: now.Add(time.Minute), StaleUntil: now.Add(time.Minute), Bytes: 8})
	}
	put("a.example.com")
	put("b.example.com")
	put("c.example.com")
	// Touch a so it is the most recent, then overflow the 4-entry budget.
	if _, kind := l.Get(key("a.example.com"), now); kind != Fresh {
		t.Fatalf("a: kind = %v, want Fresh", kind)
	}
	put("d.example.com")
	put("e.example.com") // 5 inserts, budget 4

	if got := l.Stats().Entries; got > 4 {
		t.Fatalf("entries = %d, want <= 4 (per-shard budget)", got)
	}
	if _, kind := l.Get(key("a.example.com"), now); kind != Fresh {
		t.Errorf("a was recently used, want it kept")
	}
	if _, kind := l.Get(key("b.example.com"), now); kind != Miss {
		t.Errorf("b was the least recently used, want it evicted")
	}
	if l.Stats().Evictions == 0 {
		t.Error("evictions = 0, want > 0")
	}
}

func TestLRUDeleteWhereAndFlush(t *testing.T) {
	l := NewLRU(100, 2)
	now := newClock().Now()
	names := []string{"x.example.com", "y.example.com", "z.other.net"}
	for _, n := range names {
		l.Put(key(n), &Entry{TTL: 30, CreatedAt: now, ExpiresAt: now.Add(30 * time.Second), StaleUntil: now.Add(time.Minute), Bytes: 4})
	}
	dropped := l.DeleteWhere(func(k Key, _ *Entry) bool {
		return k.Name == "x.example.com" || k.Name == "y.example.com"
	})
	if dropped != 2 {
		t.Errorf("DeleteWhere = %d, want 2", dropped)
	}
	if _, kind := l.Get(key("x.example.com"), now); kind != Miss {
		t.Error("x.example.com still present after DeleteWhere")
	}
	if _, kind := l.Get(key("z.other.net"), now); kind != Fresh {
		t.Error("z.other.net must survive an unrelated DeleteWhere")
	}
	l.Flush()
	if got := l.Stats().Entries; got != 0 {
		t.Errorf("entries after Flush = %d, want 0", got)
	}
	if l.Stats().Bytes != 0 {
		t.Errorf("bytes after Flush = %d, want 0", l.Stats().Bytes)
	}
}

func TestLRUWalkVisitsEveryEntryOnce(t *testing.T) {
	l := NewLRU(400, 4) // budget is per shard: 100 entries per shard
	now := newClock().Now()
	for i := range 200 {
		n := "n" + strconv.Itoa(i) + ".example.com"
		l.Put(key(n), &Entry{TTL: 30, CreatedAt: now, ExpiresAt: now.Add(30 * time.Second), StaleUntil: now.Add(time.Minute), Bytes: 4})
	}
	seen := make(map[string]int, 200)
	l.Walk(func(k Key, _ *Entry) bool {
		seen[k.Name]++
		return true
	})
	if len(seen) != 200 {
		t.Fatalf("Walk visited %d names, want 200", len(seen))
	}
	for n, c := range seen {
		if c != 1 {
			t.Fatalf("%s visited %d times, want 1", n, c)
		}
	}
	// Returning false stops the walk.
	count := 0
	l.Walk(func(Key, *Entry) bool { count++; return false })
	if count != 1 {
		t.Errorf("Walk with stop visited %d entries, want 1", count)
	}
}

func TestSetCapacityEvictsDownToTheNewBudget(t *testing.T) {
	l := NewLRU(100, 1)
	now := newClock().Now()
	for i := range 50 {
		n := "n" + strconv.Itoa(i) + ".example.com"
		l.Put(key(n), &Entry{TTL: 30, CreatedAt: now, ExpiresAt: now.Add(30 * time.Second), StaleUntil: now.Add(time.Minute), Bytes: 4})
	}
	l.SetCapacity(10, 1)
	if got := l.Stats().Entries; got > 10 {
		t.Errorf("entries after SetCapacity(10) = %d, want <= 10", got)
	}
	if got := l.Stats().Shards; got != 1 {
		t.Errorf("shards after SetCapacity(10, 1) = %d, want 1", got)
	}
	// The cache must still work after a resize.
	l.Put(key("after-resize.example.com"), &Entry{TTL: 30, CreatedAt: now, ExpiresAt: now.Add(30 * time.Second), StaleUntil: now.Add(time.Minute), Bytes: 4})
	if _, kind := l.Get(key("after-resize.example.com"), now); kind != Fresh {
		t.Error("a Put after a resize must be readable")
	}
}

func TestLRUConcurrentPutGetResize(t *testing.T) {
	// Exercises the resizing path: SetCapacity rewrites the per-shard budget
	// while other goroutines insert and read. Run with -race.
	l := NewLRU(512, 8)
	now := newClock().Now()
	var wg sync.WaitGroup
	stop := make(chan struct{})
	wg.Add(1)
	go func() {
		defer wg.Done()
		for {
			select {
			case <-stop:
				return
			default:
				l.SetCapacity(200, 4)
				l.SetCapacity(4096, 16)
			}
		}
	}()
	for g := range 8 {
		wg.Add(1)
		go func(g int) {
			defer wg.Done()
			for i := range 500 {
				n := "g" + strconv.Itoa(g) + "-n" + strconv.Itoa(i) + ".example.com"
				k := key(n)
				l.Put(k, &Entry{TTL: 30, CreatedAt: now, ExpiresAt: now.Add(30 * time.Second), StaleUntil: now.Add(time.Minute), Bytes: 4})
				e, _ := l.Get(k, now)
				if e != nil {
					_ = e.Hits.Load()
				}
			}
		}(g)
	}
	// Readers keep walking while the resizer runs.
	wg.Add(1)
	go func() {
		defer wg.Done()
		for {
			select {
			case <-stop:
				return
			default:
				l.Walk(func(Key, *Entry) bool { return true })
			}
		}
	}()
	// Let the background loops run, then release them and wait for every
	// goroutine: closing before the wait is what keeps this from deadlocking.
	time.Sleep(20 * time.Millisecond)
	close(stop)
	wg.Wait()
	if l.Stats().Bytes < 0 {
		t.Errorf("byte accounting went negative: %d", l.Stats().Bytes)
	}
}

func TestKeyPackUnpackRoundTrip(t *testing.T) {
	cases := []Key{
		{Name: "example.com", Type: dns.TypeA, Class: dns.ClassINET},
		{Name: "a.b.c.d.example.com", Type: dns.TypeAAAA, Class: dns.ClassINET, DO: true},
		{Name: "x", Type: dns.TypeMX, Class: dns.ClassCHAOS, Scope: 3},
		{Name: "wildcard.example.com", Type: dns.TypeTXT, Class: dns.ClassINET, DO: true, Scope: 7},
	}
	for _, k := range cases {
		buf, err := packKey(nil, k)
		if err != nil {
			t.Fatalf("packKey(%v): %v", k, err)
		}
		got, rest, err := unpackKey(buf)
		if err != nil {
			t.Fatalf("unpackKey(%v): %v", k, err)
		}
		if len(rest) != 0 {
			t.Errorf("unpackKey(%v) left %d trailing bytes, want 0", k, len(rest))
		}
		if got != k {
			t.Errorf("round trip: got %+v, want %+v", got, k)
		}
	}
	// The wire name survives, so the canonical form comes back normalised.
	buf, err := packKey(nil, Key{Name: "Example.COM", Type: dns.TypeA, Class: dns.ClassINET})
	if err != nil {
		t.Fatalf("packKey: %v", err)
	}
	got, _, err := unpackKey(buf)
	if err != nil {
		t.Fatalf("unpackKey: %v", err)
	}
	if got.Name != "example.com" {
		t.Errorf("unpacked name = %q, want %q", got.Name, "example.com")
	}
}

func TestKeyPackRejectsBadInput(t *testing.T) {
	if _, err := packKey(nil, Key{Name: string(make([]byte, 256*4+1))}); !errors.Is(err, errKeyTooLong) {
		t.Errorf("packKey(oversized name) error = %v, want errKeyTooLong", err)
	}
	if _, err := packKey(nil, Key{Name: "a..b", Type: dns.TypeA}); err == nil {
		t.Error("packKey(name with an empty label) error = nil, want an error")
	}
	if _, _, err := unpackKey([]byte{0x05, 0x00}); !errors.Is(err, errBadKey) {
		t.Errorf("unpackKey(truncated) error = %v, want errBadKey", err)
	}
	if _, _, err := unpackKey(nil); !errors.Is(err, errBadKey) {
		t.Errorf("unpackKey(empty) error = %v, want errBadKey", err)
	}
}

func TestQueryKeyCanonicalises(t *testing.T) {
	a := QueryKey("Example.COM.", dns.TypeA, dns.ClassINET, false, 0)
	b := QueryKey("example.com", dns.TypeA, dns.ClassINET, false, 0)
	if a != b {
		t.Errorf("QueryKey did not canonicalise: %+v != %+v", a, b)
	}
	if a.Name != "example.com" {
		t.Errorf("canonical name = %q, want %q", a.Name, "example.com")
	}
}

func TestEntryAgeTTL(t *testing.T) {
	now := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
	cases := []struct {
		name string
		e    Entry
		now  time.Time
		want uint32
	}{
		{"fresh", Entry{TTL: 100, CreatedAt: now}, now.Add(30 * time.Second), 70},
		{"expiring", Entry{TTL: 100, CreatedAt: now}, now.Add(100 * time.Second), 0},
		{"past", Entry{TTL: 100, CreatedAt: now}, now.Add(500 * time.Second), 0},
		{"zero created keeps the stored TTL", Entry{TTL: 100}, now, 100},
		{"sub-second ageing truncates to 0 elapsed", Entry{TTL: 1, CreatedAt: now}, now.Add(400 * time.Millisecond), 1},
	}
	for i := range cases {
		// Take the address: Entry now carries an atomic counter, so the
		// struct must never be copied.
		tc := &cases[i]
		got := tc.e.AgeTTL(tc.now)
		if got != tc.want {
			t.Errorf("%s: AgeTTL = %d, want %d", tc.name, got, tc.want)
		}
	}
}

func TestManagerTTLFor(t *testing.T) {
	clk := newClock()
	m := NewManager(Options{MinTTL: 0, MaxTTL: time.Hour, NegativeTTL: 30 * time.Second, Now: clk.Now}, nil)

	soa := func(ttl uint32) *dns.SOA {
		return &dns.SOA{Hdr: dns.RR_Header{Name: "example.com.", Rrtype: dns.TypeSOA, Class: dns.ClassINET, Ttl: ttl}, Minttl: ttl}
	}

	t.Run("answers use the lowest TTL and clamp", func(t *testing.T) {
		m := NewManager(Options{MinTTL: 10 * time.Second, MaxTTL: time.Hour, Now: clk.Now}, nil)
		msg := answer(t, "a.example.com", 300, "1.2.3.4")
		msg.Answer = append(msg.Answer, &dns.A{
			Hdr: dns.RR_Header{Name: "a.example.com.", Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 5},
			A:   net.ParseIP("1.2.3.5").To4(),
		})
		ttl, neg := m.TTLFor(msg)
		if neg {
			t.Error("neg = true for a positive answer")
		}
		if ttl != 10 {
			t.Errorf("ttl = %d, want 10 (5s clamped up to MinTTL)", ttl)
		}
	})

	t.Run("clamp down to MaxTTL", func(t *testing.T) {
		m := NewManager(Options{MaxTTL: 60 * time.Second, Now: clk.Now}, nil)
		ttl, _ := m.TTLFor(answer(t, "a.example.com", 86400, "1.2.3.4"))
		if ttl != 60 {
			t.Errorf("ttl = %d, want 60 (clamped to MaxTTL)", ttl)
		}
	})

	t.Run("NXDOMAIN takes the SOA TTL", func(t *testing.T) {
		m := NewManager(Options{NegativeTTL: 30 * time.Second, Now: clk.Now}, nil)
		msg := new(dns.Msg)
		msg.SetQuestion("nope.example.com.", dns.TypeA)
		msg.Response = true
		msg.Rcode = dns.RcodeNameError
		msg.Ns = []dns.RR{soa(45)}
		ttl, neg := m.TTLFor(msg)
		if !neg {
			t.Error("neg = false for NXDOMAIN")
		}
		if ttl != 45 {
			t.Errorf("ttl = %d, want 45", ttl)
		}
	})

	t.Run("NODATA takes the SOA TTL and counts as negative", func(t *testing.T) {
		m := NewManager(Options{NegativeTTL: 30 * time.Second, Now: clk.Now}, nil)
		msg := new(dns.Msg)
		msg.SetQuestion("nope.example.com.", dns.TypeA)
		msg.Response = true
		msg.Ns = []dns.RR{soa(12)}
		ttl, neg := m.TTLFor(msg)
		if !neg || ttl != 12 {
			t.Errorf("NODATA: ttl = %d neg = %v, want 12/true", ttl, neg)
		}
	})

	t.Run("NODATA without an SOA falls back to NegativeTTL", func(t *testing.T) {
		m := NewManager(Options{NegativeTTL: 30 * time.Second, Now: clk.Now}, nil)
		msg := new(dns.Msg)
		msg.SetQuestion("nope.example.com.", dns.TypeA)
		msg.Response = true
		ttl, neg := m.TTLFor(msg)
		if !neg || ttl != 30 {
			t.Errorf("ttl = %d neg = %v, want 30/true", ttl, neg)
		}
	})

	t.Run("truncated answers are never cached", func(t *testing.T) {
		msg := answer(t, "big.example.com", 600, "1.2.3.4")
		msg.Truncated = true
		if ttl, _ := m.TTLFor(msg); ttl != 0 {
			t.Errorf("ttl = %d for a truncated answer, want 0", ttl)
		}
	})

	t.Run("transient rcodes are never cached", func(t *testing.T) {
		for _, rc := range []int{dns.RcodeServerFailure, dns.RcodeRefused, dns.RcodeNotAuth} {
			msg := answer(t, "a.example.com", 600, "1.2.3.4")
			msg.Rcode = rc
			if ttl, _ := m.TTLFor(msg); ttl != 0 {
				t.Errorf("rcode %d: ttl = %d, want 0", rc, ttl)
			}
		}
	})

	t.Run("nil message", func(t *testing.T) {
		if ttl, _ := m.TTLFor(nil); ttl != 0 {
			t.Errorf("ttl = %d, want 0", ttl)
		}
	})
}

func TestManagerPutResponseGetAndAgeing(t *testing.T) {
	clk := newClock()
	m := NewManager(Options{MaxTTL: time.Hour, Now: clk.Now}, nil)
	k := key("www.example.com")
	m.PutResponse(k, answer(t, "www.example.com", 120, "93.184.216.34"), "127.0.0.53:53")

	res := m.Get(k)
	if res.Kind != Fresh {
		t.Fatalf("kind = %v, want Fresh", res.Kind)
	}
	if res.TTL != 120 {
		t.Errorf("TTL = %d, want 120", res.TTL)
	}
	if res.Upstream != "127.0.0.53:53" {
		t.Errorf("upstream = %q, want the address it came from", res.Upstream)
	}
	if got := res.Msg.Answer[0].(*dns.A).A.String(); got != "93.184.216.34" {
		t.Errorf("answer = %s, want 93.184.216.34", got)
	}

	// The stored copy is private: mutating the returned message must not
	// change what the next reader sees.
	res.Msg.Answer[0].(*dns.A).A = net.ParseIP("9.9.9.9").To4()
	again := m.Get(k)
	if got := again.Msg.Answer[0].(*dns.A).A.String(); got != "93.184.216.34" {
		t.Errorf("answer after mutating the first result = %s, want an independent copy", got)
	}

	// TTLs age down monotonically.
	clk.Advance(30 * time.Second)
	if got := m.Get(k).TTL; got != 90 {
		t.Errorf("TTL after 30s = %d, want 90", got)
	}
	if got := m.Get(k).Msg.Answer[0].Header().Ttl; got != 90 {
		t.Errorf("record TTL after 30s = %d, want 90", got)
	}
}

func TestManagerServeStale(t *testing.T) {
	clk := newClock()
	m := NewManager(Options{MaxTTL: time.Hour, ServeStale: time.Minute, Now: clk.Now}, nil)
	k := key("stale.example.com")
	m.PutResponse(k, answer(t, "stale.example.com", 60, "10.0.0.1"), "up")

	clk.Advance(90 * time.Second) // 30s past the TTL, inside the stale window
	res := m.Get(k)
	if res.Kind != Stale {
		t.Fatalf("kind = %v, want Stale", res.Kind)
	}
	if !res.ServedStale {
		t.Error("ServedStale = false on a stale hit")
	}
	if res.TTL != 0 {
		t.Errorf("TTL = %d on a stale hit, want 0 (the handler forces TTL 1)", res.TTL)
	}
	clk.Advance(61 * time.Second) // now past TTL+ServeStale
	if res := m.Get(k); res.Kind != Miss {
		t.Errorf("kind = %v past the stale window, want Miss", res.Kind)
	}
}

func TestManagerServeStaleDisabled(t *testing.T) {
	clk := newClock()
	m := NewManager(Options{MaxTTL: time.Hour, Now: clk.Now}, nil)
	k := key("nostale.example.com")
	m.PutResponse(k, answer(t, "nostale.example.com", 60, "10.0.0.2"), "up")
	clk.Advance(61 * time.Second)
	if res := m.Get(k); res.Kind != Miss {
		t.Errorf("kind = %v, want Miss when ServeStale is 0", res.Kind)
	}
}

func TestManagerPutResponseRespectsZeroTTL(t *testing.T) {
	clk := newClock()
	m := NewManager(Options{MaxTTL: time.Hour, Now: clk.Now}, nil)
	k := key("nocache.example.com")
	m.PutResponse(k, answer(t, "nocache.example.com", 0, "10.0.0.3"), "up")
	if res := m.Get(k); res.Kind != Miss {
		t.Errorf("kind = %v, want Miss: a zero TTL means do not cache", res.Kind)
	}
}

func TestManagerDeleteNameAndSubtree(t *testing.T) {
	clk := newClock()
	m := NewManager(Options{MaxTTL: time.Hour, Now: clk.Now}, nil)
	names := []string{"blocked.example.com", "deep.sub.blocked.example.com", "keep.example.com", "notblocked.example.com"}
	for _, n := range names {
		m.PutResponse(key(n), answer(t, n, 600, "10.0.0.4"), "up")
	}
	// "sub.blocked.example.com" must not be caught by the "blocked.example.com"
	// subtree walk, but "blocked.example.com" must not catch
	// "notblocked.example.com" either.
	if n := m.DeleteName("Blocked.Example.Com."); n != 1 {
		t.Errorf("DeleteName = %d, want 1 (canonicalised, one name)", n)
	}
	if res := m.Get(key("blocked.example.com")); res.Kind != Miss {
		t.Error("deleted name still served")
	}
	if n := m.DeleteSubtree("deep.sub.blocked.example.com"); n != 1 {
		t.Errorf("DeleteSubtree = %d, want 1", n)
	}
	for _, n := range []string{"keep.example.com", "notblocked.example.com"} {
		if res := m.Get(key(n)); res.Kind != Fresh {
			t.Errorf("%s was dropped by an unrelated delete", n)
		}
	}
}

func TestManagerTopEntriesCountsAndFlush(t *testing.T) {
	clk := newClock()
	m := NewManager(Options{MaxTTL: time.Hour, Now: clk.Now}, nil)
	for range 5 {
		m.PutResponse(key("hot.example.com"), answer(t, "hot.example.com", 300, "10.0.0.5"), "up")
	}
	for range 2 {
		m.PutResponse(key("warm.example.com"), answer(t, "warm.example.com", 300, "10.0.0.6"), "up")
	}
	top := m.TopEntries(10)
	if len(top) < 2 {
		t.Fatalf("TopEntries returned %d rows, want at least 2", len(top))
	}
	// Each Put counts too, so a name queried once still shows up.
	var hot, warm uint64
	for _, e := range top {
		switch e.Name {
		case "hot.example.com":
			hot = e.Hits
		case "warm.example.com":
			warm = e.Hits
		}
	}
	if hot < 5 || warm < 2 {
		t.Errorf("top hits: hot=%d warm=%d, want >=5 and >=2", hot, warm)
	}
	if top[0].Name != "hot.example.com" {
		t.Errorf("first row = %s, want the hottest name first", top[0].Name)
	}
	m.Flush()
	if got := len(m.TopEntries(10)); got != 0 {
		t.Errorf("TopEntries after Flush = %d rows, want 0", got)
	}
}

func TestManagerStatsShape(t *testing.T) {
	clk := newClock()
	m := NewManager(Options{Size: 100, Shards: 4, MaxTTL: time.Hour, MinTTL: time.Second, Now: clk.Now}, nil)
	m.PutResponse(key("a.example.com"), answer(t, "a.example.com", 60, "10.0.0.7"), "up")
	m.Get(key("a.example.com"))
	m.Get(key("miss.example.com"))
	s := m.Stats()
	if !s.Enabled {
		t.Error("Enabled = false on a live manager")
	}
	if s.Entries != 1 {
		t.Errorf("Entries = %d, want 1", s.Entries)
	}
	if s.Hits != 1 || s.Misses != 1 {
		t.Errorf("hits = %d misses = %d, want 1/1", s.Hits, s.Misses)
	}
	if s.Size != 100 || s.Shards != 4 {
		t.Errorf("size/shards = %d/%d, want 100/4", s.Size, s.Shards)
	}
	if s.MemoryHuman == "" || s.Bytes <= 0 {
		t.Errorf("memory accounting: bytes = %d human = %q, want a positive byte count and a human string", s.Bytes, s.MemoryHuman)
	}
	if s.Persist.Enabled {
		t.Error("Persist.Enabled = true without a store")
	}
}

func TestHumanBytes(t *testing.T) {
	cases := []struct {
		in   int64
		want string
	}{
		{0, "0 B"},
		{512, "512 B"},
		{1024, "1.0 KiB"},
		{1536, "1.5 KiB"},
		{1024 * 1024, "1.0 MiB"},
		{3 * 1024 * 1024 * 1024 / 2, "1.50 GiB"},
		{-1024, "-1.0 KiB"},
	}
	for _, tc := range cases {
		if got := HumanBytes(tc.in); got != tc.want {
			t.Errorf("HumanBytes(%d) = %q, want %q", tc.in, got, tc.want)
		}
	}
}

func TestStoreSnapshotAndLoad(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "cache.db")
	st, err := OpenStore(StoreOptions{Path: path})
	if err != nil {
		t.Fatalf("OpenStore: %v", err)
	}
	defer st.Close()

	clk := newClock()
	now := clk.Now()
	l := NewLRU(100, 2)
	for i := range 25 {
		n := "n" + strconv.Itoa(i) + ".example.com"
		wire, err := answer(t, n, 600, "10.0.0.8").Pack()
		if err != nil {
			t.Fatalf("pack: %v", err)
		}
		l.Put(key(n), &Entry{
			Wire: wire, TTL: 600, Origin: "upstream", Upstream: "1.1.1.1:53",
			CreatedAt: now, ExpiresAt: now.Add(600 * time.Second), StaleUntil: now.Add(630 * time.Second), Bytes: len(wire),
		})
	}
	written, err := st.Snapshot(l, now)
	if err != nil {
		t.Fatalf("Snapshot: %v", err)
	}
	if written != 25 {
		t.Errorf("Snapshot wrote %d entries, want 25", written)
	}
	if fi, err := os.Stat(path); err != nil || fi.Size() == 0 {
		t.Fatalf("snapshot file missing or empty: %v", err)
	}

	// A fresh LRU restored from the file keeps the answers and their metadata.
	restored := NewLRU(100, 2)
	loaded, skipped, err := st.Load(restored, now)
	if err != nil {
		t.Fatalf("Load: %v", err)
	}
	if loaded != 25 || skipped != 0 {
		t.Errorf("Load = %d loaded, %d skipped, want 25/0", loaded, skipped)
	}
	e, kind := restored.Get(key("n7.example.com"), now)
	if e == nil || kind != Fresh {
		t.Fatalf("restored entry: e = %v kind = %v, want non-nil/Fresh", e, kind)
	}
	if e.Upstream != "1.1.1.1:53" || e.Origin != "upstream" {
		t.Errorf("restored metadata: origin = %q upstream = %q, want upstream/1.1.1.1:53", e.Origin, e.Upstream)
	}
	// The wire is intact.
	msg := new(dns.Msg)
	if err := msg.Unpack(e.Wire); err != nil {
		t.Fatalf("unpack restored wire: %v", err)
	}
	if len(msg.Answer) != 1 {
		t.Errorf("restored answer has %d records, want 1", len(msg.Answer))
	}
}

func TestStoreLoadSkipsExpiredEntries(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "cache.db")
	st, err := OpenStore(StoreOptions{Path: path})
	if err != nil {
		t.Fatalf("OpenStore: %v", err)
	}
	defer st.Close()

	clk := newClock()
	now := clk.Now()
	l := NewLRU(100, 1)
	for i, ttl := range []uint32{600, 600, 5} {
		n := "n" + strconv.Itoa(i) + ".example.com"
		wire, _ := answer(t, n, ttl, "10.0.0.9").Pack()
		l.Put(key(n), &Entry{
			Wire: wire, TTL: ttl, CreatedAt: now, ExpiresAt: now.Add(time.Duration(ttl) * time.Second),
			StaleUntil: now.Add(time.Duration(ttl)*time.Second + 30*time.Second), Bytes: len(wire),
		})
	}
	if _, err := st.Snapshot(l, now); err != nil {
		t.Fatalf("Snapshot: %v", err)
	}
	// Load 10s later: the 5s entry is dead, the two 600s entries are fresh.
	restored := NewLRU(100, 1)
	loaded, skipped, err := st.Load(restored, now.Add(10*time.Second))
	if err != nil {
		t.Fatalf("Load: %v", err)
	}
	if loaded != 2 {
		t.Errorf("loaded = %d, want 2", loaded)
	}
	if skipped != 1 {
		t.Errorf("skipped = %d, want 1", skipped)
	}
	if _, kind := restored.Get(key("n2.example.com"), now.Add(10*time.Second)); kind != Miss {
		t.Errorf("expired entry came back after a reload (kind = %v)", kind)
	}
}

func TestStoreTruncateAndStats(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "cache.db")
	st, err := OpenStore(StoreOptions{Path: path})
	if err != nil {
		t.Fatalf("OpenStore: %v", err)
	}
	defer st.Close()
	clk := newClock()
	now := clk.Now()
	l := NewLRU(100, 1)
	wire, _ := answer(t, "a.example.com", 600, "10.0.0.10").Pack()
	l.Put(key("a.example.com"), &Entry{Wire: wire, TTL: 600, CreatedAt: now, ExpiresAt: now.Add(600 * time.Second), StaleUntil: now.Add(630 * time.Second), Bytes: len(wire)})
	if _, err := st.Snapshot(l, now); err != nil {
		t.Fatalf("Snapshot: %v", err)
	}
	s := st.Stats()
	if !s.Enabled || s.Path != path || s.Entries != 1 || s.Writes == 0 || s.Bytes <= 0 {
		t.Errorf("Stats = %+v, want enabled with 1 entry and at least one write", s)
	}
	if !s.LastSnapshot.IsZero() {
		t.Logf("LastSnapshot = %v", s.LastSnapshot)
	}
	if err := st.Truncate(); err != nil {
		t.Fatalf("Truncate: %v", err)
	}
	if s := st.Stats(); s.Entries != 0 {
		t.Errorf("entries after Truncate = %d, want 0", s.Entries)
	}
	restored := NewLRU(100, 1)
	loaded, _, err := st.Load(restored, now)
	if err != nil {
		t.Fatalf("Load after Truncate: %v", err)
	}
	if loaded != 0 {
		t.Errorf("loaded = %d after Truncate, want 0", loaded)
	}
}

// TestManagerFlushDropsSnapshot checks the interaction the blocklist API
// depends on: flushing the cache must also clear the persisted snapshot, or a
// restart would resurrect answers an operator just invalidated.
func TestManagerFlushDropsSnapshot(t *testing.T) {
	dir := t.TempDir()
	st, err := OpenStore(StoreOptions{Path: filepath.Join(dir, "cache.db")})
	if err != nil {
		t.Fatalf("OpenStore: %v", err)
	}
	defer st.Close()
	clk := newClock()
	m := NewManager(Options{MaxTTL: time.Hour, Now: clk.Now}, st)
	k := key("flushed.example.com")
	m.PutResponse(k, answer(t, "flushed.example.com", 600, "10.0.0.11"), "up")
	if _, err := st.Snapshot(m.LRU(), clk.Now()); err != nil {
		t.Fatalf("Snapshot: %v", err)
	}
	m.Flush()
	restored := NewLRU(10, 1)
	loaded, _, err := st.Load(restored, clk.Now())
	if err != nil {
		t.Fatalf("Load: %v", err)
	}
	if loaded != 0 {
		t.Errorf("loaded = %d entries after Flush, want 0", loaded)
	}
	if res := m.Get(k); res.Kind != Miss {
		t.Errorf("kind = %v after Flush, want Miss", res.Kind)
	}
}

func TestLRUStatsHitRate(t *testing.T) {
	l := NewLRU(100, 1)
	now := newClock().Now()
	l.Put(key("a.example.com"), &Entry{TTL: 60, CreatedAt: now, ExpiresAt: now.Add(time.Minute), StaleUntil: now.Add(90 * time.Second), Bytes: 4})
	for range 3 {
		l.Get(key("a.example.com"), now)
	}
	for range 1 {
		l.Get(key("nothere.example.com"), now)
	}
	s := l.Stats()
	if s.Hits != 3 || s.Misses != 1 {
		t.Fatalf("hits/misses = %d/%d, want 3/1", s.Hits, s.Misses)
	}
	if s.HitRate < 0.74 || s.HitRate > 0.76 {
		t.Errorf("HitRate = %v, want ~0.75", s.HitRate)
	}
}

func BenchmarkLRUGet(b *testing.B) {
	l := NewLRU(1<<16, 16)
	now := time.Now()
	names := make([]string, 4096)
	for i := range names {
		names[i] = "n" + strconv.Itoa(i) + ".example.com"
		wire, _ := answer(&testing.T{}, names[i], 600, "10.0.0.12").Pack()
		l.Put(key(names[i]), &Entry{Wire: wire, TTL: 600, CreatedAt: now, ExpiresAt: now.Add(600 * time.Second), StaleUntil: now.Add(630 * time.Second), Bytes: len(wire)})
	}
	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		i := rand.Int()
		for pb.Next() {
			l.Get(key(names[i%len(names)]), now)
			i += 7
		}
	})
}

func BenchmarkLRUPut(b *testing.B) {
	l := NewLRU(1<<16, 16)
	now := time.Now()
	b.ResetTimer()
	for i := 0; b.Loop(); i++ {
		l.Put(key("n"+strconv.Itoa(i)+".example.com"), &Entry{TTL: 600, CreatedAt: now, ExpiresAt: now.Add(600 * time.Second), StaleUntil: now.Add(630 * time.Second), Bytes: 64})
	}
}

func BenchmarkManagerGetUnpack(b *testing.B) {
	clk := newClock()
	m := NewManager(Options{MaxTTL: time.Hour, Now: clk.Now}, nil)
	names := make([]string, 4096)
	for i := range names {
		names[i] = "n" + strconv.Itoa(i) + ".example.com"
		m.PutResponse(key(names[i]), answer(&testing.T{}, names[i], 600, "10.0.0.13"), "up")
	}
	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		i := rand.Int()
		for pb.Next() {
			m.Get(key(names[i%len(names)]))
			i += 7
		}
	})
}

// TestStoreSurvivesProcessRestart is the test the warm-restart path exists for:
// write a snapshot, close the file completely, open a brand new Store on the
// same path and load it into a fresh LRU. Anything that only works because the
// same *Store is still holding the file open fails here.
func TestStoreSurvivesProcessRestart(t *testing.T) {
	path := filepath.Join(t.TempDir(), "cache.db")
	clk := newClock()
	now := clk.Now()

	st, err := OpenStore(StoreOptions{Path: path})
	if err != nil {
		t.Fatalf("OpenStore: %v", err)
	}
	warm := NewLRU(100, 2)
	for i := range 5 {
		n := "r" + strconv.Itoa(i) + ".example.com"
		wire, _ := answer(t, n, 600, "10.1.0."+strconv.Itoa(i)).Pack()
		warm.Put(key(n), &Entry{
			Wire: wire, TTL: 600, Origin: "upstream", Upstream: "9.9.9.9:53",
			CreatedAt: now, ExpiresAt: now.Add(600 * time.Second), StaleUntil: now.Add(630 * time.Second),
			Bytes: len(wire),
		})
	}
	if written, err := st.Snapshot(warm, now); err != nil || written != 5 {
		t.Fatalf("Snapshot = %d, %v; want 5, nil", written, err)
	}
	if err := st.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}

	// Second "process".
	restarted, err := OpenStore(StoreOptions{Path: path})
	if err != nil {
		t.Fatalf("reopen: %v", err)
	}
	defer restarted.Close()
	cold := NewLRU(100, 2)
	loaded, skipped, err := restarted.Load(cold, clk.Now())
	if err != nil {
		t.Fatalf("Load: %v", err)
	}
	if loaded != 5 || skipped != 0 {
		t.Fatalf("Load = %d loaded, %d skipped, want 5/0", loaded, skipped)
	}
	e, kind := cold.Get(key("r3.example.com"), clk.Now())
	if e == nil || kind != Fresh {
		t.Fatalf("restored entry: entry=%v kind=%v, want non-nil/Fresh", e, kind)
	}
	if e.Upstream != "9.9.9.9:53" {
		t.Errorf("Upstream = %q, want 9.9.9.9:53", e.Upstream)
	}
	// The TTL a client sees must not jump back up across the restart: the
	// loader rewrites it to the time actually left.
	clk.Advance(120 * time.Second)
	if got, want := e.AgeTTL(clk.Now()), uint32(480); got != want {
		t.Errorf("AgeTTL after 120s = %d, want %d", got, want)
	}
}

// TestStoreLoadSkipsDamagedRecords feeds Load a file where individual records
// have been scribbled on, truncated or written by a future version. A snapshot
// is a cache, not a source of truth: every one of these must cost exactly one
// entry and none may fail the whole load.
func TestStoreLoadSkipsDamagedRecords(t *testing.T) {
	path := filepath.Join(t.TempDir(), "cache.db")
	clk := newClock()
	now := clk.Now()

	st, err := OpenStore(StoreOptions{Path: path})
	if err != nil {
		t.Fatalf("OpenStore: %v", err)
	}
	l := NewLRU(100, 1)
	for i := range 4 {
		n := "d" + strconv.Itoa(i) + ".example.com"
		wire, _ := answer(t, n, 600, "10.2.0."+strconv.Itoa(i)).Pack()
		l.Put(key(n), &Entry{
			Wire: wire, TTL: 600, CreatedAt: now, ExpiresAt: now.Add(600 * time.Second),
			StaleUntil: now.Add(630 * time.Second), Bytes: len(wire),
		})
	}
	if written, err := st.Snapshot(l, now); err != nil || written != 4 {
		t.Fatalf("Snapshot = %d, %v; want 4, nil", written, err)
	}
	// Three different kinds of damage, applied to the first three records
	// (the keys are packed, so the cursor is the only way to find them).
	damaged := map[int][]byte{
		0: futureVersionRecord,                  // written by a hypothetical v2
		1: {recordVersion, 0, 0x01, 0x02, 0x03}, // truncated mid-varint
		2: corruptRecord(t, encodeRecord(nil, &Entry{ // flipped bit
			Wire: []byte{0xde, 0xad}, TTL: 1, CreatedAt: now, ExpiresAt: now, StaleUntil: now,
		})),
	}
	if err := st.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketCache)
		cur := b.Cursor()
		k, _ := cur.First()
		for i := 0; k != nil; i++ {
			if bad, ok := damaged[i]; ok {
				if err := b.Put(k, bad); err != nil {
					return err
				}
			}
			k, _ = cur.Next()
		}
		return nil
	}); err != nil {
		t.Fatalf("seed damaged records: %v", err)
	}
	if err := st.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}

	again, err := OpenStore(StoreOptions{Path: path})
	if err != nil {
		t.Fatalf("reopen: %v", err)
	}
	defer again.Close()
	restored := NewLRU(100, 1)
	loaded, skipped, err := again.Load(restored, clk.Now())
	if err != nil {
		t.Fatalf("Load on a damaged file returned %v, want the good records to survive", err)
	}
	if loaded != 1 {
		t.Errorf("loaded = %d, want 1 (only the untouched record)", loaded)
	}
	if skipped != 3 {
		t.Errorf("skipped = %d, want 3", skipped)
	}
	if _, kind := restored.Get(key("d3.example.com"), clk.Now()); kind != Fresh {
		t.Errorf("intact record did not survive: kind = %v", kind)
	}
	if _, kind := restored.Get(key("d0.example.com"), clk.Now()); kind != Miss {
		t.Errorf("damaged record was loaded anyway: kind = %v", kind)
	}
}

// futureVersionRecord is a well-formed record from a hypothetical v2 writer.
var futureVersionRecord = func() []byte {
	buf := []byte{recordVersion + 1, 0}
	buf = binary.AppendUvarint(buf, 0)
	return buf
}()

// corruptRecord flips one bit inside the body so the checksum no longer matches.
func corruptRecord(t *testing.T, rec []byte) []byte {
	t.Helper()
	out := append([]byte(nil), rec...)
	out[len(out)/2] ^= 0x40
	return out
}

// TestRecordEncodeDecode pins the on-disk record format itself. The round trip
// through bbolt used to break because the encoder and decoder disagreed; this
// test fails at compile-review time rather than at warm-restart time.
func TestRecordEncodeDecode(t *testing.T) {
	now := newClock().Now()
	in := &Entry{
		Wire:       []byte{0x00, 0x01, 0x02, 0x03},
		TTL:        600,
		Origin:     "upstream",
		Upstream:   "1.1.1.1:53",
		CreatedAt:  now,
		ExpiresAt:  now.Add(600 * time.Second),
		StaleUntil: now.Add(630 * time.Second),
		Neg:        true,
		Truncated:  true,
	}
	out, err := decodeRecord(encodeRecord(nil, in))
	if err != nil {
		t.Fatalf("decodeRecord: %v", err)
	}
	if string(out.Wire) != string(in.Wire) {
		t.Errorf("Wire = % x, want % x", out.Wire, in.Wire)
	}
	for _, f := range []struct {
		name      string
		got, want any
	}{
		{"TTL", out.TTL, in.TTL},
		{"Origin", out.Origin, in.Origin},
		{"Upstream", out.Upstream, in.Upstream},
		{"Neg", out.Neg, in.Neg},
		{"Truncated", out.Truncated, in.Truncated},
	} {
		if f.got != f.want {
			t.Errorf("%s = %v, want %v", f.name, f.got, f.want)
		}
	}
	// Times come back as wall-clock instants rebuilt from Unix nanoseconds, so
	// they compare with Equal rather than == (the location bits differ).
	for _, f := range []struct {
		name      string
		got, want time.Time
	}{
		{"CreatedAt", out.CreatedAt, in.CreatedAt},
		{"ExpiresAt", out.ExpiresAt, in.ExpiresAt},
		{"StaleUntil", out.StaleUntil, in.StaleUntil},
	} {
		if !f.got.Equal(f.want) {
			t.Errorf("%s = %v, want %v", f.name, f.got, f.want)
		}
	}

	// A record whose times are all zero must survive too: the varints encode
	// them as 0 and come back as the zero time, not 1970.
	zero, err := decodeRecord(encodeRecord(nil, &Entry{Wire: []byte{0xff}}))
	if err != nil {
		t.Fatalf("decodeRecord(zero times): %v", err)
	}
	if !zero.ExpiresAt.IsZero() || !zero.CreatedAt.IsZero() {
		t.Errorf("zero times came back as %v/%v", zero.ExpiresAt, zero.CreatedAt)
	}

	for name, bad := range map[string][]byte{
		"empty":          nil,
		"header only":    {recordVersion, 0},
		"no wire":        encodeRecord(nil, &Entry{}),
		"bad checksum":   corruptRecord(t, encodeRecord(nil, in)),
		"bad version":    futureVersionRecord,
		"truncated tail": {recordVersion, 0, 0x80},
	} {
		if _, err := decodeRecord(bad); !errors.Is(err, errBadRecord) {
			t.Errorf("decodeRecord(%s) error = %v, want errBadRecord", name, err)
		}
	}
}