package whois

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

func TestRDAPClientRetryOn5xx(t *testing.T) {
	var hit int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		n := atomic.AddInt32(&hit, 1)
		if n == 1 {
			http.Error(w, "temporary", http.StatusInternalServerError)
			return
		}
		_, _ = w.Write([]byte(`{"objectClassName":"domain","ldhName":"example.com"}`))
	}))
	defer srv.Close()

	c, err := NewRDAPClientWithBootstrap(&RDAPBootstrap{
		Version: "1.0",
		Services: []RDAPService{
			{TLDs: []string{"com"}, URLs: []string{srv.URL}},
		},
	})
	if err != nil {
		t.Fatalf("NewRDAPClientWithBootstrap() error: %v", err)
	}
	c.SetRetryPolicy(RDAPRetryPolicy{
		MaxAttempts:    2,
		BaseDelay:      time.Millisecond,
		MaxDelay:       5 * time.Millisecond,
		RetryOn429:     true,
		RetryOn5xx:     true,
		RetryOnNetwork: true,
	})

	resp, err := c.QueryDomain(context.Background(), "example.com")
	if err != nil {
		t.Fatalf("QueryDomain() error: %v", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("unexpected status: %d", resp.StatusCode)
	}
	if got := atomic.LoadInt32(&hit); got != 2 {
		t.Fatalf("expected retry hit=2, got=%d", got)
	}
}

func TestRDAPClientQueryIPAndASN(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/ip/1.1.1.1":
			_, _ = w.Write([]byte(`{"objectClassName":"ip network","handle":"NET-1-1-1-0-1"}`))
		case "/autnum/13335":
			_, _ = w.Write([]byte(`{"objectClassName":"autnum","handle":"AS13335"}`))
		default:
			http.NotFound(w, r)
		}
	}))
	defer srv.Close()

	c, err := NewRDAPClientWithBootstrap(&RDAPBootstrap{
		Version: "1.0",
		Services: []RDAPService{
			{TLDs: []string{"com"}, URLs: []string{"https://rdap.example.com/"}},
		},
	})
	if err != nil {
		t.Fatalf("NewRDAPClientWithBootstrap() error: %v", err)
	}
	c.SetIPServers(srv.URL)
	c.SetASNServers(srv.URL)

	ipResp, err := c.QueryIP(context.Background(), "1.1.1.1")
	if err != nil {
		t.Fatalf("QueryIP() error: %v", err)
	}
	if ipResp.StatusCode != http.StatusOK || ipResp.Domain != "1.1.1.1" {
		t.Fatalf("unexpected ip response: status=%d domain=%q", ipResp.StatusCode, ipResp.Domain)
	}

	asnResp, err := c.QueryASN(context.Background(), "AS13335")
	if err != nil {
		t.Fatalf("QueryASN() error: %v", err)
	}
	if asnResp.StatusCode != http.StatusOK || asnResp.Domain != "AS13335" {
		t.Fatalf("unexpected asn response: status=%d domain=%q", asnResp.StatusCode, asnResp.Domain)
	}
}