package loadbalance

import (
	"context"
	"testing"

	"github.com/coredns/coredns/plugin"
	"github.com/coredns/coredns/plugin/pkg/dnstest"
	"github.com/coredns/coredns/plugin/test"

	"github.com/miekg/dns"
)

func TestLoadBalance(t *testing.T) {
	rm := RoundRobin{Next: handler()}

	// the first X records must be cnames after this test
	tests := []struct {
		answer        []dns.RR
		extra         []dns.RR
		cnameAnswer   int
		cnameExtra    int
		addressAnswer int
		addressExtra  int
		mxAnswer      int
		mxExtra       int
	}{
		{
			answer: []dns.RR{
				test.CNAME("cname1.region2.skydns.test.	300	IN	CNAME		cname2.region2.skydns.test."),
				test.CNAME("cname2.region2.skydns.test.	300	IN	CNAME		cname3.region2.skydns.test."),
				test.CNAME("cname5.region2.skydns.test.	300	IN	CNAME		cname6.region2.skydns.test."),
				test.CNAME("cname6.region2.skydns.test.	300	IN	CNAME		endpoint.region2.skydns.test."),
				test.A("endpoint.region2.skydns.test.		300	IN	A			10.240.0.1"),
				test.MX("mx.region2.skydns.test.			300	IN	MX		1	mx1.region2.skydns.test."),
				test.MX("mx.region2.skydns.test.			300	IN	MX		2	mx2.region2.skydns.test."),
				test.MX("mx.region2.skydns.test.			300	IN	MX		3	mx3.region2.skydns.test."),
			},
			cnameAnswer:   4,
			addressAnswer: 1,
			mxAnswer:      3,
		},
		{
			answer: []dns.RR{
				test.A("endpoint.region2.skydns.test.		300	IN	A			10.240.0.1"),
				test.MX("mx.region2.skydns.test.			300	IN	MX		1	mx1.region2.skydns.test."),
				test.CNAME("cname.region2.skydns.test.	300	IN	CNAME		endpoint.region2.skydns.test."),
			},
			cnameAnswer:   1,
			addressAnswer: 1,
			mxAnswer:      1,
		},
		{
			answer: []dns.RR{
				test.MX("mx.region2.skydns.test.			300	IN	MX		1	mx1.region2.skydns.test."),
				test.A("endpoint.region2.skydns.test.		300	IN	A			10.240.0.1"),
				test.A("endpoint.region2.skydns.test.		300	IN	A			10.240.0.2"),
				test.MX("mx.region2.skydns.test.			300	IN	MX		1	mx2.region2.skydns.test."),
				test.CNAME("cname2.region2.skydns.test.	300	IN	CNAME		cname3.region2.skydns.test."),
				test.A("endpoint.region2.skydns.test.		300	IN	A			10.240.0.3"),
				test.MX("mx.region2.skydns.test.			300	IN	MX		1	mx3.region2.skydns.test."),
			},
			extra: []dns.RR{
				test.A("endpoint.region2.skydns.test.		300	IN	A			10.240.0.1"),
				test.AAAA("endpoint.region2.skydns.test.	300	IN	AAAA		::1"),
				test.MX("mx.region2.skydns.test.			300	IN	MX		1	mx1.region2.skydns.test."),
				test.CNAME("cname2.region2.skydns.test.	300	IN	CNAME		cname3.region2.skydns.test."),
				test.MX("mx.region2.skydns.test.			300	IN	MX		1	mx2.region2.skydns.test."),
				test.A("endpoint.region2.skydns.test.		300	IN	A			10.240.0.3"),
				test.AAAA("endpoint.region2.skydns.test.	300	IN	AAAA		::2"),
				test.MX("mx.region2.skydns.test.			300	IN	MX		1	mx3.region2.skydns.test."),
			},
			cnameAnswer:   1,
			cnameExtra:    1,
			addressAnswer: 3,
			addressExtra:  4,
			mxAnswer:      3,
			mxExtra:       3,
		},
	}

	rec := dnstest.NewRecorder(&test.ResponseWriter{})

	for i, test := range tests {
		req := new(dns.Msg)
		req.SetQuestion("region2.skydns.test.", dns.TypeSRV)
		req.Answer = test.answer
		req.Extra = test.extra

		_, err := rm.ServeDNS(context.TODO(), rec, req)
		if err != nil {
			t.Errorf("Test %d: Expected no error, but got %s", i, err)
			continue
		}

		cname, address, mx, sorted := countRecords(rec.Msg.Answer)
		if !sorted {
			t.Errorf("Test %d: Expected CNAMEs, then AAAAs, then MX in Answer, but got mixed", i)
		}
		if cname != test.cnameAnswer {
			t.Errorf("Test %d: Expected %d CNAMEs in Answer, but got %d", i, test.cnameAnswer, cname)
		}
		if address != test.addressAnswer {
			t.Errorf("Test %d: Expected %d A/AAAAs in Answer, but got %d", i, test.addressAnswer, address)
		}
		if mx != test.mxAnswer {
			t.Errorf("Test %d: Expected %d MXs in Answer, but got %d", i, test.mxAnswer, mx)
		}

		cname, address, mx, sorted = countRecords(rec.Msg.Extra)
		if !sorted {
			t.Errorf("Test %d: Expected CNAMEs, then AAAAs, then MX in Extra, but got mixed", i)
		}
		if cname != test.cnameExtra {
			t.Errorf("Test %d: Expected %d CNAMEs in Extra, but got %d", i, test.cnameAnswer, cname)
		}
		if address != test.addressExtra {
			t.Errorf("Test %d: Expected %d A/AAAAs in Extra, but got %d", i, test.addressAnswer, address)
		}
		if mx != test.mxExtra {
			t.Errorf("Test %d: Expected %d MXs in Extra, but got %d", i, test.mxAnswer, mx)
		}
	}
}

