package grpctool

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/grpctool/test"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/matcher"
	"google.golang.org/grpc/encoding"
	"google.golang.org/grpc/mem"
)

var (
	_ encoding.CodecV2 = RawCodec{}
	_ encoding.CodecV2 = RawCodecWithProtoFallback{}
)

func TestRawCodec_Roundtrip(t *testing.T) {
	input := &RawFrame{Data: mem.BufferSlice{mem.SliceBuffer([]byte{1, 2, 3, 4, 5})}}
	serialized, err := RawCodec{}.Marshal(input)
	require.NoError(t, err)
	output := &RawFrame{}
	err = RawCodec{}.Unmarshal(serialized, output)
	require.NoError(t, err)
	assert.Equal(t, input.Data.Materialize(), output.Data.Materialize())
}

func TestRawCodec_BadType(t *testing.T) {
	serialized, err := RawCodec{}.Marshal(&test.Request{})
	require.EqualError(t, err, "RawCodec.Marshal(): unexpected source message type: *test.Request")
	assert.Empty(t, serialized)

	output := &test.Request{}
	err = RawCodec{}.Unmarshal(mem.BufferSlice{mem.SliceBuffer([]byte{1, 2, 3, 4, 5})}, output)
	require.EqualError(t, err, "RawCodec.Unmarshal(): unexpected target message type: *test.Request")
}

func TestRawCodec_BufferFreeMarshal(t *testing.T) {
	codecs := []encoding.CodecV2{
		RawCodec{},
		RawCodecWithProtoFallback{},
	}
	for _, codec := range codecs {
		t.Run(codec.Name(), func(t *testing.T) {
			pool := &testPool{
				allocated: make(map[*[]byte]struct{}),
			}
			bufferSlice := pool.Get(128 * 1024)
			buffer := mem.NewBuffer(bufferSlice, pool)
			input := &RawFrame{
				Data: mem.BufferSlice{buffer},
			}

			serialized, err := codec.Marshal(input)
			buffer.Free() // we used it for marshaling, don't need it anymore.
			require.NoError(t, err)

			// use serialized

			serialized.Free() // no longer need it, free

			assert.Empty(t, pool.allocated)
		})
	}
}

func TestRawCodec_BufferFreeUnmarshal(t *testing.T) {
	codecs := []encoding.CodecV2{
		RawCodec{},
		RawCodecWithProtoFallback{},
	}
	for _, codec := range codecs {
		t.Run(codec.Name(), func(t *testing.T) {
			pool := &testPool{
				allocated: make(map[*[]byte]struct{}),
			}
			bufferSlice := pool.Get(128 * 1024)
			buffer := mem.NewBuffer(bufferSlice, pool)
			out := &RawFrame{}

			err := RawCodec{}.Unmarshal(mem.BufferSlice{buffer}, out)
			buffer.Free() // we used it for unmarshaling, don't need it anymore.
			require.NoError(t, err)

			// use out

			out.Data.Free() // no longer need it, free

			assert.Empty(t, pool.allocated)
		})
	}
}

func TestRawCodecWithProtoFallback_RoundtripRaw(t *testing.T) {
	input := &RawFrame{Data: mem.BufferSlice{mem.SliceBuffer([]byte{1, 2, 3, 4, 5})}}
	serialized, err := RawCodecWithProtoFallback{}.Marshal(input)
	require.NoError(t, err)
	output := &RawFrame{}
	err = RawCodecWithProtoFallback{}.Unmarshal(serialized, output)
	require.NoError(t, err)
	assert.Equal(t, input.Data.Materialize(), output.Data.Materialize())
}

func TestRawCodecWithProtoFallback_RoundtripNonRaw(t *testing.T) {
	input := &test.Request{S1: "bla"}
	serialized, err := RawCodecWithProtoFallback{}.Marshal(input)
	require.NoError(t, err)
	output := &test.Request{}
	err = RawCodecWithProtoFallback{}.Unmarshal(serialized, output)
	require.NoError(t, err)
	matcher.AssertProtoEqual(t, input, output)
}

var (
	_ mem.BufferPool = (*testPool)(nil)
)

type testPool struct {
	allocated map[*[]byte]struct{}
}

func (t *testPool) Get(length int) *[]byte {
	buf := make([]byte, length)
	t.allocated[&buf] = struct{}{}
	return &buf
}

func (t *testPool) Put(buf *[]byte) {
	if _, ok := t.allocated[buf]; !ok {
		panic("unexpected put")
	}
	delete(t.allocated, buf)
}
