package watch_aggregator //nolint:staticcheck

import (
	"bufio"
	"bytes"
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"testing"

	"buf.build/go/protovalidate"
	"github.com/coder/websocket"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/httpz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_modshared"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_stdlib"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/testlogger"
	"go.uber.org/mock/gomock"
	"k8s.io/apimachinery/pkg/util/wait"
	k8swatch "k8s.io/apimachinery/pkg/watch"
)

func TestWatchAggregator_RejectInvalidProtocol(t *testing.T) {
	// GIVEN
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	mockWatchFactory := NewMockwatchFactory(ctrl)
	mockResponseWriter := mock_stdlib.NewMockResponseWriterFlusher(ctrl)
	responseHeaders := http.Header{}
	mockConn := mock_stdlib.NewMockConn(ctrl)
	v, err := protovalidate.New()
	require.NoError(t, err)

	gomock.InOrder(
		mockResponseWriter.EXPECT().Header().Return(responseHeaders).AnyTimes(),
		mockResponseWriter.EXPECT().WriteHeader(http.StatusSwitchingProtocols),
		mockResponseWriter.EXPECT().Hijack().Return(mockConn, bufio.NewReadWriter(bufio.NewReader(mockConn), bufio.NewWriter(mockConn)), nil),
		mockResponseWriter.EXPECT().Header().Return(responseHeaders).AnyTimes(),
		mockConn.EXPECT().Write(matchClose(websocket.StatusProtocolError, fmt.Sprintf("protocol unspecified, please use the %q protocol", acceptedSubprotocolName))),
		mockConn.EXPECT().Close().AnyTimes(),
	)

	ctx := t.Context()
	r := createWatchRequest(t, ctx)
	r.Header.Set("Sec-Websocket-Protocol", "invalid-protocol")

	// WHEN
	wa := newWatchAggregator(mockWatchFactory, testlogger.New(t), mockAPI, v, func() (WebSocketInterface, error) {
		return WebSocketAccept(mockResponseWriter, r)
	})

	// THEN
	assert.Nil(t, wa)
}

func TestWatchAggregator_AcceptWebSocket(t *testing.T) {
	// GIVEN
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	mockWatchFactory := NewMockwatchFactory(ctrl)
	mockResponseWriter := mock_stdlib.NewMockResponseWriterFlusher(ctrl)
	responseHeaders := http.Header{}
	mockConn := mock_stdlib.NewMockConn(ctrl)
	v, err := protovalidate.New()
	require.NoError(t, err)

	gomock.InOrder(
		mockResponseWriter.EXPECT().Header().Return(responseHeaders).AnyTimes(),
		mockResponseWriter.EXPECT().WriteHeader(http.StatusSwitchingProtocols),
		mockResponseWriter.EXPECT().Hijack().Return(mockConn, bufio.NewReadWriter(bufio.NewReader(mockConn), bufio.NewWriter(mockConn)), nil),
		mockResponseWriter.EXPECT().Header().Return(responseHeaders).AnyTimes(),
	)

	ctx := t.Context()
	r := createWatchRequest(t, ctx)

	// WHEN
	wa := newWatchAggregator(mockWatchFactory, testlogger.New(t), mockAPI, v, func() (WebSocketInterface, error) {
		return WebSocketAccept(mockResponseWriter, r)
	})

	// THEN
	assert.NotNil(t, wa)
}

func TestWatchAggregator_ShutdownWithNoWatches(t *testing.T) {
	// GIVEN
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	mockWatchFactory := NewMockwatchFactory(ctrl)
	mockWebSocket := NewMockWebSocketInterface(ctrl)
	accept := func() (WebSocketInterface, error) { return mockWebSocket, nil }
	v, err := protovalidate.New()
	require.NoError(t, err)

	gomock.InOrder(
		mockWebSocket.EXPECT().Reader(gomock.Any()).DoAndReturn(func(ctx context.Context) (websocket.MessageType, io.Reader, error) {
			<-ctx.Done()
			return websocket.MessageText, nil, &websocket.CloseError{Code: websocket.StatusNormalClosure, Reason: ""}
		}),
		mockWebSocket.EXPECT().Close(gomock.Any(), gomock.Any()),
	)

	// WHEN
	wa := newWatchAggregator(mockWatchFactory, testlogger.New(t), mockAPI, v, accept)
	require.NotNil(t, wa)

	var wg wait.Group
	defer wg.Wait()

	auxCtx, auxCancel := context.WithCancel(t.Context())
	defer auxCancel()

	wg.StartWithContext(auxCtx, wa.Handle)
}

