/*
 * Copyright 2022 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 nphttp2

import (
	"context"
	"errors"
	"sync/atomic"
	"testing"
	"time"

	"github.com/cloudwego/kitex/internal/mocks"
	"github.com/cloudwego/kitex/pkg/consts"

	"github.com/golang/mock/gomock"

	"github.com/cloudwego/kitex/internal/mocks/netpoll"
	mocksremote "github.com/cloudwego/kitex/internal/mocks/remote"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/kerrors"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/grpc"
	grpcTransport "github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/grpc"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
	"github.com/cloudwego/kitex/pkg/streaming"
	"github.com/cloudwego/kitex/pkg/utils/kitexutil"
	"github.com/cloudwego/kitex/transport"
)

var _ remote.StreamingMetaHandler = &mockMetaHandler{}

type mockMetaHandler struct {
	onReadStream func(ctx context.Context) (context.Context, error)
}

func (m mockMetaHandler) OnConnectStream(ctx context.Context) (context.Context, error) {
	return ctx, nil
}

func (m mockMetaHandler) OnReadStream(ctx context.Context) (context.Context, error) {
	return m.onReadStream(ctx)
}

func TestServerHandler(t *testing.T) {
	var serviceName, methodName atomic.Value
	serviceName.Store("")
	methodName.Store("")
	// init
	opt := newMockServerOption()
	opt.StreamingMetaHandlers = append(opt.StreamingMetaHandlers, &mockMetaHandler{
		onReadStream: func(ctx context.Context) (context.Context, error) {
			service, _ := kitexutil.GetIDLServiceName(ctx)
			serviceName.Store(service)
			method, _ := kitexutil.GetMethod(ctx)
			methodName.Store(method)
			return ctx, nil
		},
	})
	opt.SvcSearcher = mocksremote.NewMockSvcSearcher(map[string]*serviceinfo.ServiceInfo{
		"Greeter": {
			Methods: map[string]serviceinfo.MethodInfo{
				"SayHello": serviceinfo.NewMethodInfo(func(ctx context.Context, handler, args, result interface{}) error {
					return nil
				}, func() interface{} { return nil }, func() interface{} { return nil }, false),
			},
		},
	})
	msg := newMockNewMessage()
	ri := newMockRPCInfo()
	mcfg := rpcinfo.AsMutableRPCConfig(ri.Config())
	mcfg.SetTransportProtocol(transport.PurePayload)
	mcfg.SetPayloadCodec(serviceinfo.Protobuf)
	msg.RPCInfoFunc = func() rpcinfo.RPCInfo {
		return ri
	}
	msg.RPCRoleFunc = func() remote.RPCRole {
		return remote.Server
	}
	npConn := newMockNpConn(mockAddr0)
	npConn.mockSettingFrame()
	tr, err := newMockServerTransport(npConn)
	test.Assert(t, err == nil, err)
	s := grpc.CreateStream(context.Background(), 1, func(i int) {}, "")
	serverConn := newServerConn(tr, s)
	defer serverConn.Close()

	// test NewTransHandler()
	handler, err := NewSvrTransHandlerFactory().NewTransHandler(opt)
	test.Assert(t, err == nil, err)

	// test Read()
	// mock grpc encoded msg data into stream recv buffer
	newMockStreamRecvHelloRequest(s)
	ctx, err := handler.Read(context.Background(), serverConn, msg)
	test.Assert(t, ctx != nil, ctx)
	test.Assert(t, err == nil, err)

	// test write()
	ctx, err = handler.Write(context.Background(), serverConn, msg)
	test.Assert(t, ctx != nil, ctx)
	test.Assert(t, err == nil, err)

	// test SetInvokeHandleFunc()
	var calledInvoke int32
	handler.(remote.InvokeHandleFuncSetter).SetInvokeHandleFunc(func(ctx context.Context, req, resp interface{}) (err error) {
		fromMethod, ok := ctx.Value(consts.CtxKeyMethod).(string)
		test.Assert(t, ok)
		test.Assert(t, fromMethod == "SayHello", fromMethod)
		atomic.StoreInt32(&calledInvoke, 1)
		return nil
	})

	// mock a setting frame to pass the server side preface check
	npConn.mockSettingFrame()
	// mock a headerFrame so onRead() can start working
	npConn.mockMetaHeaderFrame()
	go func() {
		// test OnActive()
		ctx, err := handler.OnActive(newMockCtxWithRPCInfo(), npConn)
		test.Assert(t, err == nil, err)

		handler.OnRead(ctx, npConn)

		// test OnInactive()
		handler.OnInactive(ctx, npConn)
		test.Assert(t, err == nil, err)
	}()

	// sleep 50 mills so server can handle metaHeader frame
	time.Sleep(time.Millisecond * 50)
	test.Assert(t, serviceName.Load().(string) == "Greeter", serviceName.Load())
	test.Assert(t, methodName.Load().(string) == "SayHello", methodName.Load())
	test.Assert(t, atomic.LoadInt32(&calledInvoke) == 1)

	// test OnError()
	handler.OnError(context.Background(), context.Canceled, npConn)

	// test SetPipeline()
	handler.SetPipeline(nil)
}

type mockStream struct {
	streaming.ServerStream
	recv func(ctx context.Context, msg interface{}) error
	send func(ctx context.Context, msg interface{}) error
}

func (s *mockStream) RecvMsg(ctx context.Context, m interface{}) error {
	return s.recv(ctx, m)
}

func (s *mockStream) SendMsg(ctx context.Context, m interface{}) error {
	return s.send(ctx, m)
}

func Test_invokeStreamUnaryHandler(t *testing.T) {
	t.Run("recv err", func(t *testing.T) {
		expectedErr := errors.New("mock err")
		s := &mockStream{
			recv: func(ctx context.Context, msg interface{}) error {
				return expectedErr
			},
		}
		var newArgsCalled, newResultCalled bool
		mi := serviceinfo.NewMethodInfo(nil,
			func() interface{} {
				newArgsCalled = true
				return nil
			},
			func() interface{} {
				newResultCalled = true
				return nil
			},
			false,
		)
		hdl := func(ctx context.Context, req, resp interface{}) (err error) {
			return nil
		}

		err := invokeStreamUnaryHandler(context.Background(), s, mi, hdl, nil)

		test.Assert(t, err == expectedErr, err)
		test.Assert(t, newArgsCalled)
		test.Assert(t, newResultCalled)
	})
	t.Run("handler err", func(t *testing.T) {
		expectedErr := errors.New("mock err")
		var newArgsCalled, newResultCalled, handlerCalled, recvCalled bool
		s := &mockStream{
			recv: func(ctx context.Context, msg interface{}) error {
				recvCalled = true
				return nil
			},
		}
		mi := serviceinfo.NewMethodInfo(nil,
			func() interface{} {
				newArgsCalled = true
				return nil
			},
			func() interface{} {
				newResultCalled = true
				return nil
			},
			false,
		)
		hdl := func(ctx context.Context, req, resp interface{}) (err error) {
			handlerCalled = true
			return expectedErr
		}

		err := invokeStreamUnaryHandler(context.Background(), s, mi, hdl, nil)

		test.Assert(t, err == expectedErr, err)
		test.Assert(t, recvCalled)
		test.Assert(t, newArgsCalled)
		test.Assert(t, newResultCalled)
		test.Assert(t, handlerCalled)
	})

	t.Run("biz err", func(t *testing.T) {
		expectedErr := kerrors.NewBizStatusError(100, "mock biz error")
		var newArgsCalled, newResultCalled, handlerCalled, recvCalled, sendCalled bool
		s := &mockStream{
			recv: func(ctx context.Context, msg interface{}) error {
				recvCalled = true
				return nil
			},
			send: func(ctx context.Context, msg interface{}) error {
				sendCalled = true
				return nil
			},
		}
		mi := serviceinfo.NewMethodInfo(nil,
			func() interface{} {
				newArgsCalled = true
				return nil
			},
			func() interface{} {
				newResultCalled = true
				return nil
			},
			false,
		)
		hdl := func(ctx context.Context, req, resp interface{}) (err error) {
			handlerCalled = true
			return nil
		}

		ivk := rpcinfo.NewInvocation("test", "test")
		ivk.SetBizStatusErr(expectedErr)
		ri := rpcinfo.NewRPCInfo(nil, nil, ivk, nil, nil)

		err := invokeStreamUnaryHandler(context.Background(), s, mi, hdl, ri)

		test.Assert(t, err == nil, err)
		test.Assert(t, recvCalled)
		test.Assert(t, newArgsCalled)
		test.Assert(t, newResultCalled)
		test.Assert(t, handlerCalled)
		test.Assert(t, !sendCalled)
	})

	t.Run("send err", func(t *testing.T) {
		expectedErr := errors.New("mock err")
		var newArgsCalled, newResultCalled, handlerCalled, recvCalled, sendCalled bool
		s := &mockStream{
			recv: func(ctx context.Context, msg interface{}) error {
				recvCalled = true
				return nil
			},
			send: func(ctx context.Context, msg interface{}) error {
				sendCalled = true
				return expectedErr
			},
		}
		mi := serviceinfo.NewMethodInfo(nil,
			func() interface{} {
				newArgsCalled = true
				return nil
			},
			func() interface{} {
				newResultCalled = true
				return nil
			},
			false,
		)
		hdl := func(ctx context.Context, req, resp interface{}) (err error) {
			handlerCalled = true
			return nil
		}

		ivk := rpcinfo.NewInvocation("test", "test")
		ri := rpcinfo.NewRPCInfo(nil, nil, ivk, nil, nil)

		err := invokeStreamUnaryHandler(context.Background(), s, mi, hdl, ri)

		test.Assert(t, err == expectedErr, err)
		test.Assert(t, recvCalled)
		test.Assert(t, newArgsCalled)
		test.Assert(t, newResultCalled)
		test.Assert(t, handlerCalled)
		test.Assert(t, sendCalled)
	})
}

func TestSvrTransHandlerProtocolMatch(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	th := &svrTransHandler{}
	// netpoll reader
	// 1. success
	reader := netpoll.NewMockReader(ctrl)
	reader.EXPECT().Peek(prefaceReadAtMost).Times(1).Return(grpcTransport.ClientPreface, nil)
	conn := netpoll.NewMockConnection(ctrl)
	conn.EXPECT().Reader().AnyTimes().Return(reader)
	err := th.ProtocolMatch(context.Background(), conn)
	test.Assert(t, err == nil, err)
	// 2. failed, no reader
	conn = netpoll.NewMockConnection(ctrl)
	conn.EXPECT().Reader().AnyTimes().Return(nil)
	err = th.ProtocolMatch(context.Background(), conn)
	test.Assert(t, err != nil, err)
	// 3. failed, wrong preface
	failedReader := netpoll.NewMockReader(ctrl)
	failedReader.EXPECT().Peek(prefaceReadAtMost).Times(1).Return([]byte{}, nil)
	conn = netpoll.NewMockConnection(ctrl)
	conn.EXPECT().Reader().AnyTimes().Return(failedReader)
	err = th.ProtocolMatch(context.Background(), conn)
	test.Assert(t, err != nil, err)

	// non-netpoll reader, all failed
	rawConn := &mocks.Conn{}
	err = th.ProtocolMatch(context.Background(), rawConn)
	test.Assert(t, err != nil, err)
}

func Test_parseGraceAndPollTime(t *testing.T) {
	// without timeout in ctx
	ctx := context.Background()
	graceTime, pollTime := parseGraceAndPollTime(ctx)
	test.Assert(t, graceTime == defaultGraceTime, graceTime)
	test.Assert(t, pollTime == defaultMaxPollTime, pollTime)

	// with timeout longer than default timeout
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	graceTime, pollTime = parseGraceAndPollTime(ctx)
	test.Assert(t, graceTime > defaultGraceTime, graceTime)
	// defaultMaxPollTime is the max poll time
	test.Assert(t, pollTime == defaultMaxPollTime, pollTime)

	// with timeout shorter than default timeout
	ctx, cancel = context.WithTimeout(context.Background(), 400*time.Millisecond)
	defer cancel()
	graceTime, pollTime = parseGraceAndPollTime(ctx)
	test.Assert(t, graceTime < defaultGraceTime, graceTime)
	test.Assert(t, pollTime < defaultMaxPollTime, pollTime)
}
