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

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

	"github.com/cloudwego/kitex/pkg/logid"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/metadata"
	"github.com/cloudwego/kitex/pkg/remote/transmeta"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
)

// singletons .
var (
	MetainfoClientHandler = new(metainfoClientHandler)
	MetainfoServerHandler = new(metainfoServerHandler)

	_ remote.MetaHandler          = MetainfoClientHandler
	_ remote.StreamingMetaHandler = MetainfoClientHandler
	_ remote.MetaHandler          = MetainfoServerHandler
	_ remote.StreamingMetaHandler = MetainfoServerHandler
)

type metainfoClientHandler struct{}

func (mi *metainfoClientHandler) OnConnectStream(ctx context.Context) (context.Context, error) {
	// gRPC send meta when connection is establishing
	// so put metainfo into metadata before sending http2 headers
	ri := rpcinfo.GetRPCInfo(ctx)
	if isGRPC(ri) {
		// append kitex metainfo into metadata list
		// kitex metainfo key starts with " rpc-transit-xxx "
		md, ok := metadata.FromOutgoingContext(ctx)
		if !ok {
			md = metadata.MD{}
		}
		metainfo.ToHTTPHeader(ctx, metainfo.HTTPHeader(md))
		if streamLogID := logid.GetStreamLogID(ctx); streamLogID != "" {
			md[transmeta.HTTPStreamLogID] = []string{streamLogID}
		}
		ctx = metadata.NewOutgoingContext(ctx, md)
	}
	return ctx, nil
}

func (mi *metainfoClientHandler) OnReadStream(ctx context.Context) (context.Context, error) {
	return ctx, nil
}

func (mi *metainfoClientHandler) WriteMeta(ctx context.Context, sendMsg remote.Message) (context.Context, error) {
	if metainfo.HasMetaInfo(ctx) {
		kvs := make(map[string]string)
		metainfo.SaveMetaInfoToMap(ctx, kvs)
		sendMsg.TransInfo().PutTransStrInfo(kvs)
	}
	return ctx, nil
}

func (mi *metainfoClientHandler) ReadMeta(ctx context.Context, recvMsg remote.Message) (context.Context, error) {
	if info := recvMsg.TransInfo(); info != nil {
		if kvs := info.TransStrInfo(); len(kvs) > 0 {
			metainfo.SetBackwardValuesFromMap(ctx, kvs)
		}
	}
	return ctx, nil
}

type metainfoServerHandler struct{}

func (mi *metainfoServerHandler) ReadMeta(ctx context.Context, recvMsg remote.Message) (context.Context, error) {
	if kvs := recvMsg.TransInfo().TransStrInfo(); len(kvs) > 0 {
		ctx = metainfo.SetMetaInfoFromMap(ctx, kvs)
	}
	ctx = metainfo.WithBackwardValuesToSend(ctx)
	return ctx, nil
}

func (mi *metainfoServerHandler) OnConnectStream(ctx context.Context) (context.Context, error) {
	return ctx, nil
}

func (mi *metainfoServerHandler) OnReadStream(ctx context.Context) (context.Context, error) {
	// gRPC server receives meta from http2 header when reading stream
	// read all metadata and put those with kitex-metainfo kvs into ctx
	ri := rpcinfo.GetRPCInfo(ctx)
	if isGRPC(ri) {
		// Attach kitex metainfo into context
		mdata, ok := metadata.FromIncomingContext(ctx)
		if ok {
			ctx = metainfo.FromHTTPHeader(ctx, metainfo.HTTPHeader(mdata))
			ctx = metainfo.WithBackwardValuesToSend(ctx)
			ctx = metainfo.TransferForward(ctx)
			ctx = addStreamIDToContext(ctx, mdata)
		}
	}
	return ctx, nil
}

func addStreamIDToContext(ctx context.Context, md metadata.MD) context.Context {
	streamLogIDValues := md[transmeta.HTTPStreamLogID]
	if len(streamLogIDValues) == 0 {
		// the caller has not set the stream log id, skip
		return ctx
	}
	return logid.NewCtxWithStreamLogID(ctx, streamLogIDValues[0])
}

func (mi *metainfoServerHandler) WriteMeta(ctx context.Context, sendMsg remote.Message) (context.Context, error) {
	if kvs := metainfo.AllBackwardValuesToSend(ctx); len(kvs) > 0 {
		sendMsg.TransInfo().PutTransStrInfo(kvs)
	}

	return ctx, nil
}
