package dns

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

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 TestServerPoolRoundRobin(t *testing.T) {
	a, b := &fakeLookup{}, &fakeLookup{}
	p := &ServerPool{resolvers: []Resolver{a, b}}
	for i := 0; i < 4; i++ {
		if _, err := p.LookupIP(context.Background(), "example.com"); err != nil {
			t.Fatal(err)
		}
	}
	if a.calls.Load() != 2 || b.calls.Load() != 2 {
		t.Fatalf("round robin calls: a=%d b=%d, want 2/2", a.calls.Load(), b.calls.Load())
	}
}

func TestServerPoolFailover(t *testing.T) {
	a := &fakeLookup{ipErr: &net.DNSError{Err: "timeout", IsTimeout: true}}
	b := &fakeLookup{}
	p := &ServerPool{resolvers: []Resolver{a, b}}
	ips, err := p.LookupIP(context.Background(), "example.com")
	if err != nil {
		t.Fatal(err)
	}
	if len(ips) != 1 || a.calls.Load() != 1 {
		t.Fatalf("expected failover to second server, a.calls=%d", a.calls.Load())
	}
}

func TestServerPoolNoFailoverOnNXDOMAIN(t *testing.T) {
	a := &fakeLookup{ipErr: &net.DNSError{Err: "no such host", IsNotFound: true}}
	b := &fakeLookup{}
	p := &ServerPool{resolvers: []Resolver{a, b}}
	if _, err := p.LookupIP(context.Background(), "missing.example.com"); err == nil {
		t.Fatal("expected error")
	}
	if b.calls.Load() != 0 {
		t.Fatalf("NXDOMAIN must not fail over; b.calls=%d", b.calls.Load())
	}
}