/*
 * Copyright 2023 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 kitexutil provides some util methods to get RPC information
package kitexutil

import (
	"context"
	"net"

	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/utils"
)

// GetCaller is used to get the Service Name of the caller.
// Return false if failed to get the information.
func GetCaller(ctx context.Context) (string, bool) {
	defer func() { recover() }()

	ri := rpcinfo.GetRPCInfo(ctx)
	if ri == nil {
		return "", false
	}
	return ri.From().ServiceName(), true
}

// GetCallee is used to get the Service Name of the callee.
// Return false if failed to get the information.
func GetCallee(ctx context.Context) (string, bool) {
	defer func() { recover() }()

	ri := rpcinfo.GetRPCInfo(ctx)
	if ri == nil {
		return "", false
	}
	return ri.To().ServiceName(), true
}

// GetMethod is used to get the current RPC Method name.
// Return false if failed to get the information.
func GetMethod(ctx context.Context) (string, bool) {
	defer func() { recover() }()

	ri := rpcinfo.GetRPCInfo(ctx)
	if ri == nil {
		return "", false
	}
	return ri.To().Method(), true
}

// GetCallerHandlerMethod is used to get the method name of caller.
// Only the caller is a Kitex server will have this method information by default, or you can set K_METHOD into context.Context then kitex will get it.
// Return false if failed to get the information.
func GetCallerHandlerMethod(ctx context.Context) (string, bool) {
	defer func() { recover() }()

	ri := rpcinfo.GetRPCInfo(ctx)
	if ri == nil {
		return "", false
	}
	return ri.From().Method(), true
}

// GetIDLServiceName gets the serviceName which defined in IDL.
// Return false if failed to get the information.
func GetIDLServiceName(ctx context.Context) (string, bool) {
	defer func() { recover() }()

	ri := rpcinfo.GetRPCInfo(ctx)
	if ri == nil {
		return "", false
	}
	return ri.Invocation().ServiceName(), true
}

// GetCallerAddr is used for the server to get the Address of the caller.
// Return false if failed to get the information.
func GetCallerAddr(ctx context.Context) (net.Addr, bool) {
	defer func() { recover() }()

	ri := rpcinfo.GetRPCInfo(ctx)
	if ri == nil || ri.From() == nil || ri.From().Address() == nil {
		return nil, false
	}
	return ri.From().Address(), true
}

// GetCallerIP is used for the server to get the IP of the caller.
// Return false if failed to get the information.
func GetCallerIP(ctx context.Context) (string, bool) {
	defer func() { recover() }()

	ri := rpcinfo.GetRPCInfo(ctx)
	if ri == nil || ri.From() == nil || ri.From().Address() == nil {
		return "", false
	}
	addrStr := ri.From().Address().String()
	if len(addrStr) == 0 {
		return "", false
	}

	if ip, _, err := net.SplitHostPort(addrStr); err == nil {
		return ip, true
	}
	return addrStr, true
}

// GetTransportProtocol gets the transport protocol of the request.
// Return false if failed to get the information.
func GetTransportProtocol(ctx context.Context) (string, bool) {
	defer func() { recover() }()

	ri := rpcinfo.GetRPCInfo(ctx)
	if ri == nil {
		return "", false
	}
	return ri.Config().TransportProtocol().String(), true
}

// GetRPCInfo gets the RPCInfo in ctx.
// Return false if failed to get the information.
func GetRPCInfo(ctx context.Context) (rpcinfo.RPCInfo, bool) {
	defer func() { recover() }()
	ri := rpcinfo.GetRPCInfo(ctx)
	if ri == nil {
		return nil, false
	}
	return ri, true
}

// GetRealReqFromKitexArgs assert the req to be KitexArgs and return the real request if succeeded, otherwise return nil.
// This method should be used in the middleware.
func GetRealReqFromKitexArgs(req interface{}) interface{} {
	if arg, ok := req.(utils.KitexArgs); ok {
		return arg.GetFirstArgument()
	}
	return nil
}

// GetRealRespFromKitexResult assert the req to be KitexResult and return the real response if succeeded, otherwise return nil.
// This method should be used in the middleware.
func GetRealRespFromKitexResult(resp interface{}) interface{} {
	if res, ok := resp.(utils.KitexResult); ok {
		return res.GetResult()
	}
	return nil
}
