/*
 * Copyright 2021 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package codec

import (
	"context"
	"encoding/binary"
	"net"
	"testing"

	"github.com/bytedance/gopkg/cloud/metainfo"
	"github.com/cloudwego/gopkg/protocol/ttheader"

	"github.com/cloudwego/kitex/internal/mocks"
	mocksremote "github.com/cloudwego/kitex/internal/mocks/remote"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/discovery"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/codec/perrors"
	"github.com/cloudwego/kitex/pkg/remote/transmeta"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/rpcinfo/remoteinfo"
	tm "github.com/cloudwego/kitex/pkg/transmeta"
	"github.com/cloudwego/kitex/transport"
)

var mockPayloadLen = 100

func TestTTHeaderCodec(t *testing.T) {
	for _, tb := range transportBuffers {
		t.Run(tb.Name, func(t *testing.T) {
			ctx := context.Background()
			sendMsg := initClientSendMsg(transport.TTHeader)

			// encode
			buf := tb.NewBuffer()
			totalLenField, err := ttHeaderCodec.encode(ctx, sendMsg, buf)
			binary.BigEndian.PutUint32(totalLenField, uint32(buf.WrittenLen()-Size32+mockPayloadLen))
			test.Assert(t, err == nil, err)
			buf.Flush()

			// decode
			ctx, recvMsg := initServerRecvMsg(ctx)
			test.Assert(t, err == nil, err)
			err = ttHeaderCodec.decode(ctx, recvMsg, buf)
			test.Assert(t, err == nil, err)
			test.Assert(t, recvMsg.PayloadLen() == mockPayloadLen, recvMsg.PayloadLen())
		})
	}
}

func TestTTHeaderCodecWithTransInfo(t *testing.T) {
	for _, tb := range transportBuffers {
		t.Run(tb.Name, func(t *testing.T) {
			ctx := context.Background()
			intKVInfo := prepareIntKVInfo()
			strKVInfo := prepareStrKVInfo()
			sendMsg := initClientSendMsg(transport.TTHeader)
			sendMsg.TransInfo().PutTransIntInfo(intKVInfo)
			sendMsg.TransInfo().PutTransStrInfo(strKVInfo)
			sendMsg.Tags()[HeaderFlagsKey] = ttheader.HeaderFlagSupportOutOfOrder

			// encode
			buf := tb.NewBuffer()
			totalLenField, err := ttHeaderCodec.encode(ctx, sendMsg, buf)
			binary.BigEndian.PutUint32(totalLenField, uint32(buf.WrittenLen()-Size32+mockPayloadLen))
			test.Assert(t, err == nil, err)
			buf.Flush()

			// decode
			ctx, recvMsg := initServerRecvMsg(ctx)
			err = ttHeaderCodec.decode(ctx, recvMsg, buf)
			test.Assert(t, err == nil, err)
			test.Assert(t, recvMsg.PayloadLen() == mockPayloadLen, recvMsg.PayloadLen())

			intKVInfoRecv := recvMsg.TransInfo().TransIntInfo()
			strKVInfoRecv := recvMsg.TransInfo().TransStrInfo()
			test.DeepEqual(t, intKVInfoRecv, intKVInfo)
			test.DeepEqual(t, strKVInfoRecv, strKVInfo)
			flag := recvMsg.Tags()[HeaderFlagsKey]
			test.Assert(t, flag != nil)
			test.Assert(t, flag == ttheader.HeaderFlagSupportOutOfOrder)
		})
	}
}

func TestTTHeaderCodecWithTransInfoWithGDPRToken(t *testing.T) {
	for _, tb := range transportBuffers {
		t.Run(tb.Name, func(t *testing.T) {
			ctx := context.Background()
			intKVInfo := prepareIntKVInfo()
			strKVInfo := prepareStrKVInfoWithGDPRToken()
			sendMsg := initClientSendMsg(transport.TTHeader)
			sendMsg.TransInfo().PutTransIntInfo(intKVInfo)
			sendMsg.TransInfo().PutTransStrInfo(strKVInfo)
			sendMsg.Tags()[HeaderFlagsKey] = ttheader.HeaderFlagSupportOutOfOrder

			// encode
			buf := tb.NewBuffer()
			totalLenField, err := ttHeaderCodec.encode(ctx, sendMsg, buf)
			binary.BigEndian.PutUint32(totalLenField, uint32(buf.WrittenLen()-Size32+mockPayloadLen))
			test.Assert(t, err == nil, err)
			buf.Flush()

			// decode
			ctx, recvMsg := initServerRecvMsg(ctx)
			err = ttHeaderCodec.decode(ctx, recvMsg, buf)
			test.Assert(t, err == nil, err)
			test.Assert(t, recvMsg.PayloadLen() == mockPayloadLen, recvMsg.PayloadLen())

			intKVInfoRecv := recvMsg.TransInfo().TransIntInfo()
			strKVInfoRecv := recvMsg.TransInfo().TransStrInfo()
			test.DeepEqual(t, intKVInfoRecv, intKVInfo)
			test.DeepEqual(t, strKVInfoRecv, strKVInfo)
			flag := recvMsg.Tags()[HeaderFlagsKey]
			test.Assert(t, flag != nil)
			test.Assert(t, flag == ttheader.HeaderFlagSupportOutOfOrder)
		})
	}
}

func TestTTHeaderCodecWithTransInfoFromMetaInfoGDPRToken(t *testing.T) {
	for _, tb := range transportBuffers {
		t.Run(tb.Name, func(t *testing.T) {
			ctx := context.Background()
			intKVInfo := prepareIntKVInfo()
			ctx = metainfo.WithValue(ctx, "gdpr-token", "test token")
			sendMsg := initClientSendMsg(transport.TTHeader)
			sendMsg.TransInfo().PutTransIntInfo(intKVInfo)
			ctx, err := tm.MetainfoClientHandler.WriteMeta(ctx, sendMsg)
			test.Assert(t, err == nil)
			sendMsg.Tags()[HeaderFlagsKey] = ttheader.HeaderFlagSupportOutOfOrder

			// encode
			buf := tb.NewBuffer()
			totalLenField, err := ttHeaderCodec.encode(ctx, sendMsg, buf)
			binary.BigEndian.PutUint32(totalLenField, uint32(buf.WrittenLen()-Size32+mockPayloadLen))
			test.Assert(t, err == nil, err)
			buf.Flush()

			// decode
			ctx, recvMsg := initServerRecvMsg(ctx)
			err = ttHeaderCodec.decode(ctx, recvMsg, buf)
			test.Assert(t, err == nil, err)
			test.Assert(t, recvMsg.PayloadLen() == mockPayloadLen, recvMsg.PayloadLen())

			intKVInfoRecv := recvMsg.TransInfo().TransIntInfo()
			strKVInfoRecv := recvMsg.TransInfo().TransStrInfo()
			test.DeepEqual(t, intKVInfoRecv, intKVInfo)
			test.DeepEqual(t, strKVInfoRecv, map[string]string{transmeta.GDPRToken: "test token"})
			flag := recvMsg.Tags()[HeaderFlagsKey]
			test.Assert(t, flag != nil)
			test.Assert(t, flag == ttheader.HeaderFlagSupportOutOfOrder)
		})
	}
}

func TestFillBasicInfoOfTTHeader(t *testing.T) {
	for _, tb := range transportBuffers {
		ctx := context.Background()
		mockAddr := "mock address"
		// 1. server side fill from address
		// encode
		sendMsg := initClientSendMsg(transport.TTHeader)
		sendMsg.TransInfo().TransStrInfo()[transmeta.HeaderTransRemoteAddr] = mockAddr
		sendMsg.TransInfo().TransIntInfo()[transmeta.FromService] = mockServiceName
		buf := tb.NewBuffer()
		totalLenField, err := ttHeaderCodec.encode(ctx, sendMsg, buf)
		binary.BigEndian.PutUint32(totalLenField, uint32(buf.WrittenLen()-Size32+mockPayloadLen))
		test.Assert(t, err == nil, err)
		buf.Flush()
		// decode
		ctx, recvMsg := initServerRecvMsg(ctx)
		err = ttHeaderCodec.decode(ctx, recvMsg, buf)
		test.Assert(t, err == nil, err)
		test.Assert(t, recvMsg.TransInfo().TransStrInfo()[transmeta.HeaderTransRemoteAddr] == mockAddr)
		test.Assert(t, recvMsg.RPCInfo().From().Address().String() == mockAddr)
		test.Assert(t, recvMsg.RPCInfo().From().ServiceName() == mockServiceName)

		// 2. client side fill to address
		// encode
		sendMsg = initServerSendMsg(transport.TTHeader)
		sendMsg.TransInfo().TransStrInfo()[transmeta.HeaderTransRemoteAddr] = mockAddr
		buf = tb.NewBuffer()
		totalLenField, err = ttHeaderCodec.encode(ctx, sendMsg, buf)
		binary.BigEndian.PutUint32(totalLenField, uint32(buf.WrittenLen()-Size32+mockPayloadLen))
		test.Assert(t, err == nil, err)
		buf.Flush()
		// decode
		recvMsg = initClientRecvMsg()
		toInfo := remoteinfo.AsRemoteInfo(recvMsg.RPCInfo().To())
		toInfo.SetInstance(&mockInst{})
		err = ttHeaderCodec.decode(ctx, recvMsg, buf)
		test.Assert(t, err == nil, err)
		test.Assert(t, recvMsg.TransInfo().TransStrInfo()[transmeta.HeaderTransRemoteAddr] == mockAddr)
		test.Assert(t, toInfo.Address().String() == mockAddr, toInfo.Address())
	}
}

func BenchmarkTTHeaderCodec(b *testing.B) {
	ctx := context.Background()

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		sendMsg := initClientSendMsg(transport.TTHeader)
		// encode
		out := remote.NewWriterBuffer(256)
		totalLenField, err := ttHeaderCodec.encode(ctx, sendMsg, out)
		binary.BigEndian.PutUint32(totalLenField, uint32(out.WrittenLen()-Size32+mockPayloadLen))
		test.Assert(b, err == nil, err)

		// decode
		ctx, recvMsg := initServerRecvMsg(ctx)
		buf, err := out.Bytes()
		test.Assert(b, err == nil, err)
		in := remote.NewReaderBuffer(buf)
		err = ttHeaderCodec.decode(ctx, recvMsg, in)
		test.Assert(b, recvMsg.PayloadLen() == mockPayloadLen, recvMsg.PayloadLen())
		test.Assert(b, err == nil, err)
	}
}

func BenchmarkTTHeaderWithTransInfoParallel(b *testing.B) {
	ctx := context.Background()
	intKVInfo := prepareIntKVInfo()
	strKVInfo := prepareStrKVInfo()

	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			sendMsg := initClientSendMsg(transport.TTHeader)
			sendMsg.TransInfo().PutTransIntInfo(intKVInfo)
			sendMsg.TransInfo().PutTransStrInfo(strKVInfo)
			sendMsg.Tags()[HeaderFlagsKey] = ttheader.HeaderFlagSupportOutOfOrder

			// encode
			out := remote.NewWriterBuffer(256)
			totalLenField, err := ttHeaderCodec.encode(ctx, sendMsg, out)
			binary.BigEndian.PutUint32(totalLenField, uint32(out.WrittenLen()-Size32+mockPayloadLen))
			test.Assert(b, err == nil, err)

			// decode
			ctx, recvMsg := initServerRecvMsg(ctx)
			buf, err := out.Bytes()
			test.Assert(b, err == nil, err)
			in := remote.NewReaderBuffer(buf)
			err = ttHeaderCodec.decode(ctx, recvMsg, in)
			test.Assert(b, err == nil, err)
			test.Assert(b, recvMsg.PayloadLen() == mockPayloadLen, recvMsg.PayloadLen())

			intKVInfoRecv := recvMsg.TransInfo().TransIntInfo()
			strKVInfoRecv := recvMsg.TransInfo().TransStrInfo()
			test.DeepEqual(b, intKVInfoRecv, intKVInfo)
			test.DeepEqual(b, strKVInfoRecv, strKVInfo)
			flag := recvMsg.Tags()[HeaderFlagsKey]
			test.Assert(b, flag != nil)
			test.Assert(b, flag == ttheader.HeaderFlagSupportOutOfOrder)
		}
	})
}

func BenchmarkTTHeaderCodecParallel(b *testing.B) {
	ctx := context.Background()

	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			sendMsg := initClientSendMsg(transport.TTHeader)
			// encode
			out := remote.NewWriterBuffer(256)
			totalLenField, err := ttHeaderCodec.encode(ctx, sendMsg, out)
			binary.BigEndian.PutUint32(totalLenField, uint32(out.WrittenLen()-Size32+mockPayloadLen))
			test.Assert(b, err == nil, err)

			// decode
			ctx, recvMsg := initServerRecvMsg(ctx)
			buf, err := out.Bytes()
			test.Assert(b, err == nil, err)
			in := remote.NewReaderBuffer(buf)
			err = ttHeaderCodec.decode(ctx, recvMsg, in)
			test.Assert(b, err == nil, err)
			test.Assert(b, recvMsg.PayloadLen() == mockPayloadLen, recvMsg.PayloadLen())
		}
	})
}

var (
	mockServiceName = "mock service"
	mockMethod      = "mock"
)

func mockCliRPCInfo() rpcinfo.RPCInfo {
	return rpcinfo.NewRPCInfo(
		rpcinfo.EmptyEndpointInfo(),
		remoteinfo.NewRemoteInfo(&rpcinfo.EndpointBasicInfo{ServiceName: mockServiceName}, mockMethod).ImmutableView(),
		rpcinfo.NewInvocation("", mockMethod),
		rpcinfo.NewRPCConfig(), rpcinfo.NewRPCStats())
}

func mockSvrRPCInfo() rpcinfo.RPCInfo {
	return rpcinfo.NewRPCInfo(rpcinfo.EmptyEndpointInfo(),
		rpcinfo.FromBasicInfo(&rpcinfo.EndpointBasicInfo{ServiceName: mockServiceName}),
		rpcinfo.NewServerInvocation(),
		rpcinfo.NewRPCConfig(), rpcinfo.NewRPCStats())
}

type mockMsg struct {
	msg string
}

func initServerRecvMsgWithMockMsg(ctx context.Context) (context.Context, remote.Message) {
	svcSearcher := mocksremote.NewDefaultSvcSearcher()
	ri := mockSvrRPCInfo()
	ctx = remote.WithServiceSearcher(ctx, svcSearcher)
	req := &mockMsg{}
	return ctx, remote.NewMessage(req, ri, remote.Call, remote.Server)
}

func initServerRecvMsg(ctx context.Context) (context.Context, remote.Message) {
	svcSearcher := mocksremote.NewDefaultSvcSearcher()
	ri := mockSvrRPCInfo()
	ctx = remote.WithServiceSearcher(ctx, svcSearcher)
	msg := remote.NewMessage(nil, ri, remote.Call, remote.Server)
	return ctx, msg
}

func initClientSendMsg(tp transport.Protocol, payloadLen ...int) remote.Message {
	req := &mockMsg{}
	if len(payloadLen) != 0 {
		req.msg = string(make([]byte, payloadLen[0]))
	}

	svcInfo := mocks.ServiceInfo()
	ri := mockCliRPCInfo()
	msg := remote.NewMessage(req, ri, remote.Call, remote.Client)
	mcfg := rpcinfo.AsMutableRPCConfig(msg.RPCInfo().Config())
	mcfg.SetTransportProtocol(tp)
	mcfg.SetPayloadCodec(svcInfo.PayloadCodec)
	return msg
}

func initServerSendMsg(tp transport.Protocol, payloadLen ...int) remote.Message {
	resp := &mockMsg{}
	if len(payloadLen) != 0 {
		resp.msg = string(make([]byte, payloadLen[0]))
	}
	msg := remote.NewMessage(resp, mockSvrRPCInfo(), remote.Reply, remote.Server)
	mcfg := rpcinfo.AsMutableRPCConfig(msg.RPCInfo().Config())
	mcfg.SetTransportProtocol(tp)
	mcfg.SetPayloadCodec(mocks.ServiceInfo().PayloadCodec)
	return msg
}

func initClientRecvMsg() remote.Message {
	resp := &mockMsg{}
	msg := remote.NewMessage(resp, mockCliRPCInfo(), remote.Reply, remote.Client)
	return msg
}

var _ discovery.Instance = &mockInst{}

type mockInst struct {
	addr net.Addr
}

func (m *mockInst) Address() net.Addr {
	return m.addr
}

func (m *mockInst) RefreshInstanceWithAddr(addr net.Addr) discovery.Instance {
	m.addr = addr
	return m
}

func (m *mockInst) Weight() int {
	return 10
}

func (m *mockInst) Tag(key string) (value string, exist bool) {
	return
}

func prepareIntKVInfo() map[uint16]string {
	kvInfo := map[uint16]string{
		transmeta.FromService: "mockFromService",
		transmeta.FromMethod:  "mockFromMethod",
		transmeta.ToService:   "mockToService",
		transmeta.ToMethod:    "mockToMethod",
	}
	return kvInfo
}

func prepareStrKVInfo() map[string]string {
	kvInfo := map[string]string{transmeta.HeaderIDLServiceName: mocks.MockServiceName}
	return kvInfo
}

func prepareStrKVInfoWithGDPRToken() map[string]string {
	kvInfo := map[string]string{
		transmeta.GDPRToken:             "mockToken",
		transmeta.HeaderTransRemoteAddr: "mockRemoteAddr",
		transmeta.HeaderIDLServiceName:  mocks.MockServiceName,
	}
	return kvInfo
}

// // TODO 是否提供buf.writeInt8/16/32方法，否则得先计算，然后malloc，最后write，待确认频繁malloc是否有影响
// 暂时不删除，测试一次malloc, 和多次malloc差异
//
//lint:ignore U1000 until encode2 is used
func (t ttHeader) encode2(ctx context.Context, message remote.Message, payloadBuf, out remote.ByteBuffer) error {
	tm := message.TransInfo()

	// 1. header meta
	headerMeta, err := out.Malloc(ttheader.TTHeaderMetaSize)
	if err != nil {
		return perrors.NewProtocolError(err)
	}

	totalLenField := headerMeta[0:4]
	headerSizeField := headerMeta[12:14]

	binary.BigEndian.PutUint32(headerMeta[4:8], ttheader.TTHeaderMagic+uint32(getFlags(message)))
	binary.BigEndian.PutUint32(headerMeta[8:12], uint32(message.RPCInfo().Invocation().SeqID()))

	var transformIDs []uint8
	// 2. header info, calculate size
	var headerInfo []byte
	var headerInfoSize int
	// PROTOCOL ID(u8) + NUM TRANSFORMS(always 0)(u8) + TRANSFORM IDs([]u8)
	headerInfoSize = 1 + 1 + len(transformIDs)
	// str kv info
	if len(tm.TransStrInfo()) > 0 {
		headerInfoSize += 3 // INFO ID TYPE(u8) + NUM HEADERS(u16)
		for k, v := range tm.TransStrInfo() {
			headerInfoSize += 4 + len(k) + len(v)
		}
	}
	// int kv info
	if len(tm.TransIntInfo()) > 0 {
		headerInfoSize += 3 // INFO ID TYPE(u8) + NUM HEADERS(u16)
		for _, v := range tm.TransIntInfo() {
			headerInfoSize += 4 + len(v) // key is uint16
		}
	}
	// padding blank
	padding := (4 - headerInfoSize%4) % 4
	headerInfoSize += padding
	binary.BigEndian.PutUint16(headerSizeField, uint16(headerInfoSize/4))
	totalLen := ttheader.TTHeaderMetaSize - Size32 + headerInfoSize + payloadBuf.WrittenLen()
	binary.BigEndian.PutUint32(totalLenField, uint32(totalLen))

	// 3.  header info, malloc and write
	if headerInfo, err = out.Malloc(headerInfoSize); err != nil {
		return perrors.NewProtocolError(err)
	}
	headerInfo[0] = byte(getProtocolID(message.RPCInfo().Config().PayloadCodec()))
	headerInfo[1] = byte(len(transformIDs))
	hdIdx := 2
	for tid := range transformIDs {
		// TODO 需确认下transformId的编码格式
		headerInfo[hdIdx] = byte(tid)
		hdIdx++
	}
	// str kv info
	if len(tm.TransStrInfo()) > 0 {
		// INFO ID TYPE(u8) + NUM HEADERS(u16)
		headerInfo[hdIdx] = byte(ttheader.InfoIDKeyValue)
		binary.BigEndian.PutUint16(headerInfo[hdIdx+1:hdIdx+3], uint16(len(tm.TransStrInfo())))
		hdIdx += 3
		for key, value := range tm.TransStrInfo() {
			binary.BigEndian.PutUint16(headerInfo[hdIdx:hdIdx+2], uint16(len(key)))
			hdIdx += 2
			hdIdx += copy(headerInfo[hdIdx:], key)

			binary.BigEndian.PutUint16(headerInfo[hdIdx:hdIdx+2], uint16(len(value)))
			hdIdx += 2
			hdIdx += copy(headerInfo[hdIdx:], value)
		}
	}
	// int kv info
	if len(tm.TransIntInfo()) > 0 {
		// INFO ID TYPE(u8) + NUM HEADERS(u16)
		headerInfo[hdIdx] = byte(ttheader.InfoIDIntKeyValue)
		binary.BigEndian.PutUint16(headerInfo[hdIdx+1:hdIdx+3], uint16(len(tm.TransIntInfo())))
		hdIdx += 3
		for key, value := range tm.TransIntInfo() {
			binary.BigEndian.PutUint16(headerInfo[hdIdx:hdIdx+2], key)
			hdIdx += 2

			binary.BigEndian.PutUint16(headerInfo[hdIdx:hdIdx+2], uint16(len(value)))
			hdIdx += 2
			hdIdx += copy(headerInfo[hdIdx:], value)
		}
	}
	// padding = (4 - headerInfoSize%4) % 4
	for padding := (4 - hdIdx%4) % 4; padding > 0; padding-- {
		headerInfo[hdIdx] = byte(0)
		hdIdx++
	}
	if hdIdx != headerInfoSize {
		return perrors.NewProtocolErrorWithMsg("the size that header info write is not equal with malloc size")
	}

	if err := out.AppendBuffer(payloadBuf); err != nil {
		return perrors.NewProtocolError(err)
	}
	return err
}

func TestHeaderFlags(t *testing.T) {
	// case 1: correct HeaderFlags
	msg := initClientSendMsg(transport.TTHeader)
	msg.Tags()[HeaderFlagsKey] = ttheader.HeaderFlagSASL | ttheader.HeaderFlagSupportOutOfOrder
	hfs := getFlags(msg)
	test.Assert(t, hfs == ttheader.HeaderFlagSASL|ttheader.HeaderFlagSupportOutOfOrder)

	// case 2: invalid type for HeaderFlagsKey
	msg = initClientSendMsg(transport.TTHeader)
	msg.Tags()[HeaderFlagsKey] = 1
	hfs = getFlags(msg)
	test.Assert(t, hfs == 0)

	// case 3: setFlags then get Flags
	msg = initClientSendMsg(transport.TTHeader)
	setFlags(ttheader.HeaderFlagSupportOutOfOrder, msg)
	hfs = getFlags(msg)
	test.Assert(t, hfs == ttheader.HeaderFlagSupportOutOfOrder, hfs)
}
