package traverse

import (
	"net"
	"testing"

	"github.com/miekg/dns"
)

func TestResponseProcessNil(t *testing.T) {
	ref := NewReferral("example.com", dns.TypeA, ".", 0, 1.0, nil)
	r := NewResponse(ref, net.ParseIP("1.2.3.4"), nil)
	r.Process(nil)
	if r.Type != RespError {
		t.Errorf("Type = %d, want %d", r.Type, RespError)
	}
}

func TestResponseClassifyAnswer(t *testing.T) {
	msg := new(dns.Msg)
	msg.SetReply(new(dns.Msg))
	msg.Answer = append(msg.Answer, &dns.A{
		Hdr: dns.RR_Header{Name: "example.com.", Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 300},
		A:   net.ParseIP("93.184.216.34"),
	})

	ref := NewReferral("example.com", dns.TypeA, ".", 0, 1.0, nil)
	r := NewResponse(ref, net.ParseIP("1.2.3.4"), nil)
	r.Process(msg)

	if r.Type != RespAnswer {
		t.Errorf("Type = %d, want %d", r.Type, RespAnswer)
	}
	if r.Decoded == nil {
		t.Fatal("Decoded should not be nil")
	}
}

func TestResponseClassifyReferral(t *testing.T) {
	msg := new(dns.Msg)
	msg.Rcode = dns.RcodeSuccess
	msg.Authoritative = false
	msg.Ns = append(msg.Ns, &dns.NS{
		Hdr: dns.RR_Header{Name: "com.", Rrtype: dns.TypeNS, Class: dns.ClassINET, Ttl: 172800},
		Ns:  "a.gtld-servers.net.",
	})
	msg.Extra = append(msg.Extra, &dns.A{
		Hdr: dns.RR_Header{Name: "a.gtld-servers.net.", Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 172800},
		A:   net.ParseIP("192.5.6.30"),
	})

	ref := NewReferral("example.com", dns.TypeA, ".", 0, 1.0, nil)
	cache := NewInfoCache(nil)
	r := NewResponse(ref, net.ParseIP("1.2.3.4"), cache)
	r.Process(msg)

	if r.Type != RespReferral {
		t.Errorf("Type = %d, want %d", r.Type, RespReferral)
	}
}

func TestResponseClassifyNXDOMAIN(t *testing.T) {
	msg := new(dns.Msg)
	msg.Rcode = dns.RcodeNameError

	ref := NewReferral("example.com", dns.TypeA, ".", 0, 1.0, nil)
	r := NewResponse(ref, net.ParseIP("1.2.3.4"), nil)
	r.Process(msg)

	if r.Type != RespNXDOMAIN {
		t.Errorf("Type = %d, want %d", r.Type, RespNXDOMAIN)
	}
}

func TestResponseClassifyNODATA(t *testing.T) {
	msg := new(dns.Msg)
	msg.Rcode = dns.RcodeSuccess
	msg.Authoritative = true
	msg.Ns = append(msg.Ns, &dns.SOA{
		Hdr: dns.RR_Header{Name: "example.com.", Rrtype: dns.TypeSOA, Class: dns.ClassINET, Ttl: 3600},
	})

	ref := NewReferral("example.com", dns.TypeA, ".", 0, 1.0, nil)
	r := NewResponse(ref, net.ParseIP("1.2.3.4"), nil)
	r.Process(msg)

	if r.Type != RespNODATA {
		t.Errorf("Type = %d, want %d", r.Type, RespNODATA)
	}
}

func TestResponseClassifySERVFAIL(t *testing.T) {
	msg := new(dns.Msg)
	msg.Rcode = dns.RcodeServerFailure

	ref := NewReferral("example.com", dns.TypeA, ".", 0, 1.0, nil)
	r := NewResponse(ref, net.ParseIP("1.2.3.4"), nil)
	r.Process(msg)

	if r.Type != RespSERVFAIL {
		t.Errorf("Type = %d, want %d", r.Type, RespSERVFAIL)
	}
}

func TestResponseCNAMEFollow(t *testing.T) {
	msg := new(dns.Msg)
	msg.SetReply(new(dns.Msg))
	msg.Answer = append(msg.Answer,
		&dns.CNAME{
			Hdr:    dns.RR_Header{Name: "www.example.com.", Rrtype: dns.TypeCNAME, Class: dns.ClassINET},
			Target: "example.com.",
		},
	)

	ref := NewReferral("www.example.com", dns.TypeA, ".", 0, 1.0, nil)
	r := NewResponse(ref, net.ParseIP("1.2.3.4"), nil)
	r.Process(msg)

	if r.Type != RespCNAMEFollow {
		t.Errorf("Type = %d, want %d", r.Type, RespCNAMEFollow)
	}
}

