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

import (
	"context"
	"strconv"
	"testing"
	"time"

	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/transmeta"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
	"github.com/cloudwego/kitex/transport"
)

func TestIsTTHeader(t *testing.T) {
	t.Run("with ttheader", func(t *testing.T) {
		ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("", ""), rpcinfo.NewRPCConfig(), rpcinfo.NewRPCStats())
		msg := remote.NewMessage(nil, ri, remote.Call, remote.Server)
		mcfg := rpcinfo.AsMutableRPCConfig(ri.Config())
		mcfg.SetTransportProtocol(transport.TTHeader)
		mcfg.SetPayloadCodec(serviceinfo.Thrift)
		test.Assert(t, isTTHeader(msg))
	})
	t.Run("with ttheader framed", func(t *testing.T) {
		ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("", ""), rpcinfo.NewRPCConfig(), rpcinfo.NewRPCStats())
		msg := remote.NewMessage(nil, ri, remote.Call, remote.Server)
		mcfg := rpcinfo.AsMutableRPCConfig(ri.Config())
		mcfg.SetTransportProtocol(transport.TTHeaderFramed)
		mcfg.SetPayloadCodec(serviceinfo.Thrift)
		test.Assert(t, isTTHeader(msg))
	})
}

func TestTTHeaderClientWriteMetainfo(t *testing.T) {
	ctx := context.Background()

	cfg := rpcinfo.NewRPCConfig()
	cfgMutable := rpcinfo.AsMutableRPCConfig(cfg)
	cfgMutable.SetRPCTimeout(time.Millisecond * 100)
	cfgMutable.SetConnectTimeout(time.Millisecond * 1000)
	cfgMutable.LockConfig(rpcinfo.BitRPCTimeout)
	cfgMutable.LockConfig(rpcinfo.BitConnectTimeout)

	fromInfo := rpcinfo.NewEndpointInfo("fromServiceName", "fromMethod", nil, nil)
	toInfo := rpcinfo.NewEndpointInfo("toServiceName", "toMethod", nil, nil)
	ri := rpcinfo.NewRPCInfo(fromInfo, toInfo, rpcinfo.NewInvocation("", ""), cfg, rpcinfo.NewRPCStats())
	msg := remote.NewMessage(nil, ri, remote.Call, remote.Client)
	ri.Invocation().(rpcinfo.InvocationSetter).SetServiceName("toServiceName")

	// pure payload, no effect
	cfgMutable.SetTransportProtocol(transport.PurePayload)
	cfgMutable.SetPayloadCodec(serviceinfo.Thrift)
	ctx, err := ClientTTHeaderHandler.WriteMeta(ctx, msg)
	kvs := msg.TransInfo().TransIntInfo()
	test.Assert(t, err == nil)
	test.Assert(t, len(kvs) == 0)
	strKvs := msg.TransInfo().TransStrInfo()
	test.Assert(t, len(strKvs) == 0)

	// ttheader
	cfgMutable.SetTransportProtocol(transport.TTHeader)
	cfgMutable.SetPayloadCodec(serviceinfo.Thrift)
	_, err = ClientTTHeaderHandler.WriteMeta(ctx, msg)
	test.Assert(t, err == nil)
	kvs = msg.TransInfo().TransIntInfo()
	test.Assert(t, len(kvs) > 0)
	test.Assert(t, kvs[transmeta.FromService] == fromInfo.ServiceName())
	test.Assert(t, kvs[transmeta.FromMethod] == fromInfo.Method())
	test.Assert(t, kvs[transmeta.ToService] == toInfo.ServiceName())
	test.Assert(t, kvs[transmeta.ToMethod] == toInfo.Method())
	test.Assert(t, kvs[transmeta.MsgType] == strconv.Itoa(int(remote.Call)))
	test.Assert(t, kvs[transmeta.TransportType] == unframedTransportType)
	test.Assert(t, kvs[transmeta.RPCTimeout] == "100")
	test.Assert(t, kvs[transmeta.ConnectTimeout] == "1000")
	strKvs = msg.TransInfo().TransStrInfo()
	test.Assert(t, len(strKvs) == 1)
	test.Assert(t, strKvs[transmeta.HeaderIDLServiceName] == "toServiceName")
}

func TestTTHeaderServerReadMetainfo(t *testing.T) {
	ctx := context.Background()
	ri := rpcinfo.NewRPCInfo(rpcinfo.EmptyEndpointInfo(), nil, rpcinfo.NewInvocation("", ""),
		rpcinfo.NewRPCConfig(), rpcinfo.NewRPCStats())
	msg := remote.NewMessage(nil, ri, remote.Call, remote.Client)

	hd := map[uint16]string{
		transmeta.FromService: "fromService",
		transmeta.FromMethod:  "fromMethod",
		transmeta.RPCTimeout:  "100",
	}
	msg.TransInfo().PutTransIntInfo(hd)

	mcfg := rpcinfo.AsMutableRPCConfig(ri.Config())
	mcfg.SetTransportProtocol(transport.PurePayload)
	mcfg.SetPayloadCodec(serviceinfo.Thrift)
	_, err := ServerTTHeaderHandler.ReadMeta(ctx, msg)
	test.Assert(t, err == nil)
	fromEndPoint := msg.RPCInfo().From()
	test.Assert(t, fromEndPoint.ServiceName() == "")
	test.Assert(t, fromEndPoint.Method() == "")

	mcfg.SetTransportProtocol(transport.TTHeader)
	mcfg.SetPayloadCodec(serviceinfo.Thrift)
	_, err = ServerTTHeaderHandler.ReadMeta(ctx, msg)
	test.Assert(t, err == nil)
	test.Assert(t, msg.RPCInfo().From().ServiceName() == hd[transmeta.FromService])
	test.Assert(t, msg.RPCInfo().From().Method() == hd[transmeta.FromMethod])
	test.Assert(t, ri.Config().RPCTimeout() == 100*time.Millisecond)
}

func TestTTHeaderServerWriteMetainfo(t *testing.T) {
	ctx := context.Background()
	ri := rpcinfo.NewRPCInfo(nil, rpcinfo.NewEndpointInfo("", "mock", nil, nil), rpcinfo.NewInvocation("", ""),
		rpcinfo.NewRPCConfig(), rpcinfo.NewRPCStats())
	msg := remote.NewMessage(nil, ri, remote.Call, remote.Client)

	mcfg := rpcinfo.AsMutableRPCConfig(ri.Config())
	mcfg.SetTransportProtocol(transport.PurePayload)
	mcfg.SetPayloadCodec(serviceinfo.Thrift)
	_, err := ServerTTHeaderHandler.WriteMeta(ctx, msg)
	test.Assert(t, err == nil)
	kvs := msg.TransInfo().TransIntInfo()
	test.Assert(t, len(kvs) == 0)

	mcfg.SetTransportProtocol(transport.TTHeader)
	mcfg.SetPayloadCodec(serviceinfo.Thrift)
	_, err = ServerTTHeaderHandler.WriteMeta(ctx, msg)
	test.Assert(t, err == nil)
	kvs = msg.TransInfo().TransIntInfo()
	test.Assert(t, kvs[transmeta.MsgType] == strconv.Itoa(int(remote.Call)))
}
