package whois

import (
	"bufio"
	"fmt"
	"net"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

func TestWhoisReferralLoopGuard(t *testing.T) {
	old, hadOld := defaultWhoisMap["loop"]
	defer func() {
		if hadOld {
			defaultWhoisMap["loop"] = old
		} else {
			delete(defaultWhoisMap, "loop")
		}
	}()

	var addrA, addrB string
	var hitA, hitB atomic.Int32

	addrB, stopB := startWhoisServerWithHandler(t, &hitB, func(_ string) string {
		return strings.Join([]string{
			"Domain Name: EXAMPLE.LOOP",
			"Registrar: LOOP-REG-B",
			"Whois Server: " + addrA,
			"",
		}, "\n")
	})
	defer stopB()

	addrA, stopA := startWhoisServerWithHandler(t, &hitA, func(_ string) string {
		return strings.Join([]string{
			"Domain Name: EXAMPLE.LOOP",
			"Registrar: LOOP-REG-A",
			"Whois Server: " + addrB,
			"",
		}, "\n")
	})
	defer stopA()

	defaultWhoisMap["loop"] = addrA

	c := NewClient()
	r, err := c.WhoisWithOptions("example.loop", QueryOptions{
		Level:            QueryBoth,
		ReferralMaxDepth: 8,
	})
	if err != nil {
		t.Fatalf("WhoisWithOptions() error: %v", err)
	}
	if !r.Exists() {
		t.Fatal("expected exists=true")
	}
	if got := hitA.Load(); got != 1 {
		t.Fatalf("expected A hit once (loop prevented), got=%d", got)
	}
	if got := hitB.Load(); got != 1 {
		t.Fatalf("expected B hit once, got=%d", got)
	}
}

func TestWhoisNegativeCacheNotFound(t *testing.T) {
	clearNegativeCache()
	old, hadOld := defaultWhoisMap["nfcache"]
	defer func() {
		if hadOld {
			defaultWhoisMap["nfcache"] = old
		} else {
			delete(defaultWhoisMap, "nfcache")
		}
		clearNegativeCache()
	}()

	var hit atomic.Int32
	addr, stop := startWhoisServerWithHandler(t, &hit, func(_ string) string {
		return strings.Join([]string{
			"No match for \"EXAMPLE.NFCACHE\"",
			"",
		}, "\n")
	})
	defer stop()
	defaultWhoisMap["nfcache"] = addr

	c := NewClient().SetNegativeCacheTTL(time.Minute)
	r1, err := c.WhoisWithOptions("example.nfcache", QueryOptions{Level: QueryAuto})
	if err != nil {
		t.Fatalf("first WhoisWithOptions() error: %v", err)
	}
	r2, err := c.WhoisWithOptions("example.nfcache", QueryOptions{Level: QueryAuto})
	if err != nil {
		t.Fatalf("second WhoisWithOptions() error: %v", err)
	}
	if r1.Exists() || r2.Exists() {
		t.Fatal("expected cached not-found results")
	}
	if got := hit.Load(); got != 1 {
		t.Fatalf("expected one whois query hit with negative cache, got=%d", got)
	}
}

func startWhoisServerWithHandler(t *testing.T, hit *atomic.Int32, handler func(query string) string) (addr string, shutdown func()) {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen mock whois failed: %v", err)
	}
	done := make(chan struct{})
	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				select {
				case <-done:
					return
				default:
					return
				}
			}
			go func(c net.Conn) {
				defer c.Close()
				_ = c.SetDeadline(time.Now().Add(5 * time.Second))
				line, _ := bufio.NewReader(c).ReadString('\n')
				if hit != nil {
					hit.Add(1)
				}
				resp := handler(strings.TrimSpace(line))
				if resp != "" {
					_, _ = fmt.Fprint(c, resp)
				}
			}(conn)
		}
	}()
	return ln.Addr().String(), func() {
		close(done)
		_ = ln.Close()
	}
}