func TestResponseCNAMEWithFinalAnswer(t *testing.T) {
	msg := new(dns.Msg)
	msg.SetReply(new(dns.Msg))
	msg.Answer = append(msg.Answer,
		&dns.CNAME{
			Hdr:    dns.RR_Header{Name: "www.example.com.", Rrtype: dns.TypeCNAME, Class: dns.ClassINET},
			Target: "example.com.",
		},
		&dns.A{
			Hdr: dns.RR_Header{Name: "example.com.", Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 300},
			A:   net.ParseIP("93.184.216.34"),
		},
	)

	ref := NewReferral("www.example.com", dns.TypeA, ".", 0, 1.0, nil)
	r := NewResponse(ref, net.ParseIP("1.2.3.4"), nil)
	r.Process(msg)

	if r.Type != RespAnswer {
		t.Errorf("Type = %d, want %d (CNAME with final A answer)", r.Type, RespAnswer)
	}
}

func TestResponseChildReferrals(t *testing.T) {
	msg := new(dns.Msg)
	msg.Rcode = dns.RcodeSuccess
	msg.Authoritative = false
	msg.Ns = append(msg.Ns,
		&dns.NS{Hdr: dns.RR_Header{Name: "example.com.", Rrtype: dns.TypeNS}, Ns: "a.gtld-servers.net."},
		&dns.NS{Hdr: dns.RR_Header{Name: "example.com.", Rrtype: dns.TypeNS}, Ns: "b.gtld-servers.net."},
	)
	msg.Extra = append(msg.Extra,
		&dns.A{Hdr: dns.RR_Header{Name: "a.gtld-servers.net.", Rrtype: dns.TypeA}, A: net.ParseIP("192.5.6.30")},
		&dns.A{Hdr: dns.RR_Header{Name: "b.gtld-servers.net.", Rrtype: dns.TypeA}, A: net.ParseIP("192.33.14.30")},
	)

	ref := NewReferral("example.com", dns.TypeA, ".", 0, 1.0, nil)
	cache := NewInfoCache(nil)
	r := NewResponse(ref, net.ParseIP("198.41.0.4"), cache)
	r.Process(msg)

	children := r.ChildReferrals()
	if len(children) != 2 {
		t.Fatalf("expected 2 child referrals, got %d", len(children))
	}

	if children[0].Name != "example.com." {
		t.Errorf("child[0] name = %q, want example.com.", children[0].Name)
	}
	if children[0].Bailiwick != "a.gtld-servers.net." {
		t.Errorf("child[0] bailiwick = %q, want a.gtld-servers.net.", children[0].Bailiwick)
	}
	if children[0].Prob != 0.5 {
		t.Errorf("child[0] prob = %f, want 0.5", children[0].Prob)
	}
	if children[0].Depth != 1 {
		t.Errorf("child[0] depth = %d, want 1", children[0].Depth)
	}
	if !children[0].HasAddresses() {
		t.Error("child[0] should have glue addresses")
	}

	if !children[1].HasAddresses() {
		t.Error("child[1] should have glue addresses")
	}

	nsNames := cache.LookupNS("example.com.")
	if len(nsNames) != 2 {
		t.Errorf("expected 2 NS in cache, got %d", len(nsNames))
	}
}

func TestResponseChildReferralsNonReferral(t *testing.T) {
	msg := new(dns.Msg)
	msg.SetReply(new(dns.Msg))
	msg.Answer = append(msg.Answer, &dns.A{
		Hdr: dns.RR_Header{Rrtype: dns.TypeA},
		A:   net.ParseIP("1.2.3.4"),
	})

	ref := NewReferral("example.com", dns.TypeA, ".", 0, 1.0, nil)
	r := NewResponse(ref, net.ParseIP("1.2.3.4"), nil)
	r.Process(msg)

	if children := r.ChildReferrals(); children != nil {
		t.Error("non-referral should not produce child referrals")
	}
}

