package synthetic

import (
	"context"
	"fmt"
	"github.com/coredns/coredns/plugin/pkg/dnstest"
	"github.com/coredns/coredns/plugin/test"
	"github.com/miekg/dns"
	"net"

	"testing"
)

func TestServeDNSv4(t *testing.T) {
	ip, ipNet, err := net.ParseCIDR("192.0.2.0/24")
	if err != nil {
		log.Fatal(err)
	}

	s := synthetic{
		Next: test.ErrorHandler(),
		Config: syntheticConfig{
			net:     []*net.IPNet{{IP: ip, Mask: ipNet.Mask}},
			forward: "example.com",
			prefix:  "ip-",
		},
	}

	testCases := []struct {
		qname      string
		qtype      uint16
		wantrcode  int
		wantAnswer []string
		wantTTL    uint32
	}{
		{
			qname:      "ip-192-0-2-0.example.com",
			qtype:      dns.TypeA,
			wantrcode:  dns.RcodeSuccess,
			wantAnswer: []string{"192.0.2.0"},
			wantTTL:    0,
		},
		{
			qname:      "ip-192-0-2-255.example.com",
			qtype:      dns.TypeA,
			wantrcode:  dns.RcodeSuccess,
			wantAnswer: []string{"192.0.2.255"},
			wantTTL:    0,
		},
		{
			qname:      "ip-203-0-113-100.example.com",
			qtype:      dns.TypeA,
			wantrcode:  dns.RcodeServerFailure,
			wantAnswer: nil,
			wantTTL:    0,
		},
		{
			qname:      "123.2.0.192.in-addr.arpa.",
			qtype:      dns.TypePTR,
			wantrcode:  dns.RcodeSuccess,
			wantAnswer: []string{"ip-192-0-2-123.example.com."},
			wantTTL:    0,
		},
		{
			qname:      "ip-2001-db8-abcd--.example.com",
			qtype:      dns.TypeA,
			wantrcode:  dns.RcodeServerFailure,
			wantAnswer: nil,
			wantTTL:    0,
		},
	}

	for i, tc := range testCases {
		errorMsgPrefix := fmt.Sprintf("Test case %v for '%v' failed. Expected", i, tc.qname)
		ctx := context.TODO()
		w := dnstest.NewRecorder(&test.ResponseWriter{})
		r := new(dns.Msg)
		r.SetQuestion(tc.qname, tc.qtype)

		rc, err := s.ServeDNS(ctx, w, r)

		if err != nil {
			t.Errorf("%v no error, but got %v", errorMsgPrefix, err)
		}
		if rc != tc.wantrcode {
			t.Errorf("%v rcode %v, but got %v", errorMsgPrefix, tc.wantrcode, rc)
		}
		if len(w.Msg.Answer) == 0 && tc.wantAnswer != nil {
			t.Errorf("%v an answer, but got none", errorMsgPrefix)
			continue
		}
		if tc.wantAnswer == nil {
			if len(w.Msg.Answer) > 0 {
				t.Errorf("%v no answer, but got %v", errorMsgPrefix, w.Msg.Answer[0])
			}
			continue
		}
		if w.Msg.Answer[0].Header().Ttl != 0 {
			t.Errorf("%v TTL to be 0, but got %v", errorMsgPrefix, w.Msg.Answer[0].Header().Ttl)
		}
		if w.Msg.Answer[0].Header().Name != tc.qname {
			t.Errorf("%v Name to be %s, but got %s", errorMsgPrefix, tc.qname, w.Msg.Answer[0].Header().Name)
		}
		if w.Msg.Answer[0].Header().Rrtype != tc.qtype {
			t.Errorf("%v Type to be %d, but got %d", errorMsgPrefix, tc.qtype, w.Msg.Answer[0].Header().Rrtype)
		}
		if w.Msg.Answer[0].Header().Ttl != tc.wantTTL {
			t.Errorf("%v TTL to be %d, but got %d", errorMsgPrefix, tc.wantTTL, w.Msg.Answer[0].Header().Ttl)
		}
		switch tc.qtype {
		case dns.TypeA:
			if w.Msg.Answer[0].(*dns.A).A.String() != tc.wantAnswer[0] {
				t.Errorf("%v answer %s, but got %s", errorMsgPrefix, tc.wantAnswer[0], w.Msg.Answer[0].(*dns.A).A.String())
			}
		case dns.TypePTR:
			if w.Msg.Answer[0].(*dns.PTR).Ptr != tc.wantAnswer[0] {
				t.Errorf("%v answer %s, but got %s", errorMsgPrefix, tc.wantAnswer[0], w.Msg.Answer[0].(*dns.PTR).Ptr)
			}
		}
	}
}

