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

from algo_backend.exception import (
    BasicApiId,
    DefaultApiErrorCode,
)
from algo_backend.metrics import PrometheusTimeCostMetricSetting
from algo_backend.schema import AbstractRespVo
from .exception_to_vo import gen_vo_from_exception

logger = logging.getLogger(__name__)


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

    def decorator(
        func: Callable[..., Awaitable[AbstractRespVo]],
    ) -> Callable[..., Awaitable[AbstractRespVo]]:
        @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: AbstractRespVo = 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()}"
                )

                vo_cls: type(AbstractRespVo) = get_type_hints(func).get(
                    "return"
                )  # 这里可能会失败，因为无法强制用户的类型
                return gen_vo_from_exception(
                    vo_cls, e, api_name=_name, request_id=request_id, api_id=api_id
                )

        return wrapper

    return decorator if func is None else decorator(func)
