package dns

import (
	"context"
	"net"
	"sync/atomic"
	"testing"
	"time"
)

type fakeLookup struct {
	calls    atomic.Int32
	cnameErr error
	ipErr    error
}

func (f *fakeLookup) LookupCNAME(ctx context.Context, host string) (string, error) {
	f.calls.Add(1)
	if f.cnameErr != nil {
		return "", f.cnameErr
	}
	return host, nil
}

func (f *fakeLookup) LookupIP(ctx context.Context, host string) ([]net.IP, error) {
	f.calls.Add(1)
	if f.ipErr != nil {
		return nil, f.ipErr
	}
	return []net.IP{net.ParseIP("1.2.3.4")}, nil
}

func TestServerCooldown(t *testing.T) {
	s := &server{}
	now := time.Now().UnixNano()
	for i := 0; i < serverFailureThreshold; i++ {
		s.failure(now)
	}
	if s.available(now) {
		t.Fatal("server should be in cooldown after repeated failures")
	}
	if !s.available(now + int64(serverCooldown) + 1) {
		t.Fatal("server should be available after cooldown")
	}
	s.success()
	if !s.available(now) || s.failures.Load() != 0 {
		t.Fatal("success should clear failures and cooldown")
	}
}

func TestPoolFailover(t *testing.T) {
	bad := &fakeLookup{ipErr: &net.DNSError{Err: "server misbehaving", IsTemporary: true}}
	good := &fakeLookup{}
	p := &ServerPool{servers: []*server{{r: bad}, {r: good}}}
	for i := 0; i < 50; i++ {
		ips, err := p.LookupIP(context.Background(), "example.com")
		if err != nil || len(ips) != 1 {
			t.Fatalf("lookup %d failed: %v", i, err)
		}
	}
}

func TestPoolNoFailoverOnNXDOMAIN(t *testing.T) {
	a := &fakeLookup{ipErr: &net.DNSError{Err: "no such host", IsNotFound: true}}
	b := &fakeLookup{ipErr: &net.DNSError{Err: "no such host", IsNotFound: true}}
	p := &ServerPool{servers: []*server{{r: a}, {r: b}}}
	if _, err := p.LookupIP(context.Background(), "missing.example.com"); err == nil {
		t.Fatal("expected error")
	}
	if got := a.calls.Load() + b.calls.Load(); got != 1 {
		t.Fatalf("NXDOMAIN made %d lookups, want exactly 1 (no failover)", got)
	}
}

func TestWarmDisablesDeadResolvers(t *testing.T) {
	good := &fakeLookup{}
	dead := &fakeLookup{ipErr: &net.DNSError{Err: "timeout", IsTimeout: true}}
	deadServer := &server{r: dead}
	p := &ServerPool{servers: []*server{deadServer, {r: good}}}
	healthy := p.Warm(context.Background(), "example.com", 4)
	if healthy != 1 {
		t.Fatalf("healthy = %d, want 1", healthy)
	}
	if deadServer.available(time.Now().UnixNano()) {
		t.Fatal("dead resolver should be in cooldown after warmup")
	}
}