func TestWatchAggregator_StartSingleWatch(t *testing.T) {
	// GIVEN
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	mockWatchFactory := NewMockwatchFactory(ctrl)
	mockWebSocket := NewMockWebSocketInterface(ctrl)
	accept := func() (WebSocketInterface, error) { return mockWebSocket, nil }
	v, err := protovalidate.New()
	require.NoError(t, err)

	gomock.InOrder(
		// Start the watch
		mockWebSocket.EXPECT().Reader(gomock.Any()).DoAndReturn(func(ctx context.Context) (websocket.MessageType, io.Reader, error) {
			mockReader := mock_stdlib.NewMockReader(ctrl)
			mockReader.EXPECT().Read(gomock.Any()).DoAndReturn(func(b []byte) (int, error) {
				d := []byte(`{"type": "watch", "watch_id": "any", "watch_params": {"version": "v1", "resource": "pods"}}`)
				copy(b, d)
				return len(d), io.EOF
			})
			return websocket.MessageText, mockReader, nil
		}),
		// Shutdown
		mockWebSocket.EXPECT().Reader(gomock.Any()).DoAndReturn(func(ctx context.Context) (websocket.MessageType, io.Reader, error) {
			<-ctx.Done()
			return websocket.MessageText, nil, &websocket.CloseError{Code: websocket.StatusNormalClosure, Reason: ""}
		}),
		mockWebSocket.EXPECT().Close(gomock.Any(), gomock.Any()),
	)

	auxCtx, auxCancel := context.WithCancel(context.Background())
	defer auxCancel()

	mockWatchFactory.EXPECT().
		new(gomock.Any(), gomock.Any()).
		DoAndReturn(func(wr *WatchRequest, wrs watchResponseSender) watchInterface {
			mockWatch := NewMockwatchInterface(ctrl)
			mockWatch.EXPECT().
				handle(gomock.Any()).
				Do(func(ctx context.Context) {
					// we started the watch, so we can trigger the shutdown for the test
					auxCancel()
				})
			return mockWatch
		})

	// WHEN
	wa := newWatchAggregator(mockWatchFactory, testlogger.New(t), mockAPI, v, accept)
	require.NotNil(t, wa)

	var wg wait.Group
	defer wg.Wait()

	wg.StartWithContext(auxCtx, wa.Handle)
}

func TestWatchAggregator_ForwardWatchEvent(t *testing.T) {
	// GIVEN
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	mockWatchFactory := NewMockwatchFactory(ctrl)
	mockWebSocket := NewMockWebSocketInterface(ctrl)
	accept := func() (WebSocketInterface, error) { return mockWebSocket, nil }
	v, err := protovalidate.New()
	require.NoError(t, err)

	gomock.InOrder(
		// Start the watch
		mockWebSocket.EXPECT().Reader(gomock.Any()).DoAndReturn(func(ctx context.Context) (websocket.MessageType, io.Reader, error) {
			mockReader := mock_stdlib.NewMockReader(ctrl)
			mockReader.EXPECT().Read(gomock.Any()).DoAndReturn(func(b []byte) (int, error) {
				d := []byte(`{"type": "watch", "watch_id": "any", "watch_params": {"version": "v1", "resource": "pods"}}`)
				copy(b, d)
				return len(d), io.EOF
			})
			return websocket.MessageText, mockReader, nil
		}),
		// Shutdown
		mockWebSocket.EXPECT().Reader(gomock.Any()).DoAndReturn(func(ctx context.Context) (websocket.MessageType, io.Reader, error) {
			<-ctx.Done()
			return websocket.MessageText, nil, &websocket.CloseError{Code: websocket.StatusNormalClosure, Reason: ""}
		}),
	)

	auxCtx, auxCancel := context.WithCancel(context.Background())
	defer auxCancel()

	gomock.InOrder(
		mockWebSocket.EXPECT().Writer(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, mt websocket.MessageType) (io.WriteCloser, error) {
			mockWriter := mock_stdlib.NewMockWriteCloser(ctrl)
			mockWriter.EXPECT().Write(gomock.Any()).DoAndReturn(func(b []byte) (int, error) {
				expectedWatchEvent := &watchEvent{
					WatchID: "any",
					Type:    "event",
					Event: &k8sWatchEvent{
						Type:   k8swatch.Added,
						Object: nil,
					},
				}
				expected, err := json.Marshal(expectedWatchEvent)
				require.NoError(t, err)

				assert.Equal(t, expected, b)
				return len(expected), nil
			})
			mockWriter.EXPECT().Close().DoAndReturn(func() error {
				// we can now shutdown the test
				defer auxCancel()
				return nil
			})
			return mockWriter, nil
		}),
		mockWebSocket.EXPECT().Close(gomock.Any(), gomock.Any()),
	)

	mockWatchFactory.EXPECT().
		new(gomock.Any(), gomock.Any()).
		DoAndReturn(func(wr *WatchRequest, wrs watchResponseSender) watchInterface {
			mockWatch := NewMockwatchInterface(ctrl)
			mockWatch.EXPECT().
				handle(gomock.Any()).
				Do(func(ctx context.Context) {
					wrs.sendEvent("any", &k8sWatchEvent{
						Type:   k8swatch.Added,
						Object: nil,
					})
					<-ctx.Done()
				})
			return mockWatch
		})

	// WHEN
	wa := newWatchAggregator(mockWatchFactory, testlogger.New(t), mockAPI, v, accept)
	require.NotNil(t, wa)

	var wg wait.Group
	defer wg.Wait()

	wg.StartWithContext(auxCtx, wa.Handle)
}