func TestServeDNSv6(t *testing.T) {
	ip, ipNet, err := net.ParseCIDR("2001:db8:abcd::/48")
	if err != nil {
		log.Fatal(err)
	}

	s := synthetic{
		Next: test.ErrorHandler(),
		Config: syntheticConfig{
			net:     []*net.IPNet{{IP: ip, Mask: ipNet.Mask}},
			forward: "example.com.",
			ttl:     1800,
			prefix:  "ip6-",
		},
	}

	testCases := []struct {
		qname      string
		qtype      uint16
		wantrcode  int
		wantAnswer []string
		wantTTL    uint32
	}{
		{
			qname:      "ip6-2001-db8-abcd--.example.com",
			qtype:      dns.TypeAAAA,
			wantrcode:  dns.RcodeSuccess,
			wantAnswer: []string{"2001:db8:abcd::"},
			wantTTL:    1800,
		},
		{
			qname:      "ip6-2001-db8-abcd-1234-4567-890a-bcde-f123.example.com",
			qtype:      dns.TypeAAAA,
			wantrcode:  dns.RcodeSuccess,
			wantAnswer: []string{"2001:db8:abcd:1234:4567:890a:bcde:f123"},
			wantTTL:    1800,
		},
		{
			qname:      "ip6-2001-db8-1234--.example.com",
			qtype:      dns.TypeAAAA,
			wantrcode:  dns.RcodeServerFailure,
			wantAnswer: nil,
			wantTTL:    1800,
		},
		{
			qname:      "3.2.1.f.e.d.c.b.a.0.9.8.7.6.5.4.4.3.2.1.d.c.b.a.8.b.d.0.1.0.0.2.ip6.arpa.",
			qtype:      dns.TypePTR,
			wantrcode:  dns.RcodeSuccess,
			wantAnswer: []string{"ip6-2001-db8-abcd-1234-4567-890a-bcde-f123.example.com."},
			wantTTL:    1800,
		},
		{
			qname:      "ip6-192-0-2-0.example.com",
			qtype:      dns.TypeAAAA,
			wantrcode:  dns.RcodeServerFailure,
			wantAnswer: nil,
			wantTTL:    0,
		},
	}

	for i, tc := range testCases {
		errorMsgPrefix := fmt.Sprintf("Test case %v for '%v' failed. Expected", i, tc.qname)
		ctx := context.TODO()
		w := dnstest.NewRecorder(&test.ResponseWriter{})
		r := new(dns.Msg)
		r.SetQuestion(tc.qname, tc.qtype)

		rc, err := s.ServeDNS(ctx, w, r)

		if err != nil {
			t.Errorf("%v no error, but got %v", errorMsgPrefix, err)
		}
		if rc != tc.wantrcode {
			t.Errorf("%v rcode %v, but got %v", errorMsgPrefix, tc.wantrcode, rc)
		}
		if len(w.Msg.Answer) == 0 && tc.wantAnswer != nil {
			t.Errorf("%v an answer, but got none", errorMsgPrefix)
			continue
		}
		if tc.wantAnswer == nil {
			if len(w.Msg.Answer) > 0 {
				t.Errorf("%v no answer, but got %v", errorMsgPrefix, w.Msg.Answer[0])
			}
			continue
		}
		if w.Msg.Answer[0].Header().Name != tc.qname {
			t.Errorf("%v Name to be %s, but got %s", errorMsgPrefix, tc.qname, w.Msg.Answer[0].Header().Name)
		}
		if w.Msg.Answer[0].Header().Rrtype != tc.qtype {
			t.Errorf("%v Type to be %d, but got %d", errorMsgPrefix, tc.qtype, w.Msg.Answer[0].Header().Rrtype)
		}
		if w.Msg.Answer[0].Header().Ttl != tc.wantTTL {
			t.Errorf("%v TTL to be %d, but got %d", errorMsgPrefix, tc.wantTTL, w.Msg.Answer[0].Header().Ttl)
		}
		switch tc.qtype {
		case dns.TypeA:
			if w.Msg.Answer[0].(*dns.A).A.String() != tc.wantAnswer[0] {
				t.Errorf("%v answer %s, but got %s", errorMsgPrefix, tc.wantAnswer[0], w.Msg.Answer[0].(*dns.A).A.String())
			}
		case dns.TypePTR:
			if w.Msg.Answer[0].(*dns.PTR).Ptr != tc.wantAnswer[0] {
				t.Errorf("%v answer %s, but got %s", errorMsgPrefix, tc.wantAnswer[0], w.Msg.Answer[0].(*dns.PTR).Ptr)
			}
		}
	}
}

