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

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

	"github.com/cloudwego/kitex/internal/generic"
	"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/rpcinfo"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
	"github.com/cloudwego/kitex/pkg/stats"
)

// CodecType is config of the thrift codec. Priority: Frugal > FastMode > Normal
type CodecType int

const (
	// Basic can be used for disabling fastCodec and frugal
	Basic     CodecType = 0b0000
	FastWrite CodecType = 0b0001
	FastRead  CodecType = 0b0010

	FastReadWrite = FastRead | FastWrite

	FrugalWrite CodecType = 0b0100
	FrugalRead  CodecType = 0b1000

	FrugalReadWrite = FrugalWrite | FrugalRead

	EnableSkipDecoder CodecType = 0b10000
)

var (
	errEncodeMismatchMsgType = remote.NewTransErrorWithMsg(remote.InvalidProtocol,
		"encode failed, codec msg type not match with thriftCodec")
	errDecodeMismatchMsgType = remote.NewTransErrorWithMsg(remote.InvalidProtocol,
		"decode failed, codec msg type not match with thriftCodec")
)

// NewThriftCodec creates the thrift binary codec.
func NewThriftCodec() remote.PayloadCodec {
	return &thriftCodec{FastWrite | FastRead}
}

// IsThriftCodec checks if the codec is thriftCodec
func IsThriftCodec(c remote.PayloadCodec) bool {
	_, ok := c.(*thriftCodec)
	return ok
}

// NewThriftFrugalCodec creates the thrift binary codec powered by frugal.
// Eg: xxxservice.NewServer(handler, server.WithPayloadCodec(thrift.NewThriftCodecWithConfig(thrift.FastWrite | thrift.FastRead)))
func NewThriftCodecWithConfig(c CodecType) remote.PayloadCodec {
	return &thriftCodec{c}
}

// NewThriftCodecDisableFastMode creates the thrift binary codec which can control if do fast codec.
// Eg: xxxservice.NewServer(handler, server.WithPayloadCodec(thrift.NewThriftCodecDisableFastMode(true, true)))
func NewThriftCodecDisableFastMode(disableFastWrite, disableFastRead bool) remote.PayloadCodec {
	var c CodecType
	if !disableFastRead {
		c |= FastRead
	}
	if !disableFastWrite {
		c |= FastWrite
	}
	return &thriftCodec{c}
}

// thriftCodec implements PayloadCodec
type thriftCodec struct {
	CodecType
}

// IsSet returns true if t is set
func (c thriftCodec) IsSet(t CodecType) bool {
	return c.CodecType&t != 0
}

// IsDataLenDeterministic return true if dataLen > 0 or we can use SkipDecoder
func (c thriftCodec) IsDataLenDeterministic(dataLen int) bool {
	return dataLen > 0 || c.IsSet(EnableSkipDecoder)
}

// Marshal implements the remote.PayloadCodec interface.
func (c thriftCodec) Marshal(ctx context.Context, message remote.Message, out remote.ByteBuffer) error {
	// prepare info
	methodName := message.RPCInfo().Invocation().MethodName()
	if methodName == "" {
		return errors.New("empty methodName in thrift Marshal")
	}
	msgType := message.MessageType()
	seqID := message.RPCInfo().Invocation().SeqID()

	// ???? for fixing resp==nil, err==nil? don't know
	if err := codec.NewDataIfNeeded(methodName, message); err != nil {
		return err
	}
	data := message.Data()
	if message.MessageType() == remote.Exception {
		// if remote.Exception, we always use fastcodec
		if transErr, ok := data.(*remote.TransError); ok {
			ex := thrift.NewApplicationException(transErr.TypeID(), transErr.Error())
			return fastMarshal(out, methodName, msgType, seqID, ex)
		} else if err, ok := data.(error); ok {
			ex := thrift.NewApplicationException(remote.InternalError, err.Error())
			return fastMarshal(out, methodName, msgType, seqID, ex)
		} else {
			return fmt.Errorf("got %T for remote.Exception", data)
		}
	}

	typecodec := getTypeCodec(data)

	// encode with frugal codec
	if c.IsSet(FrugalWrite) && typecodec.Frugal {
		return frugalMarshal(out, methodName, msgType, seqID, data)
	}

	// encode with FastWrite
	if c.IsSet(FastWrite) && typecodec.FastCodec {
		return fastMarshal(out, methodName, msgType, seqID, data.(thrift.FastCodec))
	}

	// generic call
	if msg, ok := data.(generic.ThriftWriter); ok {
		return encodeGenericThrift(out, ctx, methodName, msgType, seqID, msg)
	}

	// fallback to old thrift way (slow)
	if typecodec.Apache {
		return apacheMarshal(out, ctx, methodName, msgType, seqID, data)
	}

	// try fallback to fastcodec or frugal even though CodecType=Basic
	if typecodec.FastCodec {
		return fastMarshal(out, methodName, msgType, seqID, data.(thrift.FastCodec))
	}
	if typecodec.Frugal {
		return frugalMarshal(out, methodName, msgType, seqID, data)
	}
	return errEncodeMismatchMsgType
}

