import time
import traceback
from functools import wraps
from typing import Callable, Awaitable

from loguru import logger

from huibiao_framework.backend import (
    BaseRespVo,
    BasicCommonException,
    BasicApiId,
    CommonStatusCode,
)
from huibiao_framework.backend.metrics.request_metrics import (
    PrometheusTimeCostMetricSetting,
)
from huibiao_framework.backend.vo import ApiErrorCodeManage


def timing_and_exception_handler(func=None, *, api_id: BasicApiId, api_name: str = ""):
    """
    装饰器：用于统计函数执行时间并捕获异常
    函数中需要包含参数reqid或者request_id
    """

    def decorator(
        func: Callable[..., Awaitable[BaseRespVo]],
    ) -> Callable[..., Awaitable[BaseRespVo]]:
        @wraps(func)
        async def wrapper(*args, **kwargs):
            start_time = time.perf_counter()

            request_id = kwargs.get("request_id", None) or kwargs.get("reqid", None)
            _name = func.__name__ or api_name

            try:
                # 执行原函数
                logger.info(f"ReqId: {request_id} | Function: {_name} | Start")
                result: BaseRespVo = await func(*args, **kwargs)
                # 计算耗时
                elapsed_time = time.perf_counter() - start_time
                logger.info(
                    f"ReqId: {request_id} | Function: {_name} | COST:{elapsed_time:.4f}s"
                )
                PrometheusTimeCostMetricSetting.api_metrics_instance().add(
                    _name, elapsed_time
                )
                return result.set_request_id(request_id)
            except Exception as e:
                # 计算耗时
                elapsed_time = time.perf_counter() - start_time
                PrometheusTimeCostMetricSetting.api_metrics_instance().add_error(_name)
                # 记录异常信息和完整堆栈
                logger.error(
                    f"ReqId: {request_id} | Function: {_name} | COST:{elapsed_time:.4f}s | Exception: {str(e)}\n"
                    f"Traceback:\n{traceback.format_exc()}"
                )
                if isinstance(e, BasicCommonException):
                    vo = BaseRespVo(
                        code=e.code, message=f"接口[{_name}]报错:{e.msg}", result=None
                    )
                else:
                    api_error_code, api_error_msg = (
                        ApiErrorCodeManage.gen_api_error_code_msg(api_id, _name, e)
                    )
                    vo = BaseRespVo(
                        code=api_error_code,
                        message=api_error_msg,
                        result=None,
                    )
                vo.set_request_id(request_id)
                return vo

        return wrapper

    return decorator if func is None else decorator(func)


def basic_exception_handler(func=None, *, api_name: str = ""):
    """
    装饰器：用于统计函数执行时间并捕获异常
    """

    def decorator(
        func: Callable[..., Awaitable[BaseRespVo]],
    ) -> Callable[..., Awaitable[BaseRespVo]]:
        @wraps(func)
        async def wrapper(*args, **kwargs):
            _name = func.__name__ or api_name
            try:
                result: BaseRespVo = await func(*args, **kwargs)
                return result
            except Exception as e:
                logger.error(
                    f"Function: {_name}| Exception: {str(e)} Traceback:\n{traceback.format_exc()}"
                )
                if isinstance(e, BasicCommonException):
                    vo = BaseRespVo(
                        code=e.code, message=f"接口[{_name}]报错:{e.msg}", result=None
                    )
                else:
                    vo = BaseRespVo(
                        code=CommonStatusCode.UNKNOWN_ERROR.code,
                        message=f"接口[{_name}]报错:{str(e)}",
                        result=None,
                    )
                return vo

        return wrapper

    return decorator if func is None else decorator(func)
