package resolver

import (
	"context"
	"errors"
	"net"
	"testing"
	"time"

	"git.k8s-git-k6n.net/mats/fancy-dns/internal/config"
)

func TestAttemptTimeoutSilentUDP(t *testing.T) {
	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer pc.Close()
	p := NewPool(poolCfg(udpServer(pc.LocalAddr().String())))
	start := time.Now()
	_, _, err = p.Exchange(context.Background(), query("silent.test."), nil, 50*time.Millisecond)
	if err == nil || time.Since(start) > 500*time.Millisecond {
		t.Fatalf("elapsed %v, error %v", time.Since(start), err)
	}
}

func TestCallerCancellationInterruptsRead(t *testing.T) {
	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer pc.Close()
	p := NewPool(poolCfg(udpServer(pc.LocalAddr().String())))
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	result := make(chan error, 1)
	go func() { _, _, err := p.Exchange(ctx, query("cancel.test."), nil, 5*time.Second); result <- err }()
	_ = pc.SetReadDeadline(time.Now().Add(time.Second))
	if _, _, err := pc.ReadFrom(make([]byte, 4096)); err != nil {
		t.Fatal(err)
	}
	cancel()
	select {
	case err := <-result:
		if p.All()[0].Stats().Failures != 0 {
			t.Fatal("caller cancellation counted as an upstream failure")
		}
		if !errors.Is(err, context.Canceled) {
			t.Fatal(err)
		}
	case <-time.After(500 * time.Millisecond):
		t.Fatal("cancellation did not interrupt read")
	}
}

func TestFallbackSharesAttemptDeadline(t *testing.T) {
	// Each leg takes 150ms: a fresh budget for TCP would incorrectly succeed.
	up := newFakeUpstream(t, fakeOpt{delay: 150 * time.Millisecond, truncateUDP: true})
	p := NewPool(poolCfg(udpServer(up.addr)))
	start := time.Now()
	_, _, err := p.Exchange(context.Background(), query("fallback.test."), nil, 250*time.Millisecond)
	if err == nil || time.Since(start) > 450*time.Millisecond || up.tcpHits.Load() != 1 {
		t.Fatalf("elapsed %v, TCP hits %d, error %v", time.Since(start), up.tcpHits.Load(), err)
	}
}

func TestFailoverAfterAttemptTimeout(t *testing.T) {
	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer pc.Close()
	up := newFakeUpstream(t, fakeOpt{})
	cfg := poolCfg(udpServer(pc.LocalAddr().String()), udpServer(up.addr))
	cfg.Policy = "failover"
	p := NewPool(cfg)
	start := time.Now()
	_, addr, err := p.Exchange(context.Background(), query("failover.test."), nil, 50*time.Millisecond)
	if err != nil || addr != up.addr || time.Since(start) > 500*time.Millisecond {
		t.Fatalf("elapsed %v, addr %s, error %v", time.Since(start), addr, err)
	}
}

func TestAttemptTimeoutSilentStream(t *testing.T) {
	for _, proto := range []string{"tcp", "tls"} {
		t.Run(proto, func(t *testing.T) {
			lc, err := net.Listen("tcp", "127.0.0.1:0")
			if err != nil {
				t.Fatal(err)
			}
			defer lc.Close()
			accepted := make(chan net.Conn, 1)
			go func() { conn, _ := lc.Accept(); accepted <- conn }()
			p := NewPool(poolCfg(config.UpstreamServer{Addr: lc.Addr().String(), Proto: proto, InsecureSkipVerify: true}))
			start := time.Now()
			_, _, err = p.Exchange(context.Background(), query("silent.test."), nil, 50*time.Millisecond)
			if err == nil || time.Since(start) > 500*time.Millisecond {
				t.Errorf("elapsed %v, error %v", time.Since(start), err)
			}
			_ = lc.Close()
			if conn := <-accepted; conn != nil {
				_ = conn.Close()
			}
		})
	}
}