func TestResponseCNAMEFollowReferral(t *testing.T) {
	msg := new(dns.Msg)
	msg.SetReply(new(dns.Msg))
	msg.Answer = append(msg.Answer,
		&dns.CNAME{
			Hdr:    dns.RR_Header{Name: "www.example.com.", Rrtype: dns.TypeCNAME},
			Target: "example.com.",
		},
	)

	ref := NewReferral("www.example.com", dns.TypeA, ".", 0, 1.0, nil)
	r := NewResponse(ref, net.ParseIP("1.2.3.4"), nil)
	r.Process(msg)

	follow := r.CNAMEFollowReferral()
	if follow == nil {
		t.Fatal("expected CNAME follow referral")
	}
	if follow.Name != "example.com." {
		t.Errorf("follow name = %q, want %q", follow.Name, "example.com.")
	}
	if follow.Depth != 1 {
		t.Errorf("follow depth = %d, want 1", follow.Depth)
	}
}

func TestResponseIsTerminal(t *testing.T) {
	tests := []struct {
		respType ResponseType
		want     bool
	}{
		{RespAnswer, true},
		{RespNODATA, true},
		{RespNXDOMAIN, true},
		{RespSERVFAIL, true},
		{RespError, true},
		{RespNSResolutionFailed, true},
		{RespReferral, false},
		{RespCNAMEFollow, false},
	}

	for _, tt := range tests {
		t.Run(tt.respType.String(), func(t *testing.T) {
			r := &Response{Type: tt.respType}
			if got := r.IsTerminal(); got != tt.want {
				t.Errorf("IsTerminal() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestResponseTypeString(t *testing.T) {
	tests := []struct {
		rt   ResponseType
		want string
	}{
		{RespReferral, "referral"},
		{RespAnswer, "answer"},
		{RespCNAMEFollow, "cname_follow"},
		{RespNODATA, "nodata"},
		{RespNXDOMAIN, "nxdomain"},
		{RespSERVFAIL, "servfail"},
		{RespError, "error"},
		{RespNSResolutionFailed, "ns_error"},
	}

	for _, tt := range tests {
		t.Run(tt.want, func(t *testing.T) {
			if got := tt.rt.String(); got != tt.want {
				t.Errorf("String() = %q, want %q", got, tt.want)
			}
		})
	}
}

func TestResponseChildReferralsProbabilityInheritance(t *testing.T) {
	msg := new(dns.Msg)
	msg.Rcode = dns.RcodeSuccess
	msg.Authoritative = false
	msg.Ns = append(msg.Ns,
		&dns.NS{Hdr: dns.RR_Header{Name: "com.", Rrtype: dns.TypeNS}, Ns: "a.gtld-servers.net."},
		&dns.NS{Hdr: dns.RR_Header{Name: "com.", Rrtype: dns.TypeNS}, Ns: "b.gtld-servers.net."},
		&dns.NS{Hdr: dns.RR_Header{Name: "com.", Rrtype: dns.TypeNS}, Ns: "c.gtld-servers.net."},
	)

	ref := NewReferral("example.com", dns.TypeA, ".", 0, 0.5, nil)
	r := NewResponse(ref, net.ParseIP("1.2.3.4"), nil)
	r.Process(msg)

	children := r.ChildReferrals()
	if len(children) != 3 {
		t.Fatalf("expected 3 children, got %d", len(children))
	}

	for _, c := range children {
		if c.Prob != 0.5/3.0 {
			t.Errorf("child prob = %f, want %f", c.Prob, 0.5/3.0)
		}
	}
}

func TestResponseChildReferralsEmptyAuthority(t *testing.T) {
	msg := new(dns.Msg)
	msg.Rcode = dns.RcodeSuccess
	msg.Authoritative = false

	ref := NewReferral("example.com", dns.TypeA, ".", 0, 1.0, nil)
	r := NewResponse(ref, net.ParseIP("1.2.3.4"), nil)
	r.Process(msg)

	children := r.ChildReferrals()
	if len(children) != 0 {
		t.Errorf("expected 0 children with empty authority, got %d", len(children))
	}
}

func TestResponseNilReferral(t *testing.T) {
	r := NewResponse(nil, net.ParseIP("1.2.3.4"), nil)
	children := r.ChildReferrals()
	if children != nil {
		t.Error("nil referral should produce no children")
	}

	follow := r.CNAMEFollowReferral()
	if follow != nil {
		t.Error("nil referral should produce no CNAME follow")
	}
}