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

import (
	"context"
	"net"

	"github.com/bytedance/gopkg/cloud/metainfo"

	"github.com/cloudwego/kitex/pkg/consts"
	"github.com/cloudwego/kitex/pkg/remote"
)

// NewTransMetaHandler to build transMetaHandler that handle transport info
func NewTransMetaHandler(mhs []remote.MetaHandler) remote.DuplexBoundHandler {
	return &transMetaHandler{mhs: mhs}
}

type transMetaHandler struct {
	mhs []remote.MetaHandler
}

// Write exec before encode
func (h *transMetaHandler) Write(ctx context.Context, conn net.Conn, sendMsg remote.Message) (context.Context, error) {
	var err error
	for _, hdlr := range h.mhs {
		ctx, err = hdlr.WriteMeta(ctx, sendMsg)
		if err != nil {
			return ctx, err
		}
	}
	return ctx, nil
}

// OnMessage exec after decode
func (h *transMetaHandler) OnMessage(ctx context.Context, args, result remote.Message) (context.Context, error) {
	var err error
	msg, isServer := getValidMsg(args, result)
	if msg == nil {
		return ctx, nil
	}

	for _, hdlr := range h.mhs {
		ctx, err = hdlr.ReadMeta(ctx, msg)
		if err != nil {
			return ctx, err
		}
	}
	if isServer && result.MessageType() != remote.Exception {
		// Pass-through method name using ctx, the method name will be used as from method in the client.
		ctx = context.WithValue(ctx, consts.CtxKeyMethod, msg.RPCInfo().To().Method())
		// TransferForward converts transient values to transient-upstream values and filters out original transient-upstream values.
		// It should be used before the context is passing from server to client.
		// reference https://github.com/bytedance/gopkg/tree/main/cloud/metainfo
		// Notice, it should be after ReadMeta().
		ctx = metainfo.TransferForward(ctx)
	}
	return ctx, nil
}

// Onactive implements the remote.InboundHandler interface.
func (h *transMetaHandler) OnActive(ctx context.Context, conn net.Conn) (context.Context, error) {
	return ctx, nil
}

// OnRead implements the remote.InboundHandler interface.
func (h *transMetaHandler) OnRead(ctx context.Context, conn net.Conn) (context.Context, error) {
	return ctx, nil
}

// OnInactive implements the remote.InboundHandler interface.
func (h *transMetaHandler) OnInactive(ctx context.Context, conn net.Conn) context.Context {
	return ctx
}

func getValidMsg(args, result remote.Message) (msg remote.Message, isServer bool) {
	if args != nil && args.RPCRole() == remote.Server {
		// server side, read arg
		return args, true
	}
	// client side, read result
	return result, false
}