func TestWatchAggregator_Unwatch(t *testing.T) {
	// GIVEN
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	mockWatchFactory := NewMockwatchFactory(ctrl)
	mockWebSocket := NewMockWebSocketInterface(ctrl)
	accept := func() (WebSocketInterface, error) { return mockWebSocket, nil }
	v, err := protovalidate.New()
	require.NoError(t, err)

	waitUnwatch := make(chan struct{})

	gomock.InOrder(
		// Start the watch
		mockWebSocket.EXPECT().Reader(gomock.Any()).DoAndReturn(func(ctx context.Context) (websocket.MessageType, io.Reader, error) {
			mockReader := mock_stdlib.NewMockReader(ctrl)
			mockReader.EXPECT().Read(gomock.Any()).DoAndReturn(func(b []byte) (int, error) {
				d := []byte(`{"type": "watch", "watch_id": "any", "watch_params": {"version": "v1", "resource": "pods"}}`)
				copy(b, d)
				return len(d), io.EOF
			})
			return websocket.MessageText, mockReader, nil
		}),
		// Stop the watch
		mockWebSocket.EXPECT().Reader(gomock.Any()).DoAndReturn(func(ctx context.Context) (websocket.MessageType, io.Reader, error) {
			<-waitUnwatch
			mockReader := mock_stdlib.NewMockReader(ctrl)
			mockReader.EXPECT().Read(gomock.Any()).DoAndReturn(func(b []byte) (int, error) {
				d := []byte(`{"type": "unwatch", "watch_id": "any"}`)
				copy(b, d)
				return len(d), io.EOF
			})
			return websocket.MessageText, mockReader, nil
		}),
		// Shutdown
		mockWebSocket.EXPECT().Reader(gomock.Any()).DoAndReturn(func(ctx context.Context) (websocket.MessageType, io.Reader, error) {
			<-ctx.Done()
			return websocket.MessageText, nil, &websocket.CloseError{Code: websocket.StatusNormalClosure, Reason: ""}
		}),
	)

	auxCtx, auxCancel := context.WithCancel(context.Background())
	defer auxCancel()

	gomock.InOrder(
		mockWebSocket.EXPECT().Writer(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, mt websocket.MessageType) (io.WriteCloser, error) {
			mockWriter := mock_stdlib.NewMockWriteCloser(ctrl)
			mockWriter.EXPECT().Write(gomock.Any()).DoAndReturn(func(b []byte) (int, error) {
				expectedWatchEvent := &watchEvent{
					WatchID: "any",
					Type:    "stop",
				}
				expected, err := json.Marshal(expectedWatchEvent)
				require.NoError(t, err)

				assert.Equal(t, expected, b)
				return len(expected), nil
			})
			mockWriter.EXPECT().Close().DoAndReturn(func() error {
				// we can now shutdown the test
				defer auxCancel()
				return nil
			})
			return mockWriter, nil
		}),
		mockWebSocket.EXPECT().Close(gomock.Any(), gomock.Any()),
	)

	mockWatchFactory.EXPECT().
		new(gomock.Any(), gomock.Any()).
		DoAndReturn(func(wr *WatchRequest, wrs watchResponseSender) watchInterface {
			mockWatch := NewMockwatchInterface(ctrl)
			mockWatch.EXPECT().
				handle(gomock.Any()).
				Do(func(ctx context.Context) {
					close(waitUnwatch)
					<-ctx.Done()
					wrs.sendStop("any")
				})
			return mockWatch
		})

	// WHEN
	wa := newWatchAggregator(mockWatchFactory, testlogger.New(t), mockAPI, v, accept)
	require.NotNil(t, wa)

	var wg wait.Group
	defer wg.Wait()

	wg.StartWithContext(auxCtx, wa.Handle)
}

