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

import (
	"context"
	"encoding/binary"
	"fmt"

	"github.com/bytedance/gopkg/lang/dirtmake"

	"github.com/cloudwego/kitex/pkg/generic/thrift"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/codec"
	"github.com/cloudwego/kitex/pkg/remote/codec/perrors"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
)

var (
	_  remote.PayloadCodec = &binaryThriftCodec{}
	wb                     = thrift.NewWriteBinary()
)

type binaryReqType = []byte

type binaryThriftCodec struct {
	thriftCodec remote.PayloadCodec
}

func (c *binaryThriftCodec) Marshal(ctx context.Context, msg remote.Message, out remote.ByteBuffer) error {
	data := msg.Data()
	if data == nil {
		return perrors.NewProtocolErrorWithMsg("invalid marshal data in rawThriftBinaryCodec: nil")
	}
	if msg.MessageType() == remote.Exception {
		if err := c.thriftCodec.Marshal(ctx, msg, out); err != nil {
			return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("rawThriftBinaryCodec Marshal exception failed, err: %s", err.Error()))
		}
		return nil
	}
	var transBuff []byte
	var ok bool
	if msg.RPCRole() == remote.Server {
		// Business error only works properly when using TTHeader and HTTP2 transmission protocols
		// If there is a business error, data.(*Result).Success will be nil, and an empty payload will be constructed here to return
		if msg.RPCInfo().Invocation().BizStatusErr() != nil {
			msg.Data().(WithCodec).SetCodec(wb)
			return thriftCodec.Marshal(ctx, msg, out)
		}
		gResult := data.(*Result)
		transBinary := gResult.Success
		if transBuff, ok = transBinary.(binaryReqType); !ok {
			return perrors.NewProtocolErrorWithMsg("invalid marshal result in rawThriftBinaryCodec: must be []byte")
		}
	} else {
		gArg := data.(*Args)
		transBinary := gArg.Request
		if transBuff, ok = transBinary.(binaryReqType); !ok {
			return perrors.NewProtocolErrorWithMsg("invalid marshal request in rawThriftBinaryCodec: must be []byte")
		}
		if err := SetSeqID(msg.RPCInfo().Invocation().SeqID(), transBuff); err != nil {
			return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("rawThriftBinaryCodec set seqID failed, err: %s", err.Error()))
		}
	}
	_, err := out.WriteBinary(transBuff)
	return err
}

func (c *binaryThriftCodec) Unmarshal(ctx context.Context, msg remote.Message, in remote.ByteBuffer) error {
	magicAndMsgType, err := codec.PeekUint32(in)
	if err != nil {
		return err
	}
	msgType := magicAndMsgType & codec.FrontMask
	if msgType == uint32(remote.Exception) {
		return c.thriftCodec.Unmarshal(ctx, msg, in)
	}
	payloadLen := msg.PayloadLen()
	transBuff := dirtmake.Bytes(payloadLen, payloadLen)
	_, err = in.ReadBinary(transBuff)
	if err != nil {
		return err
	}
	if err := readBinaryMethod(ctx, transBuff, msg); err != nil {
		return err
	}

	if err = codec.NewDataIfNeeded(serviceinfo.GenericMethod, msg); err != nil {
		return err
	}
	data := msg.Data()
	if msg.RPCRole() == remote.Server {
		gArg := data.(*Args)
		gArg.Method = msg.RPCInfo().Invocation().MethodName()
		gArg.Request = transBuff
	} else {
		gResult := data.(*Result)
		gResult.Success = transBuff
	}
	return nil
}

func (c *binaryThriftCodec) Name() string {
	return "RawThriftBinary"
}

// SetSeqID is used to reset seqID for thrift payload.
// For client side, you don't need this function, Kitex will gen seqID and set it into transport protocol to ignore
// inconsistent seqID between thrift payload and transport protocol, reset the seqID to that generated by kitex for
// client side by default.
// But for server side(binary generic server), you need to return the same seqID with upstream, it is suggested to keep
// the upstream seqID(use GetSeqID) then use SetSeqID to reset the seqID of transBuff.
func SetSeqID(seqID int32, transBuff []byte) error {
	seqID4Bytes, err := getSeqID4Bytes(transBuff)
	if err != nil {
		return err
	}
	binary.BigEndian.PutUint32(seqID4Bytes, uint32(seqID))
	return nil
}

// GetSeqID from thrift buffered binary.
func GetSeqID(transBuff []byte) (int32, error) {
	seqID4Bytes, err := getSeqID4Bytes(transBuff)
	if err != nil {
		return 0, err
	}
	seqID := binary.BigEndian.Uint32(seqID4Bytes)
	return int32(seqID), nil
}

// seqID has 4 bytes
func getSeqID4Bytes(transBuff []byte) ([]byte, error) {
	idx := 4
	ret, e := codec.Bytes2Uint32(transBuff[:idx])
	if e != nil {
		return nil, e
	}
	first4Bytes := int32(ret)
	if first4Bytes > 0 {
		return nil, perrors.NewProtocolErrorWithMsg("missing version in Thrift Message")
	}
	version := int64(first4Bytes) & codec.MagicMask
	if version != codec.ThriftV1Magic {
		return nil, perrors.NewProtocolErrorWithType(perrors.BadVersion, "bad version in Thrift Message")
	}
	idx += 4
	ret, e = codec.Bytes2Uint32(transBuff[4:idx])
	if e != nil {
		return nil, e
	}
	methodNameLen := int32(ret)
	if methodNameLen < 0 {
		return nil, perrors.InvalidDataLength
	}
	idx += int(methodNameLen)
	if len(transBuff) < idx+4 {
		return nil, perrors.NewProtocolErrorWithMsg("invalid trans buffer")
	}
	return transBuff[idx : idx+4], nil
}

func readBinaryMethod(ctx context.Context, buff []byte, msg remote.Message) error {
	bufLen := len(buff)
	if bufLen < codec.Size32*2 {
		return perrors.NewProtocolErrorWithMsg(
			fmt.Sprintf("invalid trans buffer in binaryThriftCodec Unmarshal, size=%d less than 8 bytes", bufLen))
	}
	methodLen := int(binary.BigEndian.Uint32(buff[4:8]))
	if bufLen < codec.Size32*2+methodLen || methodLen <= 0 {
		return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("method len[%d] invalid in binaryThriftCodec Unmarshal", methodLen))
	}
	method := string(buff[8:(8 + methodLen)])
	if err := codec.SetOrCheckMethodName(ctx, method, msg); err != nil {
		return perrors.NewProtocolError(err)
	}
	return nil
}
