package log

import (
	"bytes"
	"context"
	"io"
	"log"
	"strings"
	"testing"

	"github.com/coredns/coredns/plugin"
	"github.com/coredns/coredns/plugin/pkg/dnstest"
	clog "github.com/coredns/coredns/plugin/pkg/log"
	"github.com/coredns/coredns/plugin/pkg/replacer"
	"github.com/coredns/coredns/plugin/pkg/response"
	"github.com/coredns/coredns/plugin/test"

	"github.com/miekg/dns"
)

func init() { clog.Discard() }

func TestLoggedStatus(t *testing.T) {
	rule := Rule{
		NameScope: ".",
		Format:    DefaultLogFormat,
		Class:     map[response.Class]struct{}{response.All: {}},
	}

	var f bytes.Buffer
	log.SetOutput(&f)

	logger := Logger{
		Rules: []Rule{rule},
		Next:  test.ErrorHandler(),
		repl:  replacer.New(),
	}

	ctx := context.TODO()
	r := new(dns.Msg)
	r.SetQuestion("example.org.", dns.TypeA)

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

	rcode, _ := logger.ServeDNS(ctx, rec, r)
	if rcode != 2 {
		t.Errorf("Expected rcode to be 2 - was: %d", rcode)
	}

	logged := f.String()
	if !strings.Contains(logged, "A IN example.org. udp 29 false 512") {
		t.Errorf("Expected it to be logged. Logged string: %s", logged)
	}
}

func TestLoggedClassDenial(t *testing.T) {
	rule := Rule{
		NameScope: ".",
		Format:    DefaultLogFormat,
		Class:     map[response.Class]struct{}{response.Denial: {}},
	}

	var f bytes.Buffer
	log.SetOutput(&f)

	logger := Logger{
		Rules: []Rule{rule},
		Next:  test.ErrorHandler(),
		repl:  replacer.New(),
	}

	ctx := context.TODO()
	r := new(dns.Msg)
	r.SetQuestion("example.org.", dns.TypeA)

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

	logger.ServeDNS(ctx, rec, r)

	logged := f.String()
	if len(logged) != 0 {
		t.Errorf("Expected it not to be logged, but got string: %s", logged)
	}
}

func TestLoggedClassError(t *testing.T) {
	rule := Rule{
		NameScope: ".",
		Format:    DefaultLogFormat,
		Class:     map[response.Class]struct{}{response.Error: {}},
	}

	var f bytes.Buffer
	log.SetOutput(&f)

	logger := Logger{
		Rules: []Rule{rule},
		Next:  test.ErrorHandler(),
		repl:  replacer.New(),
	}

	ctx := context.TODO()
	r := new(dns.Msg)
	r.SetQuestion("example.org.", dns.TypeA)

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

	logger.ServeDNS(ctx, rec, r)

	logged := f.String()
	if !strings.Contains(logged, "SERVFAIL") {
		t.Errorf("Expected it to be logged. Logged string: %s", logged)
	}
}

func TestLoggedSynthesizesDeferredServerFailure(t *testing.T) {
	rule := Rule{
		NameScope: ".",
		Format:    DefaultLogFormat,
		Class:     map[response.Class]struct{}{response.All: {}},
	}

	var f bytes.Buffer
	log.SetOutput(&f)

	logger := Logger{
		Rules: []Rule{rule},
		repl:  replacer.New(),
	}

	ctx := context.TODO()
	r := new(dns.Msg)
	r.SetQuestion("example.org.", dns.TypeA)

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

	rcode, err := logger.ServeDNS(ctx, rec, r)
	if err == nil {
		t.Fatal("expected no next plugin error")
	}
	if rcode != dns.RcodeServerFailure {
		t.Fatalf("expected SERVFAIL rcode, got %d", rcode)
	}

	logged := f.String()
	if !strings.Contains(logged, "SERVFAIL") {
		t.Fatalf("expected SERVFAIL in log output, got %q", logged)
	}
	if strings.Contains(logged, "\"OTHERERROR ") {
		t.Fatalf("expected synthesized server error classification, got %q", logged)
	}
}

func TestLoggedSynthesizesDeferredRefused(t *testing.T) {
	rule := Rule{
		NameScope: ".",
		Format:    DefaultLogFormat,
		Class:     map[response.Class]struct{}{response.All: {}},
	}

	var f bytes.Buffer
	log.SetOutput(&f)

	logger := Logger{
		Rules: []Rule{rule},
		Next:  test.NextHandler(dns.RcodeRefused, nil),
		repl:  replacer.New(),
	}

	ctx := context.TODO()
	r := new(dns.Msg)
	r.SetQuestion("example.org.", dns.TypeA)

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

	rcode, err := logger.ServeDNS(ctx, rec, r)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if rcode != dns.RcodeRefused {
		t.Fatalf("expected REFUSED rcode, got %d", rcode)
	}

	logged := f.String()
	if !strings.Contains(logged, "REFUSED") {
		t.Fatalf("expected REFUSED in log output, got %q", logged)
	}
}