func TestLoadBalanceXFR(t *testing.T) {
	rm := RoundRobin{Next: handler()}

	answer := []dns.RR{
		test.SOA("skydns.test.	30	IN	SOA	ns.dns.skydns.test. hostmaster.skydns.test. 1542756695 7200 1800 86400 30"),
		test.MX("mx.region2.skydns.test.			300	IN	MX		1	mx1.region2.skydns.test."),
		test.A("endpoint.region2.skydns.test.		300	IN	A			10.240.0.1"),
		test.A("endpoint.region2.skydns.test.		300	IN	A			10.240.0.2"),
		test.MX("mx.region2.skydns.test.			300	IN	MX		1	mx2.region2.skydns.test."),
		test.CNAME("cname2.region2.skydns.test.	300	IN	CNAME		cname3.region2.skydns.test."),
		test.A("endpoint.region2.skydns.test.		300	IN	A			10.240.0.3"),
		test.MX("mx.region2.skydns.test.			300	IN	MX		1	mx3.region2.skydns.test."),
		test.SOA("skydns.test.	30	IN	SOA	ns.dns.skydns.test. hostmaster.skydns.test. 1542756695 7200 1800 86400 30"),
	}

	for _, xfrtype := range []uint16{dns.TypeIXFR, dns.TypeAXFR} {
		rec := dnstest.NewRecorder(&test.ResponseWriter{})
		req := new(dns.Msg)
		req.SetQuestion("skydns.test.", xfrtype)
		req.Answer = answer
		_, err := rm.ServeDNS(context.TODO(), rec, req)
		if err != nil {
			t.Errorf("Expected no error, but got %s for %s", err, dns.TypeToString[xfrtype])
			continue
		}

		if rec.Msg.Answer[0].Header().Rrtype != dns.TypeSOA {
			t.Errorf("Expected SOA record for first answer for %s", dns.TypeToString[xfrtype])
		}

		if rec.Msg.Answer[len(rec.Msg.Answer)-1].Header().Rrtype != dns.TypeSOA {
			t.Errorf("Expected SOA record for last answer for %s", dns.TypeToString[xfrtype])
		}
	}
}

func countRecords(result []dns.RR) (cname int, address int, mx int, sorted bool) {
	const (
		Start = iota
		CNAMERecords
		ARecords
		MXRecords
		Any
	)

	// The order of the records is used to determine if the round-robin actually did anything.
	sorted = true
	cname = 0
	address = 0
	mx = 0
	state := Start
	for _, r := range result {
		switch r.Header().Rrtype {
		case dns.TypeCNAME:
			sorted = sorted && state <= CNAMERecords
			state = CNAMERecords
			cname++
		case dns.TypeA, dns.TypeAAAA:
			sorted = sorted && state <= ARecords
			state = ARecords
			address++
		case dns.TypeMX:
			sorted = sorted && state <= MXRecords
			state = MXRecords
			mx++
		default:
			state = Any
		}
	}
	return
}

func handler() plugin.Handler {
	return plugin.HandlerFunc(func(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
		w.WriteMsg(r)
		return dns.RcodeSuccess, nil
	})
}