func TestWatchAggregator_ErrorInvalidWatchType(t *testing.T) {
	// GIVEN
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	mockWatchFactory := NewMockwatchFactory(ctrl)
	mockWebSocket := NewMockWebSocketInterface(ctrl)
	accept := func() (WebSocketInterface, error) { return mockWebSocket, nil }
	v, err := protovalidate.New()
	require.NoError(t, err)

	mockWebSocket.EXPECT().Reader(gomock.Any()).DoAndReturn(func(ctx context.Context) (websocket.MessageType, io.Reader, error) {
		mockReader := mock_stdlib.NewMockReader(ctrl)
		mockReader.EXPECT().Read(gomock.Any()).DoAndReturn(func(b []byte) (int, error) {
			d := []byte(`{"type": "invalid", "watch_id": "any"}`)
			copy(b, d)
			return len(d), io.EOF
		})
		return websocket.MessageText, mockReader, nil
	})

	auxCtx, auxCancel := context.WithCancel(context.Background())
	defer auxCancel()

	gomock.InOrder(
		mockWebSocket.EXPECT().Writer(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, mt websocket.MessageType) (io.WriteCloser, error) {
			mockWriter := mock_stdlib.NewMockWriteCloser(ctrl)
			mockWriter.EXPECT().Write(gomock.Any()).DoAndReturn(func(b []byte) (int, error) {
				expectedWatchEvent := &watchEvent{
					WatchID: "any",
					Type:    "error",
					Error: &watchError{
						Type:    watchRequestFailedErrorType,
						Message: "validation of watch request failed: validation error: type: value must be in list [watch, unwatch]",
					},
				}
				expected, err := json.Marshal(expectedWatchEvent)
				require.NoError(t, err)

				assert.Equal(t, expected, b)
				return len(expected), nil
			})
			mockWriter.EXPECT().Close().DoAndReturn(func() error {
				// we can now shutdown the test
				defer auxCancel()
				return nil
			})
			return mockWriter, nil
		}),
		mockWebSocket.EXPECT().Close(gomock.Any(), gomock.Any()),
	)

	// WHEN
	wa := newWatchAggregator(mockWatchFactory, testlogger.New(t), mockAPI, v, accept)
	require.NotNil(t, wa)

	var wg wait.Group
	defer wg.Wait()

	wg.StartWithContext(auxCtx, wa.Handle)
}