func encodeGenericThrift(out bufiox.Writer, ctx context.Context, method string, msgType remote.MessageType, seqID int32, msg generic.ThriftWriter) error {
	binaryWriter := thrift.NewBufferWriter(out)
	if err := binaryWriter.WriteMessageBegin(method, thrift.TMessageType(msgType), seqID); err != nil {
		return perrors.NewProtocolErrorWithErrMsg(err, fmt.Sprintf("thrift marshal, Write failed: %s", err.Error()))
	}
	binaryWriter.Recycle()
	if err := msg.Write(ctx, method, out); err != nil {
		return perrors.NewProtocolErrorWithErrMsg(err, fmt.Sprintf("thrift marshal, Write failed: %s", err.Error()))
	}
	return nil
}

// Unmarshal implements the remote.PayloadCodec interface.
func (c thriftCodec) Unmarshal(ctx context.Context, message remote.Message, in remote.ByteBuffer) error {
	br := thrift.NewBufferReader(in)
	defer br.Recycle()

	methodName, msgType, seqID, err := br.ReadMessageBegin()
	if err != nil {
		return perrors.NewProtocolErrorWithErrMsg(err, fmt.Sprintf("thrift unmarshal, ReadMessageBegin failed: %s", err.Error()))
	}
	if err = codec.UpdateMsgType(uint32(msgType), message); err != nil {
		return err
	}

	// exception message
	if message.MessageType() == remote.Exception {
		return unmarshalThriftException(in)
	}

	if err = validateMessageBeforeDecode(ctx, message, seqID, methodName); err != nil {
		return err
	}

	// decode thrift data
	data := message.Data()
	msgBeginLen := thrift.Binary.MessageBeginLength(methodName)
	dataLen := message.PayloadLen() - msgBeginLen
	// For Buffer Protocol, dataLen would be negative. Set it to zero so as not to confuse
	if dataLen < 0 {
		dataLen = 0
	}

	ri := message.RPCInfo()
	rpcinfo.Record(ctx, ri, stats.WaitReadStart, nil)
	if msg, ok := data.(generic.ThriftReader); ok {
		err = msg.Read(ctx, methodName, dataLen, in)
		if err != nil {
			err = remote.NewTransError(remote.ProtocolError, err)
		}
	} else {
		err = c.unmarshalThriftData(in, data, dataLen)
	}
	rpcinfo.Record(ctx, ri, stats.WaitReadFinish, err)
	if err != nil {
		return err
	}
	return err
}

// validateMessageBeforeDecode validate message before decode
func validateMessageBeforeDecode(ctx context.Context, message remote.Message, seqID int32, methodName string) (err error) {
	// For server side, the following error can be sent back and 'SetSeqID' should be executed first to ensure the seqID
	// is right when return Exception back.
	if err = codec.SetOrCheckSeqID(seqID, message); err != nil {
		return err
	}

	if err = codec.SetOrCheckMethodName(ctx, methodName, message); err != nil {
		return err
	}

	if err = codec.NewDataIfNeeded(methodName, message); err != nil {
		return err
	}
	return nil
}

// Name implements the remote.PayloadCodec interface.
func (c thriftCodec) Name() string {
	return serviceinfo.Thrift.String()
}
