package dnstap

import (
	"fmt"
	"net"
	"sync"
	"testing"
	"time"

	"github.com/coredns/coredns/plugin/pkg/reuseport"

	tap "github.com/dnstap/golang-dnstap"
	fs "github.com/farsightsec/golang-framestream"
	"github.com/stretchr/testify/require"
)

var (
	msgType = tap.Dnstap_MESSAGE
	tmsg    = tap.Dnstap{Type: &msgType}
)

type MockLogger struct {
	mu        sync.Mutex
	warnCount int
	warnLog   string
}

func (l *MockLogger) Warningf(format string, v ...any) {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.warnCount++
	l.warnLog += fmt.Sprintf(format, v...)
}

func (l *MockLogger) WarnCount() int {
	l.mu.Lock()
	defer l.mu.Unlock()
	return l.warnCount
}

func (l *MockLogger) WarnLog() string {
	l.mu.Lock()
	defer l.mu.Unlock()
	return l.warnLog
}

func accept(t *testing.T, l net.Listener, count int) {
	t.Helper()
	server, err := l.Accept()
	if err != nil {
		t.Fatalf("Server accepted: %s", err)
	}
	dec, err := fs.NewDecoder(server, &fs.DecoderOptions{
		ContentType:   []byte("protobuf:dnstap.Dnstap"),
		Bidirectional: true,
	})
	if err != nil {
		t.Fatalf("Server decoder: %s", err)
	}

	for range count {
		if _, err := dec.Decode(); err != nil {
			t.Errorf("Server decode: %s", err)
		}
	}

	if err := server.Close(); err != nil {
		t.Error(err)
	}
}

func TestTransport(t *testing.T) {
	transport := [2][2]string{
		{"tcp", ":0"},
		{"unix", "dnstap.sock"},
	}

	for _, param := range transport {
		l, err := reuseport.Listen(param[0], param[1])
		if err != nil {
			t.Fatalf("Cannot start listener: %s", err)
		}

		var wg sync.WaitGroup
		wg.Go(func() {
			accept(t, l, 1)
		})

		dio := newIO(param[0], l.Addr().String(), 1, 1)
		dio.tcpTimeout = 10 * time.Millisecond
		dio.flushTimeout = 30 * time.Millisecond
		dio.errorCheckInterval = 50 * time.Millisecond
		dio.connect()

		dio.Dnstap(&tmsg)

		wg.Wait()
		l.Close()
		dio.close()
	}
}

func TestRace(t *testing.T) {
	count := 10

	l, err := reuseport.Listen("tcp", ":0")
	if err != nil {
		t.Fatalf("Cannot start listener: %s", err)
	}
	defer l.Close()

	var wg sync.WaitGroup
	wg.Go(func() {
		accept(t, l, count)
	})

	dio := newIO("tcp", l.Addr().String(), 1, 1)
	dio.tcpTimeout = 10 * time.Millisecond
	dio.flushTimeout = 30 * time.Millisecond
	dio.errorCheckInterval = 50 * time.Millisecond
	dio.connect()
	defer dio.close()

	wg.Add(count)
	for range count {
		go func() {
			tmsg := tap.Dnstap_MESSAGE
			dio.Dnstap(&tap.Dnstap{Type: &tmsg})
			wg.Done()
		}()
	}
	wg.Wait()
}