// MockSuccessPlugin is a mock plugin that always responds with a successful DNS response.
type MockSuccessPlugin struct{}

// Name returns the plugin name.
func (m MockSuccessPlugin) Name() string { return "mock" }

// ServeDNS simulates a successful DNS response.
func (m MockSuccessPlugin) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
	m1 := new(dns.Msg)
	m1.SetReply(r)
	hdr := dns.RR_Header{Name: r.Question[0].Name, Rrtype: r.Question[0].Qtype, Class: r.Question[0].Qclass, Ttl: 0}
	switch r.Question[0].Qtype {
	case dns.TypeA:
		m1.Answer = append(m1.Answer, &dns.A{Hdr: hdr, A: net.ParseIP("192.0.2.100")})
	case dns.TypeAAAA:
		m1.Answer = append(m1.Answer, &dns.AAAA{Hdr: hdr, AAAA: net.ParseIP("2001:db8::100")})
	}
	w.WriteMsg(m1)
	return dns.RcodeSuccess, nil
}

func TestServeDNSNextPluginRespondsSuccess(t *testing.T) {
	ip, ipNet, err := net.ParseCIDR("192.0.2.0/24")
	if err != nil {
		log.Fatal(err)
	}
	ip6, ipNet6, err := net.ParseCIDR("2001:db8:abcd::/48")
	if err != nil {
		log.Fatal(err)
	}

	s := synthetic{
		Next: MockSuccessPlugin{},
		Config: syntheticConfig{
			net:     []*net.IPNet{{IP: ip, Mask: ipNet.Mask}, {IP: ip6, Mask: ipNet6.Mask}},
			forward: "example.com",
			prefix:  "ip-",
		},
	}

	testCases := []struct {
		qname      string
		qtype      uint16
		wantrcode  int
		wantAnswer []string
		wantTTL    uint32
	}{
		{
			qname:      "ip-192-0-2-1.example.com",
			qtype:      dns.TypeA,
			wantrcode:  dns.RcodeSuccess,
			wantAnswer: []string{"192.0.2.1"},
			wantTTL:    0,
		},
		{
			qname:      "foobar.example.com",
			qtype:      dns.TypeA,
			wantrcode:  dns.RcodeSuccess,
			wantAnswer: []string{"192.0.2.100"},
			wantTTL:    0,
		},
		{
			qname:      "ip-2001-db8-abcd--1.example.com",
			qtype:      dns.TypeAAAA,
			wantrcode:  dns.RcodeSuccess,
			wantAnswer: []string{"2001:db8:abcd::1"},
			wantTTL:    0,
		},
		{
			qname:      "foobar.example.com",
			qtype:      dns.TypeAAAA,
			wantrcode:  dns.RcodeSuccess,
			wantAnswer: []string{"2001:db8:abcd::100"},
			wantTTL:    0,
		},
		{
			qname:      "ip-2001-db8-abcd--1.example.com",
			qtype:      dns.TypeA,
			wantrcode:  dns.RcodeSuccess,
			wantAnswer: nil,
			wantTTL:    0,
		},
		{
			qname:      "ip-192-0-2-1.example.com",
			qtype:      dns.TypeAAAA,
			wantrcode:  dns.RcodeSuccess,
			wantAnswer: nil,
			wantTTL:    0,
		},
		{
			qname:      "3.2.1.f.e.d.c.b.a.0.9.8.7.6.5.4.4.3.2.1.d.c.b.a.8.b.d.0.1.0.0.2.ip6.arpa.",
			qtype:      dns.TypePTR,
			wantrcode:  dns.RcodeSuccess,
			wantAnswer: []string{"ip6-2001-db8-abcd-1234-4567-890a-bcde-f123.example.com."},
			wantTTL:    0,
		},
		{
			qname:      "123.2.0.192.in-addr.arpa.",
			qtype:      dns.TypePTR,
			wantrcode:  dns.RcodeSuccess,
			wantAnswer: []string{"ip-192-0-2-123.example.com."},
			wantTTL:    0,
		},
	}

	for i, tc := range testCases {
		errorMsgPrefix := fmt.Sprintf("Test case %v for '%v' failed. Expected", i, tc.qname)
		ctx := context.TODO()
		w := dnstest.NewRecorder(&test.ResponseWriter{})
		r := new(dns.Msg)
		r.SetQuestion(tc.qname, tc.qtype)

		rc, err := s.ServeDNS(ctx, w, r)

		if err != nil {
			t.Errorf("%v no error, but got %v", errorMsgPrefix, err)
		}
		if rc != tc.wantrcode {
			t.Errorf("%v rcode %v, but got %v", errorMsgPrefix, tc.wantrcode, rc)
		}
		if len(w.Msg.Answer) == 0 && tc.wantAnswer != nil {
			t.Errorf("%v an answer, but got none", errorMsgPrefix)
			continue
		}
		if tc.wantAnswer == nil {
			if len(w.Msg.Answer) > 0 {
				t.Errorf("%v no answer, but got %v", errorMsgPrefix, w.Msg.Answer[0])
			}
			continue
		}
		if w.Msg.Answer[0].Header().Ttl != 0 {
			t.Errorf("%v TTL to be 0, but got %v", errorMsgPrefix, w.Msg.Answer[0].Header().Ttl)
		}
		if w.Msg.Answer[0].Header().Name != tc.qname {
			t.Errorf("%v Name to be %s, but got %s", errorMsgPrefix, tc.qname, w.Msg.Answer[0].Header().Name)
		}
		if w.Msg.Answer[0].Header().Rrtype != tc.qtype {
			t.Errorf("%v Type to be %d, but got %d", errorMsgPrefix, tc.qtype, w.Msg.Answer[0].Header().Rrtype)
		}
		if w.Msg.Answer[0].Header().Ttl != tc.wantTTL {
			t.Errorf("%v TTL to be %d, but got %d", errorMsgPrefix, tc.wantTTL, w.Msg.Answer[0].Header().Ttl)
		}
		switch tc.qtype {
		case dns.TypeA:
			if w.Msg.Answer[0].(*dns.A).A.String() != tc.wantAnswer[0] {
				t.Errorf("%v answer %s, but got %s", errorMsgPrefix, tc.wantAnswer[0], w.Msg.Answer[0].(*dns.A).A.String())
			}
		}
	}
}