package traverse

import (
	"errors"
	"testing"

	"github.com/miekg/dns"
)

func cnameRR(owner, target string) dns.RR {
	return &dns.CNAME{
		Hdr:    dns.RR_Header{Name: dns.Fqdn(owner), Rrtype: dns.TypeCNAME, Class: dns.ClassINET},
		Target: dns.Fqdn(target),
	}
}

func soaRR(zone string) dns.RR {
	return &dns.SOA{
		Hdr:     dns.RR_Header{Name: dns.Fqdn(zone), Rrtype: dns.TypeSOA, Class: dns.ClassINET},
		Ns:      dns.Fqdn("ns1." + zone),
		Mbox:    dns.Fqdn("hostmaster." + zone),
		Serial:  1,
		Refresh: 3600, Retry: 600, Expire: 86400, Minttl: 300,
	}
}

func newMsg(qname string, qtype uint16, rcode int) *dns.Msg {
	m := new(dns.Msg)
	m.SetQuestion(dns.Fqdn(qname), qtype)
	m.Response = true
	m.Rcode = rcode
	return m
}

func decode(msg *dns.Msg, qname string, qtype uint16, bailiwick string) *DecodedQuery {
	return NewDecodedQuery(msg, nil, qname, dns.ClassINET, qtype, "192.0.2.1", bailiwick)
}

func TestDecodeException(t *testing.T) {
	dq := NewDecodedQuery(nil, errors.New("network timeout"), "example.com", dns.ClassINET, dns.TypeA, "192.0.2.1", "com")
	if dq.Status != StatusException {
		t.Fatalf("status = %s, want exception", dq.Status)
	}
	if dq.ExceptionMessage != "network timeout" {
		t.Errorf("exception message = %q", dq.ExceptionMessage)
	}
}

func TestDecodeNilMessageIsException(t *testing.T) {
	dq := NewDecodedQuery(nil, nil, "example.com", dns.ClassINET, dns.TypeA, "192.0.2.1", "com")
	if dq.Status != StatusException {
		t.Fatalf("status = %s, want exception", dq.Status)
	}
}

func TestDecodeErrorMessages(t *testing.T) {
	tests := []struct {
		rcode int
		want  string
	}{
		{dns.RcodeFormatError, "Format error (FORMERR)"},
		{dns.RcodeServerFailure, "Server failure (SERVFAIL)"},
		{dns.RcodeNameError, "No such domain (NXDOMAIN)"},
		{dns.RcodeNotImplemented, "Not implemented (NOTIMP)"},
		{dns.RcodeRefused, "Refused"},
		{dns.RcodeYXDomain, "YXDOMAIN"},
	}
	for _, tt := range tests {
		msg := newMsg("example.com", dns.TypeA, tt.rcode)
		dq := decode(msg, "example.com", dns.TypeA, "com")
		if dq.Status != StatusError {
			t.Errorf("rcode %d: status = %s, want error", tt.rcode, dq.Status)
		}
		if dq.ErrorMessage != tt.want {
			t.Errorf("rcode %d: message = %q, want %q", tt.rcode, dq.ErrorMessage, tt.want)
		}
	}
}

func TestDecodeAnswered(t *testing.T) {
	msg := newMsg("example.com", dns.TypeA, dns.RcodeSuccess)
	msg.Answer = append(msg.Answer, aRR("example.com", "93.184.216.34"))
	dq := decode(msg, "example.com", dns.TypeA, "example.com")
	if dq.Status != StatusAnswered {
		t.Fatalf("status = %s, want answered", dq.Status)
	}
	if len(dq.Answers) != 1 {
		t.Errorf("answers = %v", dq.Answers)
	}
	if dq.Endname != "example.com" {
		t.Errorf("endname = %q", dq.Endname)
	}
}

func TestDecodeAnsweredViaCNAMEChain(t *testing.T) {
	// In-bailiwick chain ends at a name that has the A answer.
	msg := newMsg("www.example.com", dns.TypeA, dns.RcodeSuccess)
	msg.Answer = append(msg.Answer,
		cnameRR("www.example.com", "web.example.com"),
		aRR("web.example.com", "93.184.216.34"),
	)
	dq := decode(msg, "www.example.com", dns.TypeA, "example.com")
	if dq.Status != StatusAnswered {
		t.Fatalf("status = %s, want answered", dq.Status)
	}
	if dq.Endname != "web.example.com" {
		t.Errorf("endname = %q, want web.example.com", dq.Endname)
	}
}