func TestReconnect(t *testing.T) {
	t.Run("ConnectedOnStart", func(t *testing.T) {
		// GIVEN
		// 		TCP connection available before DnsTap start up
		//		DnsTap successfully established output connection on start up
		l, err := reuseport.Listen("tcp", ":0")
		if err != nil {
			t.Fatalf("Cannot start listener: %s", err)
		}

		var wg sync.WaitGroup
		wg.Go(func() {
			accept(t, l, 1)
		})

		addr := l.Addr().String()
		logger := MockLogger{}
		dio := newIO("tcp", addr, 1, 1)
		dio.tcpTimeout = 10 * time.Millisecond
		dio.flushTimeout = 30 * time.Millisecond
		dio.errorCheckInterval = 50 * time.Millisecond
		dio.logger = &logger
		dio.connect()
		defer dio.close()

		// WHEN
		//		TCP connection closed when DnsTap is still running
		//		TCP listener starts again on the same port
		//		DnsTap send multiple messages
		dio.Dnstap(&tmsg)
		wg.Wait()

		// Close listener
		l.Close()
		// And start TCP listener again on the same port
		l, err = reuseport.Listen("tcp", addr)
		if err != nil {
			t.Fatalf("Cannot start listener: %s", err)
		}
		defer l.Close()

		wg.Go(func() {
			accept(t, l, 1)
		})

		messageCount := 5
		for range messageCount {
			time.Sleep(100 * time.Millisecond)
			dio.Dnstap(&tmsg)
		}
		wg.Wait()

		// THEN
		//		DnsTap is able to reconnect
		//		Messages can be sent eventually
		require.Eventually(t, func() bool {
			return len(dio.queue) == 0
		}, time.Second, 10*time.Millisecond, "queue should be drained by serve goroutine")
		require.Less(t, logger.WarnCount(), messageCount)
	})

	t.Run("NotConnectedOnStart", func(t *testing.T) {
		// GIVEN
		// 		No TCP connection established at DnsTap start up
		l, err := reuseport.Listen("tcp", ":0")
		if err != nil {
			t.Fatalf("Cannot start listener: %s", err)
		}
		l.Close()

		logger := MockLogger{}
		addr := l.Addr().String()
		dio := newIO("tcp", addr, 1, 1)
		dio.tcpTimeout = 10 * time.Millisecond
		dio.flushTimeout = 30 * time.Millisecond
		dio.errorCheckInterval = 50 * time.Millisecond
		dio.logger = &logger
		dio.connect()
		defer dio.close()

		// WHEN
		//		DnsTap is already running
		//		TCP listener starts on DnsTap's configured port
		//		DnsTap send multiple messages
		dio.Dnstap(&tmsg)

		l, err = reuseport.Listen("tcp", addr)
		if err != nil {
			t.Fatalf("Cannot start listener: %s", err)
		}
		defer l.Close()

		var wg sync.WaitGroup
		wg.Add(1)
		messageCount := 5
		go func() {
			accept(t, l, messageCount)
			wg.Done()
		}()

		for range messageCount {
			time.Sleep(100 * time.Millisecond)
			dio.Dnstap(&tmsg)
		}
		wg.Wait()

		// THEN
		//		DnsTap is able to reconnect
		//		Messages can be sent eventually
		require.Eventually(t, func() bool {
			return len(dio.queue) == 0
		}, time.Second, 10*time.Millisecond, "queue should be drained by serve goroutine")
		require.Less(t, logger.WarnCount(), messageCount)
	})
}

func TestFullQueueWriteFail(t *testing.T) {
	// GIVEN
	// 		DnsTap I/O with a small queue
	l, err := reuseport.Listen("unix", "dn2stap.sock")
	if err != nil {
		t.Fatalf("Cannot start listener: %s", err)
	}
	defer l.Close()

	var wg sync.WaitGroup
	wg.Go(func() {
		accept(t, l, 1)
	})

	logger := MockLogger{}
	dio := newIO("unix", l.Addr().String(), 1, 1)
	dio.flushTimeout = 500 * time.Millisecond
	dio.errorCheckInterval = 50 * time.Millisecond
	dio.logger = &logger
	dio.queue = make(chan *tap.Dnstap, 1)
	dio.connect()
	defer dio.close()

	// WHEN
	//		messages overwhelms the queue
	count := 100
	for range count {
		dio.Dnstap(&tmsg)
	}
	wg.Wait()

	// THEN
	//		Dropped messages are logged
	require.NotEqual(t, 0, logger.WarnCount())
	require.Contains(t, logger.WarnLog(), "Dropped dnstap messages")
}

func TestReconnectClosesPreviousConnection(t *testing.T) {
	// A reconnect (a second dial after a write error) must flush and close the
	// previous encoder and its connection before installing the new one.
	// Otherwise the old socket leaks on every redial.
	l, err := reuseport.Listen("tcp", ":0")
	if err != nil {
		t.Fatalf("Cannot start listener: %s", err)
	}
	defer l.Close()

	accepted := make(chan struct{}, 2)
	firstClosed := make(chan struct{})

	go func() {
		n := 0
		for {
			server, err := l.Accept()
			if err != nil {
				return
			}
			n++
			go func(server net.Conn, n int) {
				dec, err := fs.NewDecoder(server, &fs.DecoderOptions{
					ContentType:   []byte("protobuf:dnstap.Dnstap"),
					Bidirectional: true,
				})
				if err != nil {
					server.Close()
					return
				}
				accepted <- struct{}{}
				if n == 1 {
					// Block until the client closes this connection during the
					// second dial; Decode then returns an error (framestream FINISH/EOF).
					for {
						if _, err := dec.Decode(); err != nil {
							break
						}
					}
					server.Close()
					close(firstClosed)
					return
				}
				// n >= 2: leave the connection open for the rest of the test.
			}(server, n)
		}
	}()

	dio := newIO("tcp", l.Addr().String(), 1, 1)
	dio.tcpTimeout = time.Second

	// First dial establishes the initial connection.
	if err := dio.dial(); err != nil {
		t.Fatalf("first dial failed: %s", err)
	}
	<-accepted
	first := dio.enc
	if first == nil {
		t.Fatal("first dial did not set an encoder")
	}

	// Second dial simulates a reconnect: it must replace the encoder.
	if err := dio.dial(); err != nil {
		t.Fatalf("second dial failed: %s", err)
	}
	<-accepted
	if dio.enc == first {
		t.Fatal("second dial did not replace the previous encoder")
	}
	defer dio.enc.close()

	// The previous connection must have been closed by the second dial.
	select {
	case <-firstClosed:
	case <-time.After(2 * time.Second):
		t.Fatal("previous connection was not closed on reconnect (fd/goroutine leak)")
	}
}
