package tcp

import (
	"context"
	"encoding/binary"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"mpbl3p/mocks"
	"mpbl3p/proxy"
	"testing"
)

func TestFlow_Consume(t *testing.T) {
	testContent := []byte("A test string is the content of this packet.")
	testPacket := proxy.SimplePacket(testContent)
	testMac := mocks.AlmostUselessMac("abcd")
	testMac2 := mocks.AlmostUselessMac("efgh")

	t.Run("Length", func(t *testing.T) {
		testConn := mocks.NewMockPerfectBiStreamConn(100)

		flowA := NewFlowConn(context.Background(), testConn.SideA(), []proxy.MacVerifier{testMac}, []proxy.MacGenerator{testMac})

		err := flowA.Consume(context.Background(), testPacket)
		require.Nil(t, err)

		buf := make([]byte, 100)
		n, err := testConn.SideB().Read(buf)
		require.Nil(t, err)

		assert.Equal(t, len(testContent)+4+4, n)
		assert.Equal(t, uint32(len(testContent)+4), binary.LittleEndian.Uint32(buf[:len(buf)-4]))
	})

	t.Run("MultipleGeneratorsLength", func(t *testing.T) {
		testConn := mocks.NewMockPerfectBiStreamConn(100)

		flowA := NewFlowConn(context.Background(), testConn.SideA(), []proxy.MacVerifier{testMac, testMac2}, []proxy.MacGenerator{testMac, testMac2})

		err := flowA.Consume(context.Background(), testPacket)
		require.Nil(t, err)

		buf := make([]byte, 100)
		n, err := testConn.SideB().Read(buf)
		require.Nil(t, err)

		assert.Equal(t, len(testContent)+4+4+4, n)
		assert.Equal(t, uint32(len(testContent)+4+4), binary.LittleEndian.Uint32(buf[:len(buf)-4]))
	})

	t.Run("MultipleGeneratorsOrder", func(t *testing.T) {
		testConn := mocks.NewMockPerfectBiStreamConn(100)

		flowA := NewFlowConn(context.Background(), testConn.SideA(), []proxy.MacVerifier{testMac, testMac2}, []proxy.MacGenerator{testMac, testMac2})

		err := flowA.Consume(context.Background(), testPacket)
		require.Nil(t, err)

		buf := make([]byte, 100)
		n, err := testConn.SideB().Read(buf)
		require.Nil(t, err)

		assert.Equal(t, len(testContent)+4+4+4, n)
		assert.Equal(t, "abcdefgh", string(buf[n-8:n]))
	})
}

func TestFlow_Produce(t *testing.T) {
	testContent := "A test string is the content of this packet."
	testMarshalled := []byte("0000" + testContent + "abcd")
	binary.LittleEndian.PutUint32(testMarshalled, uint32(len(testMarshalled)-4))

	testMac := mocks.AlmostUselessMac("abcd")
	testMac2 := mocks.AlmostUselessMac("efgh")

	t.Run("Length", func(t *testing.T) {
		testConn := mocks.NewMockPerfectBiStreamConn(100)

		flowA := NewFlowConn(context.Background(), testConn.SideA(), []proxy.MacVerifier{testMac}, []proxy.MacGenerator{testMac})

		_, err := testConn.SideB().Write(testMarshalled)
		require.Nil(t, err)

		p, err := flowA.Produce(context.Background())
		require.Nil(t, err)
		assert.Equal(t, len(testContent), len(p.Contents()))
	})

	t.Run("Value", func(t *testing.T) {
		testConn := mocks.NewMockPerfectBiStreamConn(100)

		flowA := NewFlowConn(context.Background(), testConn.SideA(), []proxy.MacVerifier{testMac}, []proxy.MacGenerator{testMac})

		_, err := testConn.SideB().Write(testMarshalled)
		require.Nil(t, err)

		p, err := flowA.Produce(context.Background())
		require.Nil(t, err)
		assert.Equal(t, testContent, string(p.Contents()))
	})

	t.Run("MultipleVerifiersStrip", func(t *testing.T) {
		testContent := "A test string is the content of this packet."
		testMarshalled := []byte("0000" + testContent + "abcdefgh")
		binary.LittleEndian.PutUint32(testMarshalled, uint32(len(testMarshalled)-4))

		testConn := mocks.NewMockPerfectBiStreamConn(100)

		flowA := NewFlowConn(context.Background(), testConn.SideA(), []proxy.MacVerifier{testMac, testMac2}, []proxy.MacGenerator{testMac, testMac2})

		_, err := testConn.SideB().Write(testMarshalled)
		require.Nil(t, err)

		p, err := flowA.Produce(context.Background())
		require.Nil(t, err)
		assert.Equal(t, testContent, string(p.Contents()))
	})
}