func TestWatchAggregator_ErrorStartWatchSameID(t *testing.T) {
	// GIVEN
	ctrl := gomock.NewController(t)
	mockAPI := mock_modshared.NewMockAPI(ctrl)
	mockWatchFactory := NewMockwatchFactory(ctrl)
	mockWebSocket := NewMockWebSocketInterface(ctrl)
	accept := func() (WebSocketInterface, error) { return mockWebSocket, nil }
	v, err := protovalidate.New()
	require.NoError(t, err)

	gomock.InOrder(
		// Start the watch
		mockWebSocket.EXPECT().Reader(gomock.Any()).DoAndReturn(func(ctx context.Context) (websocket.MessageType, io.Reader, error) {
			mockReader := mock_stdlib.NewMockReader(ctrl)
			mockReader.EXPECT().Read(gomock.Any()).DoAndReturn(func(b []byte) (int, error) {
				d := []byte(`{"type": "watch", "watch_id": "any", "watch_params": {"version": "v1", "resource": "pods"}}`)
				copy(b, d)
				return len(d), io.EOF
			})
			return websocket.MessageText, mockReader, nil
		}),
		// Start watch with same id
		mockWebSocket.EXPECT().Reader(gomock.Any()).DoAndReturn(func(ctx context.Context) (websocket.MessageType, io.Reader, error) {
			mockReader := mock_stdlib.NewMockReader(ctrl)
			mockReader.EXPECT().Read(gomock.Any()).DoAndReturn(func(b []byte) (int, error) {
				d := []byte(`{"type": "watch", "watch_id": "any", "watch_params": {"version": "v1", "resource": "deployments"}}`)
				copy(b, d)
				return len(d), io.EOF
			})
			return websocket.MessageText, mockReader, nil
		}),
		// Shutdown
		mockWebSocket.EXPECT().Reader(gomock.Any()).DoAndReturn(func(ctx context.Context) (websocket.MessageType, io.Reader, error) {
			<-ctx.Done()
			return websocket.MessageText, nil, &websocket.CloseError{Code: websocket.StatusNormalClosure, Reason: ""}
		}),
	)

	auxCtx, auxCancel := context.WithCancel(context.Background())
	defer auxCancel()

	gomock.InOrder(
		mockWebSocket.EXPECT().Writer(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, mt websocket.MessageType) (io.WriteCloser, error) {
			mockWriter := mock_stdlib.NewMockWriteCloser(ctrl)
			mockWriter.EXPECT().Write(gomock.Any()).DoAndReturn(func(b []byte) (int, error) {
				expectedWatchEvent := &watchEvent{
					WatchID: "any",
					Type:    "error",
					Error: &watchError{
						Type:    watchRequestFailedErrorType,
						Message: "a watch with this id is already registered, stop it before starting a new one with the same id",
					},
				}
				expected, err := json.Marshal(expectedWatchEvent)
				require.NoError(t, err)

				assert.Equal(t, expected, b)
				return len(expected), nil
			})
			mockWriter.EXPECT().Close().DoAndReturn(func() error {
				// we can now shutdown the test
				defer auxCancel()
				return nil
			})
			return mockWriter, nil
		}),
		mockWebSocket.EXPECT().Close(gomock.Any(), gomock.Any()),
	)

	mockWatchFactory.EXPECT().
		new(gomock.Any(), gomock.Any()).
		DoAndReturn(func(wr *WatchRequest, wrs watchResponseSender) watchInterface {
			mockWatch := NewMockwatchInterface(ctrl)
			mockWatch.EXPECT().
				handle(gomock.Any()).
				Do(func(ctx context.Context) {
					<-ctx.Done()
				})
			return mockWatch
		})

	// WHEN
	wa := newWatchAggregator(mockWatchFactory, testlogger.New(t), mockAPI, v, accept)
	require.NotNil(t, wa)

	var wg wait.Group
	defer wg.Wait()

	wg.StartWithContext(auxCtx, wa.Handle)
}

func createWatchRequest(t *testing.T, ctx context.Context) *http.Request {
	r, err := http.NewRequestWithContext(ctx, http.MethodGet, "/watch", http.NoBody)
	require.NoError(t, err)
	key := make([]byte, 16)
	_, err = rand.Read(key)
	require.NoError(t, err)
	r.Header.Set(httpz.UpgradeHeader, "websocket")
	r.Header.Set(httpz.ConnectionHeader, "upgrade")
	r.Header.Set("Sec-WebSocket-Version", "13")
	r.Header.Set("Sec-WebSocket-Key", base64.StdEncoding.EncodeToString(key))
	r.Header.Add("Sec-Websocket-Protocol", acceptedSubprotocolName)
	return r
}

const (
	websocketCloseOpCode byte = 8
	websocketFin         byte = 1 << 7
)

type closeMatcher struct {
	statusCode websocket.StatusCode
	reason     string
}

func (m closeMatcher) String() string {
	return string(m.bytes())
}

func (m closeMatcher) bytes() []byte {
	// see https://datatracker.ietf.org/doc/html/rfc6455#section-5.2
	buf := make([]byte, 4+len(m.reason))
	buf[0] = websocketFin | websocketCloseOpCode
	buf[1] = byte(2 + len(m.reason))
	binary.BigEndian.PutUint16(buf[2:], uint16(m.statusCode))
	copy(buf[4:], m.reason)
	return buf
}

func (m closeMatcher) Matches(v any) bool {
	actual := v.([]byte)
	expected := m.bytes()

	return bytes.Equal(actual, expected)
}

func matchClose(statusCode websocket.StatusCode, reason string) gomock.Matcher {
	return closeMatcher{statusCode, reason}
}