func TestDecodeRestartOnOutOfBailiwickCNAME(t *testing.T) {
	msg := newMsg("www.example.com", dns.TypeA, dns.RcodeSuccess)
	msg.Answer = append(msg.Answer, cnameRR("www.example.com", "cdn.example.org"))
	dq := decode(msg, "www.example.com", dns.TypeA, "example.com")
	if dq.Status != StatusRestart {
		t.Fatalf("status = %s, want restart", dq.Status)
	}
	if dq.Endname != "cdn.example.org" {
		t.Errorf("endname = %q", dq.Endname)
	}
}

func TestDecodeCNAMELoop(t *testing.T) {
	msg := newMsg("a.example.com", dns.TypeA, dns.RcodeSuccess)
	msg.Answer = append(msg.Answer,
		cnameRR("a.example.com", "b.example.com"),
		cnameRR("b.example.com", "a.example.com"),
	)
	dq := decode(msg, "a.example.com", dns.TypeA, "example.com")
	if dq.Status != StatusCNAMELoop {
		t.Fatalf("status = %s, want cname_loop", dq.Status)
	}
}

func TestDecodeCNAMESelfLoop(t *testing.T) {
	msg := newMsg("a.example.com", dns.TypeA, dns.RcodeSuccess)
	msg.Answer = append(msg.Answer, cnameRR("a.example.com", "A.EXAMPLE.COM"))
	dq := decode(msg, "a.example.com", dns.TypeA, "example.com")
	if dq.Status != StatusCNAMELoop {
		t.Fatalf("status = %s, want cname_loop (case-insensitive)", dq.Status)
	}
}

func TestDecodeCNAMEChainStopsAtOutOfBailiwickOwner(t *testing.T) {
	// Ruby stops following once the CURRENT owner leaves the bailiwick, so a
	// two-hop loop through an out-of-bailiwick owner is NOT cname_loop: the
	// unfollowed target equals the qname again, leaving endname == qname and
	// an empty authority — nodata.
	msg := newMsg("www.example.com", dns.TypeA, dns.RcodeSuccess)
	msg.Answer = append(msg.Answer,
		cnameRR("www.example.com", "a.example.org"),
		cnameRR("a.example.org", "www.example.com"),
	)
	dq := decode(msg, "www.example.com", dns.TypeA, "example.com")
	if dq.Status != StatusNoData {
		t.Fatalf("status = %s, want nodata", dq.Status)
	}
	if dq.Endname != "www.example.com" {
		t.Errorf("endname = %q, want www.example.com", dq.Endname)
	}
}

func TestDecodeCNAMEOwnerEqualToBailiwickStopsChain(t *testing.T) {
	// Owner exactly equal to the bailiwick is NOT strictly inside it, so the
	// chain stops after one hop even though another CNAME exists.
	msg := newMsg("example.com", dns.TypeA, dns.RcodeSuccess)
	msg.Answer = append(msg.Answer,
		cnameRR("example.com", "a.example.com"),
		cnameRR("a.example.com", "b.example.com"),
	)
	dq := decode(msg, "example.com", dns.TypeA, "example.com")
	if dq.Status != StatusRestart {
		t.Fatalf("status = %s, want restart", dq.Status)
	}
	if dq.Endname != "a.example.com" {
		t.Errorf("endname = %q, want a.example.com (unfollowed target)", dq.Endname)
	}
}

func TestDecodeQtypeCNAMEIsAnswered(t *testing.T) {
	// qtype=CNAME: the CNAME record IS the answer; the chain is never followed.
	msg := newMsg("www.example.com", dns.TypeCNAME, dns.RcodeSuccess)
	msg.Answer = append(msg.Answer,
		cnameRR("www.example.com", "web.example.com"),
		cnameRR("web.example.com", "www.example.com"),
	)
	dq := decode(msg, "www.example.com", dns.TypeCNAME, "example.com")
	if dq.Status != StatusAnswered {
		t.Fatalf("status = %s, want answered", dq.Status)
	}
	if dq.Endname != "www.example.com" {
		t.Errorf("endname = %q", dq.Endname)
	}
}

