# -*- coding:utf-8 -*-

from functools import wraps

from gm_tracer.context import register_tracer, current_tracer, get_current_span, active_span
from helios.rpc import get_mesh_info, set_mesh_info


trace_info_attr_name = '_trace_attr_'


def init_tracer(instance=None):
    '''
    获取请求链路跟踪数据, 返回或封装实例
    :return:
        {
           'tracer':,
           'mesh_attr':
        }
    '''
    tracer = current_tracer()
    mesh_attr = get_mesh_info()
    span = get_current_span()
    trace_info = {
        'tracer': tracer,
        'mesh_attr': mesh_attr,
        'span': span,
    }
    if instance:
        setattr(instance, trace_info_attr_name, trace_info)
    return trace_info


# # 类里调用gevent并发请求时使用, 和init_tracer配合使用
# def wrap_tracer(func):
#     @wraps(func)
#     def deco(self, *args, **kwargs):
#         trace_info = self.trace_info
#         register_tracer(trace_info['tracer'])
#         set_mesh_info(trace_info['mesh_attr'])
#         return func(self, *args, **kwargs)
#     return deco


# 支持func & method
# eg:
# func:
#   @wrap_tracer
#   def func(self, request, *args, **kwargs):
#       pass
# method:
#   @wrap_tracer(**{'tracer':1, 'mesh_attr':2})
#   def method(*args, **kwargs):
#       pass

def wrap_tracer(*args, **kwargs):
    def _inner_tracer(func):
        @wraps(func)
        def deco(*f_args, **f_kwargs):
            if kwargs:
                trace_info = kwargs
            else:
                trace_info = getattr(f_args[0], trace_info_attr_name, {})
            register_tracer(trace_info['tracer'])
            set_mesh_info(trace_info['mesh_attr'])
            active_span(trace_info['span'])
            return func(*f_args, **f_kwargs)
        return deco
    if len(args) == 1 and callable(args[0]):
        return _inner_tracer(args[0])
    else:
        return _inner_tracer