/*
 * 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 thrift

import (
	"context"
	"errors"
	"testing"

	"github.com/cloudwego/gopkg/bufiox"
	"github.com/cloudwego/gopkg/protocol/thrift"
	"github.com/cloudwego/netpoll"

	"github.com/cloudwego/kitex/internal/mocks"
	mt "github.com/cloudwego/kitex/internal/mocks/thrift"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/remote"
	netpolltrans "github.com/cloudwego/kitex/pkg/remote/trans/netpoll"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
	"github.com/cloudwego/kitex/transport"
)

var (
	payloadCodec = &thriftCodec{FastWrite | FastRead}
	svcInfo      = mocks.ServiceInfo()

	transportBuffers = []struct {
		Name      string
		NewBuffer func() (bufiox.Writer, bufiox.Reader)
	}{
		{
			Name: "BytesBuffer",
			NewBuffer: func() (bufiox.Writer, bufiox.Reader) {
				conn := mocks.NewIOConn()
				return bufiox.NewDefaultWriter(conn), bufiox.NewDefaultReader(conn)
			},
		},
		{
			Name: "NetpollBuffer",
			NewBuffer: func() (bufiox.Writer, bufiox.Reader) {
				bb := netpolltrans.NewReaderWriterByteBuffer(netpoll.NewLinkBuffer(1024))
				return bb, bb
			},
		},
	}
)

func init() {
	svcInfo.Methods["mock"] = serviceinfo.NewMethodInfo(nil, newMockTestArgs, nil, false)
}

type mockWithContext struct {
	ReadFunc  func(ctx context.Context, method string, dataLen int, oprot bufiox.Reader) error
	WriteFunc func(ctx context.Context, method string, oprot bufiox.Writer) error
}

func (m *mockWithContext) Read(ctx context.Context, method string, dataLen int, oprot bufiox.Reader) error {
	if m.ReadFunc != nil {
		return m.ReadFunc(ctx, method, dataLen, oprot)
	}
	return nil
}

func (m *mockWithContext) Write(ctx context.Context, method string, oprot bufiox.Writer) error {
	if m.WriteFunc != nil {
		return m.WriteFunc(ctx, method, oprot)
	}
	return nil
}

func TestWithContext(t *testing.T) {
	for _, tb := range transportBuffers {
		t.Run(tb.Name, func(t *testing.T) {
			ctx := context.Background()

			req := &mockWithContext{WriteFunc: func(ctx context.Context, method string, oprot bufiox.Writer) error {
				return nil
			}}
			ink := rpcinfo.NewInvocation("", "mock")
			ri := rpcinfo.NewRPCInfo(nil, nil, ink, rpcinfo.NewRPCConfig(), nil)
			msg := remote.NewMessage(req, ri, remote.Call, remote.Client)
			mcfg := rpcinfo.AsMutableRPCConfig(ri.Config())
			mcfg.SetTransportProtocol(transport.TTHeader)
			mcfg.SetPayloadCodec(svcInfo.PayloadCodec)
			bw, br := tb.NewBuffer()
			bb := remote.NewByteBufferFromBufiox(bw, br)
			err := payloadCodec.Marshal(ctx, msg, bb)
			test.Assert(t, err == nil, err)
			wl := bw.WrittenLen()
			bw.Flush()

			{
				resp := &mockWithContext{ReadFunc: func(ctx context.Context, method string, dataLen int, oprot bufiox.Reader) error {
					return nil
				}}
				ink := rpcinfo.NewInvocation("", "mock")
				ri := rpcinfo.NewRPCInfo(nil, nil, ink, rpcinfo.NewRPCConfig(), nil)
				msg := remote.NewMessage(resp, ri, remote.Call, remote.Client)
				mcfg := rpcinfo.AsMutableRPCConfig(ri.Config())
				mcfg.SetTransportProtocol(transport.TTHeader)
				mcfg.SetPayloadCodec(svcInfo.PayloadCodec)
				msg.SetPayloadLen(wl)
				err = payloadCodec.Unmarshal(ctx, msg, bb)
				test.Assert(t, err == nil, err)
			}
		})
	}
}

func TestNormal(t *testing.T) {
	// msg only supports FastCodec
	for _, tb := range transportBuffers {
		t.Run(tb.Name, func(t *testing.T) {
			ctx := context.Background()
			// encode client side
			sendMsg := initSendMsg(transport.TTHeader)
			bw, br := tb.NewBuffer()
			bb := remote.NewByteBufferFromBufiox(bw, br)
			err := payloadCodec.Marshal(ctx, sendMsg, bb)
			test.Assert(t, err == nil, err)
			wl := bw.WrittenLen()
			bw.Flush()

			// decode server side
			recvMsg := initRecvMsg()
			recvMsg.SetPayloadLen(wl)
			test.Assert(t, err == nil, err)
			err = payloadCodec.Unmarshal(ctx, recvMsg, bb)
			test.Assert(t, err == nil, err)

			// compare Req Arg
			compare(t, sendMsg, recvMsg)
		})
	}

	// msg only supports Basic codec (apache codec)
	for _, tb := range transportBuffers {
		t.Run(tb.Name+"Basic", func(t *testing.T) {
			ctx := context.Background()
			// encode client side
			sendMsg := initSendMsg(transport.TTHeader)
			bw, br := tb.NewBuffer()
			bb := remote.NewByteBufferFromBufiox(bw, br)
			err := payloadCodec.Marshal(ctx, sendMsg, bb)
			test.Assert(t, err == nil, err)
			wl := bw.WrittenLen()
			bw.Flush()

			// decode server side
			recvMsg := initRecvMsg()
			recvMsg.SetPayloadLen(wl)
			test.Assert(t, err == nil, err)
			err = payloadCodec.Unmarshal(ctx, recvMsg, bb)
			test.Assert(t, err == nil, err)

			// compare Req Arg
			compare(t, sendMsg, recvMsg)
		})
	}

	// Exception case
	for _, tb := range transportBuffers {
		t.Run(tb.Name+"Ex", func(t *testing.T) {
			ctx := context.Background()
			// encode client side
			sendMsg := newMsg(remote.NewTransErrorWithMsg(1, "hello"))
			sendMsg.SetMessageType(remote.Exception)
			bw, br := tb.NewBuffer()
			bb := remote.NewByteBufferFromBufiox(bw, br)
			err := payloadCodec.Marshal(ctx, sendMsg, bb)
			test.Assert(t, err == nil, err)
			wl := bw.WrittenLen()
			bw.Flush()

			// decode server side
			recvMsg := newMsg(nil)
			recvMsg.SetPayloadLen(wl)
			test.Assert(t, err == nil, err)
			err = payloadCodec.Unmarshal(ctx, recvMsg, bb)
			test.Assert(t, err != nil)
			te, ok := err.(*remote.TransError)
			test.Assert(t, ok)
			test.Assert(t, te.TypeID() == 1 && te.Error() == "hello", te)
		})
	}
}

func BenchmarkNormalParallel(b *testing.B) {
	for _, tb := range transportBuffers {
		b.Run(tb.Name, func(b *testing.B) {
			ctx := context.Background()

			b.ResetTimer()
			b.RunParallel(func(pb *testing.PB) {
				for pb.Next() {
					// encode // client side
					sendMsg := initSendMsg(transport.TTHeader)
					bw, br := tb.NewBuffer()
					bb := remote.NewByteBufferFromBufiox(bw, br)
					err := payloadCodec.Marshal(ctx, sendMsg, bb)
					test.Assert(b, err == nil, err)
					wl := bw.WrittenLen()
					bw.Flush()

					// decode server side
					recvMsg := initRecvMsg()
					recvMsg.SetPayloadLen(wl)
					test.Assert(b, err == nil, err)
					err = payloadCodec.Unmarshal(ctx, recvMsg, bb)
					test.Assert(b, err == nil, err)

					// compare Req Arg
					sendReq := sendMsg.Data().(*mt.MockTestArgs).Req
					recvReq := recvMsg.Data().(*mt.MockTestArgs).Req
					test.Assert(b, sendReq.Msg == recvReq.Msg)
					test.Assert(b, len(sendReq.StrList) == len(recvReq.StrList))
					test.Assert(b, len(sendReq.StrMap) == len(recvReq.StrMap))
					for i, item := range sendReq.StrList {
						test.Assert(b, item == recvReq.StrList[i])
					}
					for k := range sendReq.StrMap {
						test.Assert(b, sendReq.StrMap[k] == recvReq.StrMap[k])
					}
				}
			})
		})
	}
}

func TestException(t *testing.T) {
	for _, tb := range transportBuffers {
		t.Run(tb.Name, func(t *testing.T) {
			ctx := context.Background()
			ink := rpcinfo.NewInvocation("", "mock")
			ri := rpcinfo.NewRPCInfo(nil, nil, ink, rpcinfo.NewRPCConfig(), nil)
			errInfo := "mock exception"
			transErr := remote.NewTransErrorWithMsg(remote.UnknownMethod, errInfo)
			// encode server side
			errMsg := initServerErrorMsg(transport.TTHeader, ri, transErr)
			bw, br := tb.NewBuffer()
			bb := remote.NewByteBufferFromBufiox(bw, br)
			err := payloadCodec.Marshal(ctx, errMsg, bb)
			test.Assert(t, err == nil, err)
			wl := bw.WrittenLen()
			bw.Flush()

			// decode client side
			recvMsg := initClientRecvMsg(ri)
			recvMsg.SetPayloadLen(wl)
			test.Assert(t, err == nil, err)
			err = payloadCodec.Unmarshal(ctx, recvMsg, bb)
			test.Assert(t, err != nil)
			transErr, ok := err.(*remote.TransError)
			test.Assert(t, ok, err)
			test.Assert(t, err.Error() == errInfo)
			test.Assert(t, transErr.TypeID() == remote.UnknownMethod)
		})
	}
}

func TestTransErrorUnwrap(t *testing.T) {
	errMsg := "mock err"
	transErr := remote.NewTransError(remote.InternalError, thrift.NewApplicationException(1000, errMsg))
	uwErr, ok := transErr.Unwrap().(*thrift.ApplicationException)
	test.Assert(t, ok)
	test.Assert(t, uwErr.TypeId() == 1000)
	test.Assert(t, transErr.Error() == errMsg)

	uwErr2, ok := errors.Unwrap(transErr).(*thrift.ApplicationException)
	test.Assert(t, ok)
	test.Assert(t, uwErr2.TypeId() == 1000)
	test.Assert(t, uwErr2.Error() == errMsg)
}

func TestSkipDecoder(t *testing.T) {
	testcases := []struct {
		desc     string
		codec    remote.PayloadCodec
		protocol transport.Protocol
		wantErr  bool
	}{
		{
			desc:     "SkipDecoder=false|PurePayload",
			codec:    NewThriftCodec(),
			protocol: transport.PurePayload,
		},
		{
			desc:     "SkipDecoder=false|TTHeader",
			codec:    NewThriftCodec(),
			protocol: transport.TTHeader,
		},
		{
			desc:     "SkipDecoder=true|PurePayload",
			codec:    NewThriftCodecWithConfig(FastRead | FastWrite | EnableSkipDecoder),
			protocol: transport.PurePayload,
		},
		{
			desc:     "SkipDecoder=true|TTHeader",
			codec:    NewThriftCodecWithConfig(FastRead | FastWrite | EnableSkipDecoder),
			protocol: transport.TTHeader,
		},
	}

	for _, tc := range testcases {
		for _, tb := range transportBuffers {
			t.Run(tc.desc+"#"+tb.Name, func(t *testing.T) {
				// encode client side
				sendMsg := initSendMsg(tc.protocol) // always use Basic to test skipdecodec
				bw, br := tb.NewBuffer()
				bb := remote.NewByteBufferFromBufiox(bw, br)
				err := tc.codec.Marshal(context.Background(), sendMsg, bb)
				test.Assert(t, err == nil, err)
				wl := bw.WrittenLen()
				bw.Flush()

				// decode server side
				recvMsg := initRecvMsg()
				if tc.protocol != transport.PurePayload {
					recvMsg.SetPayloadLen(wl)
				}
				err = tc.codec.Unmarshal(context.Background(), recvMsg, bb)
				if tc.wantErr {
					test.Assert(t, err != nil)
					return
				}
				test.Assert(t, err == nil, err)

				// compare Req Arg
				compare(t, sendMsg, recvMsg)
			})
		}
	}
}

func newMsg(data interface{}) remote.Message {
	ink := rpcinfo.NewInvocation("", "mock")
	ri := rpcinfo.NewRPCInfo(nil, rpcinfo.EmptyEndpointInfo(), ink, rpcinfo.NewRPCConfig(), nil)
	return remote.NewMessage(data, ri, remote.Call, remote.Client)
}

func initSendMsg(tp transport.Protocol) remote.Message {
	var _args mt.MockTestArgs // fastcodec only, if basic is true -> apachecodec
	_args.Req = prepareReq()
	msg := newMsg(&_args)
	mcfg := rpcinfo.AsMutableRPCConfig(msg.RPCInfo().Config())
	mcfg.SetTransportProtocol(tp)
	mcfg.SetPayloadCodec(svcInfo.PayloadCodec)
	return msg
}

func initRecvMsg() remote.Message {
	var _args mt.MockTestArgs // fastcodec only, if basic is true -> apachecodec
	return newMsg(&_args)
}

func compare(t *testing.T, sendMsg, recvMsg remote.Message) {
	sendReq := sendMsg.Data().(*mt.MockTestArgs).Req
	recvReq := recvMsg.Data().(*mt.MockTestArgs).Req
	test.Assert(t, sendReq.Msg == recvReq.Msg)
	test.Assert(t, len(sendReq.StrList) == len(recvReq.StrList))
	test.Assert(t, len(sendReq.StrMap) == len(recvReq.StrMap))
	for i, item := range sendReq.StrList {
		test.Assert(t, item == recvReq.StrList[i])
	}
	for k := range sendReq.StrMap {
		test.Assert(t, sendReq.StrMap[k] == recvReq.StrMap[k])
	}
}

func initServerErrorMsg(tp transport.Protocol, ri rpcinfo.RPCInfo, transErr *remote.TransError) remote.Message {
	errMsg := remote.NewMessage(transErr, ri, remote.Exception, remote.Server)
	mcfg := rpcinfo.AsMutableRPCConfig(errMsg.RPCInfo().Config())
	mcfg.SetTransportProtocol(tp)
	mcfg.SetPayloadCodec(svcInfo.PayloadCodec)
	return errMsg
}

func initClientRecvMsg(ri rpcinfo.RPCInfo) remote.Message {
	var resp interface{}
	clientRecvMsg := remote.NewMessage(resp, ri, remote.Reply, remote.Client)
	return clientRecvMsg
}

func prepareReq() *mt.MockReq {
	strMap := make(map[string]string)
	strMap["key1"] = "val1"
	strMap["key2"] = "val2"
	strList := []string{"str1", "str2"}
	req := &mt.MockReq{
		Msg:     "MockReq",
		StrMap:  strMap,
		StrList: strList,
	}
	return req
}

func newMockTestArgs() interface{} {
	return mt.NewMockTestArgs()
}