func TestDecodeQtypeANYMatchesAnyAnswer(t *testing.T) {
	msg := newMsg("example.com", dns.TypeANY, dns.RcodeSuccess)
	msg.Answer = append(msg.Answer, cnameRR("example.com", "elsewhere.example.net"))
	dq := decode(msg, "example.com", dns.TypeANY, "example.com")
	if dq.Status != StatusAnswered {
		t.Fatalf("status = %s, want answered (ANY matches CNAME)", dq.Status)
	}
}

func TestDecodeNoDataWithSOA(t *testing.T) {
	msg := newMsg("example.com", dns.TypeMX, dns.RcodeSuccess)
	msg.Ns = append(msg.Ns, soaRR("example.com"))
	dq := decode(msg, "example.com", dns.TypeMX, "example.com")
	if dq.Status != StatusNoData {
		t.Fatalf("status = %s, want nodata", dq.Status)
	}
}

func TestDecodeNoDataEmptyAuthority(t *testing.T) {
	msg := newMsg("example.com", dns.TypeMX, dns.RcodeSuccess)
	dq := decode(msg, "example.com", dns.TypeMX, "example.com")
	if dq.Status != StatusNoData {
		t.Fatalf("status = %s, want nodata", dq.Status)
	}
}

func TestDecodeNoDataSOAWinsOverNS(t *testing.T) {
	// SOA + NS in authority is a negative answer, not a referral.
	msg := newMsg("example.com", dns.TypeMX, dns.RcodeSuccess)
	msg.Ns = append(msg.Ns, soaRR("example.com"), nsRR("example.com", "ns1.example.com"))
	dq := decode(msg, "example.com", dns.TypeMX, "example.com")
	if dq.Status != StatusNoData {
		t.Fatalf("status = %s, want nodata", dq.Status)
	}
}

func TestDecodeReferral(t *testing.T) {
	msg := newMsg("www.example.com", dns.TypeA, dns.RcodeSuccess)
	msg.Ns = append(msg.Ns,
		nsRR("example.com", "NS1.Example.COM"),
		nsRR("example.com", "ns2.example.net"),
	)
	msg.Extra = append(msg.Extra, aRR("ns1.example.com", "1.2.3.4"))
	dq := decode(msg, "www.example.com", dns.TypeA, "com")
	if dq.Status != StatusReferral {
		t.Fatalf("status = %s, want referral", dq.Status)
	}
	if len(dq.AuthorityNames) != 2 || dq.AuthorityNames[0] != "ns1.example.com" || dq.AuthorityNames[1] != "ns2.example.net" {
		t.Errorf("authority names = %v", dq.AuthorityNames)
	}
}

func TestDecodeErrorBeatsAnswer(t *testing.T) {
	// rcode is checked before answers (step 3 before step 4).
	msg := newMsg("example.com", dns.TypeA, dns.RcodeServerFailure)
	msg.Answer = append(msg.Answer, aRR("example.com", "1.2.3.4"))
	dq := decode(msg, "example.com", dns.TypeA, "com")
	if dq.Status != StatusError {
		t.Fatalf("status = %s, want error", dq.Status)
	}
}

func TestDecodeCNAMEFollowedIntoNXDOMAIN(t *testing.T) {
	// CNAME followed first (step 2), then rcode (step 3): NXDOMAIN after an
	// in-message CNAME is still an error, but the loop check ran first.
	msg := newMsg("www.example.com", dns.TypeA, dns.RcodeNameError)
	msg.Answer = append(msg.Answer, cnameRR("www.example.com", "gone.example.com"))
	dq := decode(msg, "www.example.com", dns.TypeA, "example.com")
	if dq.Status != StatusError {
		t.Fatalf("status = %s, want error", dq.Status)
	}
	if dq.Endname != "gone.example.com" {
		t.Errorf("endname = %q", dq.Endname)
	}
}

