/*
 * 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"
	"fmt"

	"github.com/cloudwego/gopkg/protocol/ttheader"

	"github.com/cloudwego/kitex/pkg/klog"
	"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"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
	"github.com/cloudwego/kitex/pkg/utils"
)

// Header keys
const (
	// Header Magics
	TTHeaderMagic            = ttheader.TTHeaderMagic
	MeshHeaderMagic   uint32 = 0xFFAF0000
	MeshHeaderLenMask uint32 = 0x0000FFFF

	// HeaderMask         = 0xFFFF0000
	FlagsMask     = ttheader.FlagsMask
	MethodMask    = ttheader.MethodMask
	MaxFrameSize  = ttheader.MaxFrameSize
	MaxHeaderSize = ttheader.MaxHeaderSize
)

type HeaderFlags = ttheader.HeaderFlags

const (
	HeaderFlagsKey              string = "HeaderFlags"
	HeaderFlagSupportOutOfOrder        = ttheader.HeaderFlagSupportOutOfOrder
	HeaderFlagDuplexReverse            = ttheader.HeaderFlagDuplexReverse
	HeaderFlagSASL                     = ttheader.HeaderFlagSASL
)

const (
	TTHeaderMetaSize = ttheader.TTHeaderMetaSize
)

type ProtocolID = ttheader.ProtocolID

const (
	ProtocolIDThriftBinary    = ttheader.ProtocolIDThriftBinary
	ProtocolIDThriftCompact   = ttheader.ProtocolIDThriftCompact
	ProtocolIDThriftCompactV2 = ttheader.ProtocolIDThriftCompactV2
	ProtocolIDKitexProtobuf   = ttheader.ProtocolIDKitexProtobuf
	ProtocolIDDefault         = ttheader.ProtocolIDDefault
)

type InfoIDType = ttheader.InfoIDType

const (
	InfoIDPadding     = ttheader.InfoIDPadding
	InfoIDKeyValue    = ttheader.InfoIDKeyValue
	InfoIDIntKeyValue = ttheader.InfoIDIntKeyValue
	InfoIDACLToken    = ttheader.InfoIDACLToken
)

type ttHeader struct{}

func (t ttHeader) encode(ctx context.Context, message remote.Message, out remote.ByteBuffer) (totalLenField []byte, err error) {
	tm := message.TransInfo()
	if totalLenField, err = ttheader.Encode(ctx, ttheader.EncodeParam{
		Flags:      getFlags(message),
		SeqID:      message.RPCInfo().Invocation().SeqID(),
		ProtocolID: getProtocolID(message.RPCInfo().Config().PayloadCodec()),
		IntInfo:    tm.TransIntInfo(),
		StrInfo:    tm.TransStrInfo(),
	}, out); err != nil {
		return nil, perrors.NewProtocolError(err)
	}
	return totalLenField, nil
}

func (t ttHeader) decode(ctx context.Context, message remote.Message, in remote.ByteBuffer) (err error) {
	var param ttheader.DecodeParam
	if param, err = ttheader.Decode(ctx, in); err != nil {
		return perrors.NewProtocolError(err)
	}
	setFlags(param.Flags, message)
	if err = SetOrCheckSeqID(param.SeqID, message); err != nil {
		klog.Warnf("the seqID in TTHeader check failed, error=%s", err.Error())
		// some framework doesn't write correct seqID in TTheader, to ignore err only check it in payload
		// print log to push the downstream framework to refine it.
	}
	message.TransInfo().PutTransIntInfo(param.IntInfo)
	message.TransInfo().PutTransStrInfo(param.StrInfo)
	fillBasicInfoOfTTHeader(message)
	message.SetPayloadLen(param.PayloadLen)
	return nil
}

func readStrKVInfo(idx *int, buf []byte, info map[string]string) (has bool, err error) {
	kvSize, err := Bytes2Uint16(buf, *idx)
	*idx += 2
	if err != nil {
		return false, fmt.Errorf("error reading str kv info size: %s", err.Error())
	}
	if kvSize <= 0 {
		return false, nil
	}
	for i := uint16(0); i < kvSize; i++ {
		key, n, err := ReadString2BLen(buf, *idx)
		*idx += n
		if err != nil {
			return false, fmt.Errorf("error reading str kv info: %s", err.Error())
		}
		val, n, err := ReadString2BLen(buf, *idx)
		*idx += n
		if err != nil {
			return false, fmt.Errorf("error reading str kv info: %s", err.Error())
		}
		info[key] = val
	}
	return true, nil
}

func getFlags(message remote.Message) ttheader.HeaderFlags {
	var headerFlags ttheader.HeaderFlags
	if message.Tags() != nil && message.Tags()[HeaderFlagsKey] != nil {
		if hfs, ok := message.Tags()[HeaderFlagsKey].(ttheader.HeaderFlags); ok {
			headerFlags = hfs
		} else {
			klog.Warnf("KITEX: the type of headerFlags is invalid, %T", message.Tags()[HeaderFlagsKey])
		}
	}
	return headerFlags
}

func setFlags(flags ttheader.HeaderFlags, message remote.Message) {
	if message.MessageType() == remote.Call {
		message.Tags()[HeaderFlagsKey] = flags
	}
}

// protoID just for ttheader
func getProtocolID(ct serviceinfo.PayloadCodec) ProtocolID {
	switch ct {
	case serviceinfo.Protobuf:
		// ProtocolIDKitexProtobuf is 0x03 at old version(<=v1.9.1) , but it conflicts with ThriftCompactV2.
		// Change the ProtocolIDKitexProtobuf to 0x04 from v1.9.2. But notice! that it is an incompatible change of protocol.
		// For keeping compatible, Kitex use ProtocolIDDefault send ttheader+KitexProtobuf request to ignore the old version
		// check failed if use 0x04. It doesn't make sense, but it won't affect the correctness of RPC call because the actual
		// protocol check at checkPayload func which check payload with HEADER MAGIC bytes of payload.
		return ttheader.ProtocolIDDefault
	}
	return ttheader.ProtocolIDDefault
}

/**
 * +-------------2Byte-------------|-------------2Byte--------------+
 * +----------------------------------------------------------------+
 * |       HEADER MAGIC            |      HEADER SIZE               |
 * +----------------------------------------------------------------+
 * |       HEADER MAP SIZE         |    HEADER MAP...               |
 * +----------------------------------------------------------------+
 * |                                                                |
 * |                            PAYLOAD                             |
 * |                                                                |
 * +----------------------------------------------------------------+
 */