func TestLoggedClassNameErrorWithoutSOA(t *testing.T) {
	tests := []struct {
		name      string
		class     response.Class
		shouldLog bool
	}{
		{name: "error", class: response.Error, shouldLog: false},
		{name: "denial", class: response.Denial, shouldLog: true},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			rule := Rule{
				NameScope: ".",
				Format:    DefaultLogFormat,
				Class:     map[response.Class]struct{}{tc.class: {}},
			}

			var f bytes.Buffer
			log.SetOutput(&f)

			logger := Logger{
				Rules: []Rule{rule},
				Next: plugin.HandlerFunc(func(_ context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
					m := new(dns.Msg)
					m.SetRcode(r, dns.RcodeNameError)
					return dns.RcodeNameError, w.WriteMsg(m)
				}),
				repl: replacer.New(),
			}

			ctx := context.TODO()
			r := new(dns.Msg)
			r.SetQuestion("1.1.168.192.in-addr.arpa.", dns.TypePTR)

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

			_, err := logger.ServeDNS(ctx, rec, r)
			if err != nil {
				t.Fatal(err)
			}

			logged := f.String()
			if !tc.shouldLog && len(logged) != 0 {
				t.Errorf("Expected it not to be logged, but got string: %s", logged)
			}
			if tc.shouldLog && !strings.Contains(logged, "NXDOMAIN") {
				t.Errorf("Expected it to be logged as NXDOMAIN. Logged string: %s", logged)
			}
		})
	}
}

func TestLogged(t *testing.T) {
	tests := []struct {
		Rules           []Rule
		Domain          string
		ShouldLog       bool
		ShouldString    string
		ShouldNOTString string // for test format
	}{
		// case for NameScope
		{
			Rules: []Rule{
				{
					NameScope: "example.org.",
					Format:    DefaultLogFormat,
					Class:     map[response.Class]struct{}{response.All: {}},
				},
			},
			Domain:       "example.org.",
			ShouldLog:    true,
			ShouldString: "A IN example.org.",
		},
		{
			Rules: []Rule{
				{
					NameScope: "example.org.",
					Format:    DefaultLogFormat,
					Class:     map[response.Class]struct{}{response.All: {}},
				},
			},
			Domain:       "example.net.",
			ShouldLog:    false,
			ShouldString: "",
		},
		{
			Rules: []Rule{
				{
					NameScope: "example.org.",
					Format:    DefaultLogFormat,
					Class:     map[response.Class]struct{}{response.All: {}},
				},
				{
					NameScope: "example.net.",
					Format:    DefaultLogFormat,
					Class:     map[response.Class]struct{}{response.All: {}},
				},
			},
			Domain:       "example.net.",
			ShouldLog:    true,
			ShouldString: "A IN example.net.",
		},

		// case for format
		{
			Rules: []Rule{
				{
					NameScope: ".",
					Format:    "{type} {class}",
					Class:     map[response.Class]struct{}{response.All: {}},
				},
			},
			Domain:          "example.org.",
			ShouldLog:       true,
			ShouldString:    "A IN",
			ShouldNOTString: "example.org",
		},
		{
			Rules: []Rule{
				{
					NameScope: ".",
					Format:    "{remote}:{port}",
					Class:     map[response.Class]struct{}{response.All: {}},
				},
			},
			Domain:          "example.org.",
			ShouldLog:       true,
			ShouldString:    "10.240.0.1:40212",
			ShouldNOTString: "A IN example.org",
		},
		{
			Rules: []Rule{
				{
					NameScope: ".",
					Format:    CombinedLogFormat,
					Class:     map[response.Class]struct{}{response.All: {}},
				},
			},
			Domain:       "example.org.",
			ShouldLog:    true,
			ShouldString: "\"0\"",
		},
		{
			Rules: []Rule{
				{
					NameScope: ".",
					Format:    CombinedLogFormat,
					Class:     map[response.Class]struct{}{response.All: {}},
				},
			},
			Domain:          "foo.%s.example.org.",
			ShouldLog:       true,
			ShouldString:    "foo.%s.example.org.",
			ShouldNOTString: "%!s(MISSING)",
		},
	}

	for _, tc := range tests {
		var f bytes.Buffer
		log.SetOutput(&f)

		logger := Logger{
			Rules: tc.Rules,
			Next:  test.ErrorHandler(),
			repl:  replacer.New(),
		}

		ctx := context.TODO()
		r := new(dns.Msg)
		r.SetQuestion(tc.Domain, dns.TypeA)

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

		_, err := logger.ServeDNS(ctx, rec, r)
		if err != nil {
			t.Fatal(err)
		}

		logged := f.String()

		if !tc.ShouldLog && len(logged) != 0 {
			t.Errorf("Expected it not to be logged, but got string: %s", logged)
		}
		if tc.ShouldLog && !strings.Contains(logged, tc.ShouldString) {
			t.Errorf("Expected it to contains: %s. Logged string: %s", tc.ShouldString, logged)
		}
		if tc.ShouldLog && tc.ShouldNOTString != "" && strings.Contains(logged, tc.ShouldNOTString) {
			t.Errorf("Expected it to NOT contains: %s. Logged string: %s", tc.ShouldNOTString, logged)
		}
	}
}

func BenchmarkLogged(b *testing.B) {
	log.SetOutput(io.Discard)

	rule := Rule{
		NameScope: ".",
		Format:    DefaultLogFormat,
		Class:     map[response.Class]struct{}{response.All: {}},
	}

	logger := Logger{
		Rules: []Rule{rule},
		Next:  test.ErrorHandler(),
		repl:  replacer.New(),
	}

	ctx := context.TODO()
	r := new(dns.Msg)
	r.SetQuestion("example.org.", dns.TypeA)

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

	for b.Loop() {
		logger.ServeDNS(ctx, rec, r)
	}
}