func TestDecodeCacheablePartition(t *testing.T) {
	msg := newMsg("www.example.com", dns.TypeA, dns.RcodeSuccess)
	msg.Answer = append(msg.Answer, cnameRR("www.example.com", "cdn.example.org"))
	msg.Ns = append(msg.Ns, nsRR("example.org", "ns1.example.org"))
	msg.Extra = append(msg.Extra, aRR("ns1.example.org", "5.6.7.8"))
	opt := new(dns.OPT)
	opt.Hdr = dns.RR_Header{Name: ".", Rrtype: dns.TypeOPT}
	msg.Extra = append(msg.Extra, opt)

	dq := decode(msg, "www.example.com", dns.TypeA, "example.com")
	if len(dq.CacheableGood) != 1 {
		t.Errorf("good = %v, want just the CNAME", dq.CacheableGood)
	}
	if len(dq.CacheableBad) != 2 {
		t.Errorf("bad = %v, want NS+A for example.org", dq.CacheableBad)
	}
}

func TestInsideBailiwick(t *testing.T) {
	tests := []struct {
		name, bailiwick string
		want            bool
	}{
		{"anything.example.com", "", true},        // root bailiwick
		{"anything.example.com", ".", true},       // root as dot
		{"example.com", "example.com", true},      // exact
		{"Example.COM", "example.com", true},      // exact, case fold
		{"www.example.com", "EXAMPLE.com", true},  // suffix, case fold
		{"a.b.example.com", "example.com", true},  // deep suffix
		{"badexample.com", "example.com", false},  // label boundary
		{"example.org", "example.com", false},     // sideways
		{"com", "example.com", false},             // shallower
		{"www.example.com.", "example.com", true}, // trailing dot
	}
	for _, tt := range tests {
		if got := insideBailiwick(tt.name, tt.bailiwick); got != tt.want {
			t.Errorf("insideBailiwick(%q, %q) = %v, want %v", tt.name, tt.bailiwick, got, tt.want)
		}
	}
}

func TestIsLameReferral(t *testing.T) {
	tests := []struct {
		bailiwick, newBailiwick string
		want                    bool
	}{
		{"", "com", false},                        // root bailiwick never lame
		{"", "", false},                           // root to root
		{"com", "example.com", false},             // strictly deeper
		{"com", "a.b.example.com", false},         // much deeper
		{"COM", "example.com", false},             // case fold
		{"com", "com", true},                      // equal zone is lame
		{"com", "", true},                         // back to root is lame
		{"com", "org", true},                      // sideways is lame
		{"example.com", "com", true},              // shallower is lame
		{"example.com", "badexample.com", true},   // label boundary
		{"example.com", "www.example.com", false}, // deeper
	}
	for _, tt := range tests {
		if got := isLameReferral(tt.bailiwick, tt.newBailiwick); got != tt.want {
			t.Errorf("isLameReferral(%q, %q) = %v, want %v", tt.bailiwick, tt.newBailiwick, got, tt.want)
		}
	}
}

func TestMsgFollowCNAMEsNoChain(t *testing.T) {
	msg := newMsg("example.com", dns.TypeA, dns.RcodeSuccess)
	end, _, ok := msgFollowCNAMEs(msg, "Example.COM.", dns.TypeA, "com")
	if !ok || end != "example.com" {
		t.Errorf("end = %q ok=%v, want example.com true", end, ok)
	}
}

func TestMsgFollowCNAMEsRootBailiwickFollowsEverything(t *testing.T) {
	msg := newMsg("a.example.com", dns.TypeA, dns.RcodeSuccess)
	msg.Answer = append(msg.Answer,
		cnameRR("a.example.com", "b.example.org"),
		cnameRR("b.example.org", "c.example.net"),
		aRR("c.example.net", "1.2.3.4"),
	)
	end, targets, ok := msgFollowCNAMEs(msg, "a.example.com", dns.TypeA, "")
	if !ok || end != "c.example.net" {
		t.Errorf("end = %q ok=%v, want c.example.net true", end, ok)
	}
	if len(targets) != 2 || targets[0] != "b.example.org" || targets[1] != "c.example.net" {
		t.Errorf("chain targets = %v", targets)
	}
}