type meshHeader struct{}

//lint:ignore U1000 until encode is used
func (m meshHeader) encode(ctx context.Context, message remote.Message, payloadBuf, out remote.ByteBuffer) error {
	// do nothing, kitex just support decode meshHeader, encode protocol depend on the payload
	return nil
}

func (m meshHeader) decode(ctx context.Context, message remote.Message, in remote.ByteBuffer) error {
	headerMeta, err := in.Next(Size32)
	if err != nil {
		return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("meshHeader read header meta failed, %s", err.Error()))
	}
	if !isMeshHeader(headerMeta) {
		return perrors.NewProtocolErrorWithMsg("not MeshHeader protocol")
	}
	headerLen := Bytes2Uint16NoCheck(headerMeta[Size16:])
	var headerInfo []byte
	if headerInfo, err = in.Next(int(headerLen)); err != nil {
		return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("meshHeader read header buf failed, %s", err.Error()))
	}
	mapInfo := message.TransInfo().TransStrInfo()
	idx := 0
	if _, err = readStrKVInfo(&idx, headerInfo, mapInfo); err != nil {
		return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("meshHeader read kv info failed, %s", err.Error()))
	}
	fillBasicInfoOfTTHeader(message)
	return nil
}

// Fill basic from_info(from service, from address) which carried by ttheader to rpcinfo.
// It is better to fill rpcinfo in matahandlers in terms of design,
// but metahandlers are executed after payloadDecode, we don't know from_info when error happen in payloadDecode.
// So 'fillBasicInfoOfTTHeader' is just for getting more info to output log when decode error happen.
func fillBasicInfoOfTTHeader(msg remote.Message) {
	if msg.RPCRole() == remote.Server {
		fi := rpcinfo.AsMutableEndpointInfo(msg.RPCInfo().From())
		if fi != nil {
			if v := msg.TransInfo().TransStrInfo()[transmeta.HeaderTransRemoteAddr]; v != "" {
				fi.SetAddress(utils.NewNetAddr("tcp", v))
			}
			if v := msg.TransInfo().TransIntInfo()[transmeta.FromService]; v != "" {
				fi.SetServiceName(v)
			}
		}
		if ink, ok := msg.RPCInfo().Invocation().(rpcinfo.InvocationSetter); ok {
			if svcName, ok := msg.TransInfo().TransStrInfo()[transmeta.HeaderIDLServiceName]; ok {
				ink.SetServiceName(svcName)
			}
		}
	} else {
		ti := remoteinfo.AsRemoteInfo(msg.RPCInfo().To())
		if ti != nil {
			if v := msg.TransInfo().TransStrInfo()[transmeta.HeaderTransRemoteAddr]; v != "" {
				ti.SetRemoteAddr(utils.NewNetAddr("tcp", v))
			}
		}
	}
}
